// SPDX-License-Identifier: Apache-2.0

package core

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"errors"
	"fmt"
	"log/slog"
	"math/rand/v2"
	"sync"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/keepalive"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"

	"git.beugsiud.de/beug.k/DDI/internal/agentpki"
	"git.beugsiud.de/beug.k/DDI/internal/buildinfo"
	"git.beugsiud.de/beug.k/DDI/internal/bundle"
	agentv1 "git.beugsiud.de/beug.k/DDI/proto/gen/ddi/agent/v1"
)

// Defaults of the server connection (docs/adr/0021-agent-strecke-m3.md).
const (
	defaultPullInterval      = 60 * time.Second
	defaultHeartbeatInterval = 30 * time.Second
	keepaliveInterval        = 30 * time.Second
	keepaliveTimeout         = 10 * time.Second
	backoffMin               = time.Second
	backoffMax               = 60 * time.Second
	welcomeTimeout           = 30 * time.Second
	rpcTimeout               = 30 * time.Second
	// renewRetryInterval spaces failed renewal attempts (also with an
	// expired certificate, which the server accepts for 30 days).
	renewRetryInterval = 15 * time.Minute
	// oldConnGrace keeps the connection of a replaced certificate open for
	// calls in flight.
	oldConnGrace = 30 * time.Second
	// maxBundleMessage allows bundles up to the bundle size limit.
	maxBundleMessage = bundle.MaxTotalSize + 16<<20
)

// errReconnect ends a session after the certificate was renewed.
var errReconnect = errors.New("reconnecting with the renewed certificate")

// ServerSource is the Source backed by the management server: an mTLS gRPC
// connection with the Connect control stream and GetDesiredState polling as
// fallback (docs/adr/0007-server-agent-grpc-mtls.md). It also renews the
// client certificate.
type ServerSource struct {
	store      *Store
	agentID    string
	ca         *x509.Certificate
	server     string
	serverName string
	hostname   string
	log        *slog.Logger

	// Timing, shortened by tests.
	backoffMin, backoffMax time.Duration
	welcomeTimeout         time.Duration
	renewRetry             time.Duration
	// dialOptions are appended to the defaults (tests).
	dialOptions []grpc.DialOption

	mu                sync.Mutex
	cert              tls.Certificate
	conn              *grpc.ClientConn
	pullInterval      time.Duration
	heartbeatInterval time.Duration
	lastRenewTry      time.Time

	renewMu sync.Mutex
}

// NewServerSource creates the source for an enrolled identity.
func NewServerSource(store *Store, id *Identity, hostname string, log *slog.Logger) *ServerSource {
	return &ServerSource{
		store: store, agentID: id.AgentID, ca: id.CA, server: id.Server, serverName: id.ServerName,
		hostname: hostname, log: log, cert: id.Certificate,
		backoffMin: backoffMin, backoffMax: backoffMax, welcomeTimeout: welcomeTimeout,
		renewRetry: renewRetryInterval, pullInterval: defaultPullInterval,
		heartbeatInterval: defaultHeartbeatInterval,
	}
}

// client returns the gRPC client, creating the connection on first use.
// Connections are established lazily by gRPC.
func (s *ServerSource) client() (agentv1.AgentServiceClient, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.conn == nil {
		opts := append([]grpc.DialOption{
			grpc.WithTransportCredentials(credentials.NewTLS(agentpki.ClientTLSConfig(s.ca, s.cert, s.serverName))),
			grpc.WithKeepaliveParams(keepalive.ClientParameters{
				Time: keepaliveInterval, Timeout: keepaliveTimeout, PermitWithoutStream: true,
			}),
			grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(maxBundleMessage)),
		}, s.dialOptions...)
		conn, err := grpc.NewClient(s.server, opts...)
		if err != nil {
			return nil, fmt.Errorf("connect to %s: %w", s.server, err)
		}
		s.conn = conn
	}
	return agentv1.NewAgentServiceClient(s.conn), nil
}

// replaceConn drops the current connection (after a certificate renewal);
// the next call dials with the new certificate.
func (s *ServerSource) replaceConn() {
	s.mu.Lock()
	old := s.conn
	s.conn = nil
	s.mu.Unlock()
	if old != nil {
		time.AfterFunc(oldConnGrace, func() { _ = old.Close() })
	}
}

func (s *ServerSource) closeConn() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.conn != nil {
		_ = s.conn.Close()
		s.conn = nil
	}
}

func (s *ServerSource) intervals() (pull, heartbeat time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.pullInterval, s.heartbeatInterval
}

func (s *ServerSource) setIntervals(pull, heartbeat time.Duration) {
	s.mu.Lock()
	defer s.mu.Unlock()
	if pull > 0 {
		s.pullInterval = pull
	}
	if heartbeat > 0 {
		s.heartbeatInterval = heartbeat
	}
}

// Run implements Source: it keeps a Connect stream open, reconnecting with
// exponential backoff and jitter, and polls GetDesiredState every pull
// interval while no stream can be established.
func (s *ServerSource) Run(ctx context.Context, h SourceHandler) error {
	defer s.closeConn()
	backoff := s.backoffMin
	var lastPoll time.Time
	reachable := false
	for ctx.Err() == nil {
		s.maybeRenew(ctx)
		established, err := s.session(ctx, h)
		if ctx.Err() != nil {
			break
		}
		wait := true
		switch {
		case errors.Is(err, errReconnect):
			wait = false
			backoff = s.backoffMin
		case established:
			s.log.Warn("control stream to the server closed", "error", err)
			backoff = s.backoffMin
			reachable = false
		default:
			s.log.Warn("cannot open control stream", "server", s.server, "error", err)
			pull, _ := s.intervals()
			if time.Since(lastPoll) >= pull {
				lastPoll = time.Now()
				if err := s.poll(ctx, h); err != nil {
					s.log.Warn("cannot poll the desired state", "error", err)
					reachable = false
				} else if !reachable {
					reachable = true
					h.Connected()
				}
			}
		}
		if wait {
			if !sleepCtx(ctx, jitter(backoff)) {
				break
			}
			backoff = min(backoff*2, s.backoffMax)
		}
	}
	return ctx.Err()
}

// session runs one Connect stream. established reports whether the server
// answered the Hello with a Welcome.
func (s *ServerSource) session(ctx context.Context, h SourceHandler) (established bool, err error) {
	cl, err := s.client()
	if err != nil {
		return false, err
	}
	sctx, cancel := context.WithCancel(ctx)
	defer cancel()
	stream, err := cl.Connect(sctx)
	if err != nil {
		return false, err
	}
	defer func() { _ = stream.CloseSend() }()

	st := h.Status(sctx)
	if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_Hello{Hello: &agentv1.Hello{
		AgentVersion: buildinfo.Get().Version, AppliedRevision: st.Applied, LastGoodRevision: st.LastGood,
		Health: healthToProto(st.Health, st.CheckedAt), ServiceVersion: st.ServiceVersion, Hostname: s.hostname,
	}}}); err != nil {
		return false, err
	}

	msgs := make(chan *agentv1.ConnectResponse)
	recvErr := make(chan error, 1)
	go func() {
		for {
			msg, err := stream.Recv()
			if err != nil {
				recvErr <- err
				return
			}
			select {
			case msgs <- msg:
			case <-sctx.Done():
				return
			}
		}
	}()

	timer := time.NewTimer(s.welcomeTimeout)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return false, ctx.Err()
	case err := <-recvErr:
		return false, err
	case <-timer.C:
		return false, errors.New("server did not answer the Hello in time")
	case msg := <-msgs:
		w := msg.GetWelcome()
		if w == nil {
			return false, errors.New("server did not answer the Hello with a Welcome")
		}
		s.log.Info("connected to the server", "server", s.server, "mode", w.GetMode().String(),
			"desired_revision", w.GetDesiredRevision())
		s.welcome(w, h, true)
	}

	_, hbInterval := s.intervals()
	heartbeat := time.NewTicker(hbInterval)
	defer heartbeat.Stop()
	renew := time.NewTimer(s.untilRenewal())
	defer renew.Stop()
	for {
		select {
		case <-ctx.Done():
			return true, ctx.Err()
		case err := <-recvErr:
			return true, err
		case msg := <-msgs:
			switch m := msg.GetMessage().(type) {
			case *agentv1.ConnectResponse_Welcome:
				s.welcome(m.Welcome, h, false)
				_, hb := s.intervals()
				heartbeat.Reset(hb)
			case *agentv1.ConnectResponse_Command:
				if err := s.command(ctx, stream, m.Command, h); err != nil {
					return true, err
				}
			}
		case <-heartbeat.C:
			st := h.Status(sctx)
			if err := stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_Heartbeat{Heartbeat: &agentv1.Heartbeat{
				AppliedRevision: st.Applied, LastGoodRevision: st.LastGood, Health: healthToProto(st.Health, st.CheckedAt),
			}}}); err != nil {
				return true, err
			}
		case <-renew.C:
			if s.maybeRenew(ctx) {
				return true, errReconnect
			}
			renew.Reset(s.untilRenewal())
		}
	}
}

// welcome applies the settings of a Welcome and announces its desired
// revision.
func (s *ServerSource) welcome(w *agentv1.Welcome, h SourceHandler, first bool) {
	s.setIntervals(w.GetPullInterval().AsDuration(), w.GetHeartbeatInterval().AsDuration())
	if first {
		h.Connected()
	}
	h.Desired(w.GetDesiredRevision(), w.GetMode() != agentv1.OperatingMode_OPERATING_MODE_MANUAL)
}

// command acknowledges and executes a server command.
func (s *ServerSource) command(ctx context.Context, stream agentv1.AgentService_ConnectClient, c *agentv1.Command, h SourceHandler) error {
	ack := func(accepted bool, msg string) error {
		return stream.Send(&agentv1.ConnectRequest{Message: &agentv1.ConnectRequest_CommandAck{CommandAck: &agentv1.CommandAck{
			CommandId: c.GetCommandId(), Accepted: accepted, Message: msg,
		}}})
	}
	log := s.log.With("command_id", c.GetCommandId())
	switch a := c.GetAction().(type) {
	case *agentv1.Command_ApplyRevision:
		rev, dep := a.ApplyRevision.GetRevision(), a.ApplyRevision.GetDeploymentId()
		if rev == 0 {
			log.Warn("apply command without revision")
			return ack(false, "revision is missing")
		}
		log.Info("apply command received", "revision", rev, "deployment_id", dep)
		if err := ack(true, ""); err != nil {
			return err
		}
		h.Push(rev, dep)
		return nil
	case *agentv1.Command_RenewCertificate:
		log.Info("certificate renewal requested by the server")
		if err := ack(true, ""); err != nil {
			return err
		}
		if err := s.renew(ctx); err != nil {
			log.Error("certificate renewal failed", "error", err)
			return nil
		}
		return errReconnect
	default:
		log.Warn("unknown command")
		return ack(false, "unknown command")
	}
}

// poll asks for the desired state once.
func (s *ServerSource) poll(ctx context.Context, h SourceHandler) error {
	cl, err := s.client()
	if err != nil {
		return err
	}
	st := h.Status(ctx)
	rctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	resp, err := cl.GetDesiredState(rctx, &agentv1.GetDesiredStateRequest{AppliedRevision: st.Applied})
	if err != nil {
		return err
	}
	s.setIntervals(resp.GetPullInterval().AsDuration(), 0)
	rev, dep := resp.GetDesiredRevision(), resp.GetDeploymentId()
	s.log.Debug("desired state polled", "desired_revision", rev, "mode", resp.GetMode().String(), "deployment_id", dep)
	if dep != "" && rev != 0 {
		// An open push: applied in every mode (docs/adr/0022-ausrollen-pull-push-m4.md).
		s.log.Info("open push received by polling", "revision", rev, "deployment_id", dep)
		h.Push(rev, dep)
		return nil
	}
	h.Desired(rev, resp.GetMode() != agentv1.OperatingMode_OPERATING_MODE_MANUAL)
	return nil
}

// leaf returns the current client certificate.
func (s *ServerSource) leaf() *x509.Certificate {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.cert.Leaf
}

// renewalDue returns when the certificate should be renewed: when half of
// its lifetime has passed.
func renewalDue(c *x509.Certificate) time.Time {
	return c.NotBefore.Add(c.NotAfter.Sub(c.NotBefore) / 2)
}

// untilRenewal returns the time until the next renewal attempt.
func (s *ServerSource) untilRenewal() time.Duration {
	due := renewalDue(s.leaf())
	s.mu.Lock()
	last := s.lastRenewTry
	s.mu.Unlock()
	if !last.IsZero() && time.Now().After(due) {
		due = last.Add(s.renewRetry)
	}
	return max(time.Until(due), 0)
}

// maybeRenew renews the certificate if half of its lifetime has passed
// (or it expired), spacing failed attempts by renewRetry.
func (s *ServerSource) maybeRenew(ctx context.Context) bool {
	leaf := s.leaf()
	now := time.Now()
	if now.Before(renewalDue(leaf)) {
		return false
	}
	s.mu.Lock()
	if !s.lastRenewTry.IsZero() && now.Sub(s.lastRenewTry) < s.renewRetry {
		s.mu.Unlock()
		return false
	}
	s.lastRenewTry = now
	s.mu.Unlock()
	if err := s.renew(ctx); err != nil {
		if now.After(leaf.NotAfter) {
			s.log.Error("client certificate expired and renewal failed; the agent must be enrolled again if this persists",
				"not_after", leaf.NotAfter, "error", err)
		} else {
			s.log.Warn("certificate renewal failed", "not_after", leaf.NotAfter, "error", err)
		}
		return false
	}
	return true
}

// renew obtains a certificate for a new key and stores it.
func (s *ServerSource) renew(ctx context.Context) error {
	s.renewMu.Lock()
	defer s.renewMu.Unlock()
	key, csr, err := newKeyAndCSR(s.agentID)
	if err != nil {
		return err
	}
	cl, err := s.client()
	if err != nil {
		return err
	}
	rctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	resp, err := cl.RenewCertificate(rctx, &agentv1.RenewCertificateRequest{CsrDer: csr})
	if err != nil {
		return err
	}
	if chain := resp.GetCaCertificatesDer(); len(chain) > 0 {
		if root, err := x509.ParseCertificate(chain[len(chain)-1]); err != nil || !root.Equal(s.ca) {
			s.log.Warn("server sent a different CA certificate; keeping the enrolled CA")
		}
	}
	cert, err := checkIssued(resp.GetCertificateDer(), s.ca, s.agentID, key)
	if err != nil {
		return err
	}
	if err := s.store.SaveCertificate(key, cert); err != nil {
		return err
	}
	s.mu.Lock()
	s.cert = keyPair(key, cert)
	s.lastRenewTry = time.Time{}
	s.mu.Unlock()
	s.replaceConn()
	s.log.Info("client certificate renewed", "not_after", cert.NotAfter)
	return nil
}

// FetchBundle implements Source.
func (s *ServerSource) FetchBundle(ctx context.Context, revision uint64) (*Bundle, error) {
	cl, err := s.client()
	if err != nil {
		return nil, err
	}
	resp, err := cl.FetchBundle(ctx, &agentv1.FetchBundleRequest{Revision: revision})
	if err != nil {
		return nil, err
	}
	b := &Bundle{Manifest: resp.GetManifest(), Signature: resp.GetSignature()}
	for _, a := range resp.GetArtifacts() {
		b.Artifacts = append(b.Artifacts, BundleArtifact{Path: a.GetPath(), Content: a.GetContent()})
	}
	return b, nil
}

// ReportDeployment implements Source.
func (s *ServerSource) ReportDeployment(ctx context.Context, r DeploymentReport) error {
	cl, err := s.client()
	if err != nil {
		return err
	}
	_, err = cl.ReportDeployment(ctx, &agentv1.ReportDeploymentRequest{
		DeploymentId: r.DeploymentID, Revision: r.Revision, Trigger: triggerToProto(r.Trigger),
		Outcome: outcomeToProto(r.Outcome), Message: r.Message, ActiveRevision: r.ActiveRevision,
		StartedAt: timestamppb.New(r.StartedAt), FinishedAt: timestamppb.New(r.FinishedAt),
	})
	if refusedReportCodes[status.Code(err)] {
		return fmt.Errorf("%w: %w", ErrReportRefused, err)
	}
	return err
}

// refusedReportCodes are the answers to ReportDeployment that a retry
// cannot change: the server considers the report invalid or does not
// expect it (any more).
var refusedReportCodes = map[codes.Code]bool{
	codes.InvalidArgument:    true,
	codes.NotFound:           true,
	codes.AlreadyExists:      true,
	codes.FailedPrecondition: true,
	codes.OutOfRange:         true,
}

// ReportLeases implements Source.
func (s *ServerSource) ReportLeases(ctx context.Context, b LeaseBatch) error {
	cl, err := s.client()
	if err != nil {
		return err
	}
	req := &agentv1.ReportLeasesRequest{Sequence: b.Sequence, FullSync: b.FullSync, Leases: make([]*agentv1.Lease, len(b.Leases))}
	for i, l := range b.Leases {
		req.Leases[i] = leaseToProto(l)
	}
	rctx, cancel := context.WithTimeout(ctx, rpcTimeout)
	defer cancel()
	resp, err := cl.ReportLeases(rctx, req)
	if err != nil {
		return err
	}
	if ack := resp.GetAcknowledgedSequence(); ack != b.Sequence {
		s.log.Debug("lease report acknowledged with another sequence", "sequence", b.Sequence, "acknowledged", ack)
	}
	return nil
}

func healthToProto(h Health, at time.Time) *agentv1.HealthStatus {
	out := &agentv1.HealthStatus{State: healthStateToProto(h.State), Message: h.Message}
	if !at.IsZero() {
		out.CheckedAt = timestamppb.New(at)
	}
	for _, c := range h.Checks {
		out.Checks = append(out.Checks, &agentv1.HealthCheck{Name: c.Name, State: healthStateToProto(c.State), Message: c.Message})
	}
	return out
}

func healthStateToProto(s HealthState) agentv1.HealthState {
	switch s {
	case HealthOK:
		return agentv1.HealthState_HEALTH_STATE_OK
	case HealthDegraded:
		return agentv1.HealthState_HEALTH_STATE_DEGRADED
	case HealthFailed:
		return agentv1.HealthState_HEALTH_STATE_FAILED
	default:
		return agentv1.HealthState_HEALTH_STATE_UNSPECIFIED
	}
}

func triggerToProto(t Trigger) agentv1.DeploymentTrigger {
	switch t {
	case TriggerPull:
		return agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_PULL
	case TriggerPush:
		return agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_PUSH
	default:
		return agentv1.DeploymentTrigger_DEPLOYMENT_TRIGGER_UNSPECIFIED
	}
}

func outcomeToProto(o Outcome) agentv1.DeploymentOutcome {
	switch o {
	case OutcomeSucceeded:
		return agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_SUCCEEDED
	case OutcomeRejected:
		return agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_REJECTED
	case OutcomeRolledBack:
		return agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_ROLLED_BACK
	case OutcomeFailed:
		return agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_FAILED
	default:
		return agentv1.DeploymentOutcome_DEPLOYMENT_OUTCOME_UNSPECIFIED
	}
}

func leaseToProto(l Lease) *agentv1.Lease {
	out := &agentv1.Lease{
		Ip: l.IP.String(), Mac: l.MAC, Hostname: l.Hostname, State: leaseStateToProto(l.State),
		SubnetId: l.SubnetID, ClientId: l.ClientID,
	}
	if !l.Expires.IsZero() {
		out.ExpiresAt = timestamppb.New(l.Expires)
	}
	if !l.Updated.IsZero() {
		out.UpdatedAt = timestamppb.New(l.Updated)
	}
	return out
}

func leaseStateToProto(s LeaseState) agentv1.LeaseState {
	switch s {
	case LeaseActive:
		return agentv1.LeaseState_LEASE_STATE_ACTIVE
	case LeaseExpired:
		return agentv1.LeaseState_LEASE_STATE_EXPIRED
	case LeaseReleased:
		return agentv1.LeaseState_LEASE_STATE_RELEASED
	case LeaseDeclined:
		return agentv1.LeaseState_LEASE_STATE_DECLINED
	default:
		return agentv1.LeaseState_LEASE_STATE_UNSPECIFIED
	}
}

// jitter returns d randomized by ±20 %.
func jitter(d time.Duration) time.Duration {
	return time.Duration(float64(d) * (0.8 + 0.4*rand.Float64())) //nolint:gosec // jitter needs no cryptographic randomness
}

// sleepCtx waits for d; it returns false if ctx ended first.
func sleepCtx(ctx context.Context, d time.Duration) bool {
	t := time.NewTimer(d)
	defer t.Stop()
	select {
	case <-ctx.Done():
		return false
	case <-t.C:
		return true
	}
}