// agent.go — the gRPC agent plane: the long-lived bidi stream each host holds.
//
// mTLS is verified against the FLEET CA (the host's reused TPM identity cert),
// distinct from the admin plane's dedicated admin CA. The host's cert SERIAL is
// the inventory key, taken from the peer cert — never self-asserted — so a host
// can only report/receive its own state.
//
// Concurrency model per Session: one receive loop (host->manager) in the handler
// goroutine + one send loop draining a per-connection channel (manager->host).
// The channel decouples admin-initiated pushes (SetDesired, Command) from the
// stream's Send, which must not be called concurrently. Inventory is the source
// of truth; the live stream is only the low-latency delivery path — a host that
// is offline or whose send buffer is full simply gets the current desired-state
// on its next (re)connect.
package main

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"log"
	"net"
	"os"
	"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/peer"
	"google.golang.org/grpc/status"

	fleetpb "fleet-manager/proto"
)

// MetricsSink receives each health snapshot for time-series storage. The no-op
// sink is the default; the InfluxDB sink lands in checkpoint 4.
type MetricsSink interface {
	WriteSnapshot(serial, hostnqn, starship string, s *fleetpb.HealthSnapshot, at time.Time)
}

type noopSink struct{}

func (noopSink) WriteSnapshot(string, string, string, *fleetpb.HealthSnapshot, time.Time) {}

// agentConn is one live host stream's downstream (manager->host) side.
type agentConn struct {
	serial string
	send   chan *fleetpb.ManagerDown
	done   chan struct{}
	once   sync.Once
}

func (c *agentConn) close() { c.once.Do(func() { close(c.done) }) }

// agentHub implements the FleetAgent gRPC service and holds the live-stream
// registry that admin actions push through.
type agentHub struct {
	fleetpb.UnimplementedFleetAgentServer
	store Store
	sink  MetricsSink

	mu    sync.Mutex
	conns map[string]*agentConn // serial -> live connection
}

func newAgentHub(store Store, sink MetricsSink) *agentHub {
	if sink == nil {
		sink = noopSink{}
	}
	return &agentHub{store: store, sink: sink, conns: map[string]*agentConn{}}
}

// register installs a connection, evicting any prior stream for the same serial
// (a reconnect supersedes a lingering old one).
func (h *agentHub) register(c *agentConn) {
	h.mu.Lock()
	if old := h.conns[c.serial]; old != nil {
		old.close()
	}
	h.conns[c.serial] = c
	h.mu.Unlock()
}

func (h *agentHub) deregister(c *agentConn) {
	h.mu.Lock()
	if h.conns[c.serial] == c { // don't clobber a newer reconnect
		delete(h.conns, c.serial)
	}
	h.mu.Unlock()
	c.close()
}

// pushDown delivers a message to a host's live stream if connected. Returns false
// if the host is offline or its buffer is backed up (it will reconcile on its
// next connect from persisted desired-state — no message is "lost" that matters).
func (h *agentHub) pushDown(serial string, m *fleetpb.ManagerDown) bool {
	h.mu.Lock()
	c := h.conns[serial]
	h.mu.Unlock()
	if c == nil {
		return false
	}
	select {
	case c.send <- m:
		return true
	case <-c.done:
		return false
	default:
		return false // slow consumer; drop — desired-state is durable in inventory
	}
}

// pushDesired sends a host its current desired-state over the live stream.
func (h *agentHub) pushDesired(serial string, d Desired) bool {
	return h.pushDown(serial, desiredMsg(d))
}

// Session is the bidirectional stream RPC.
func (h *agentHub) Session(stream fleetpb.FleetAgent_SessionServer) error {
	serial, err := peerSerial(stream.Context())
	if err != nil {
		return status.Error(codes.Unauthenticated, err.Error())
	}
	// A host must be enrolled (inventory record exists) before it can stream.
	host, ok := h.store.Get(serial)
	if !ok {
		return status.Errorf(codes.PermissionDenied, "host %s is not enrolled", serial)
	}

	c := &agentConn{serial: serial, send: make(chan *fleetpb.ManagerDown, 16), done: make(chan struct{})}
	h.register(c)
	h.store.SetOnline(serial, true, time.Now().UTC())
	log.Printf("agent connected: %s (%s)", serial, host.Identity.StarshipName)
	defer func() {
		h.deregister(c)
		h.store.SetOnline(serial, false, time.Now().UTC())
		log.Printf("agent disconnected: %s", serial)
	}()

	// Deliver current desired-state immediately on connect.
	c.send <- desiredMsg(host.Desired)

	// Sender goroutine: the ONLY caller of stream.Send.
	sendErr := make(chan error, 1)
	go func() {
		for {
			select {
			case m := <-c.send:
				if err := stream.Send(m); err != nil {
					sendErr <- err
					return
				}
			case <-c.done:
				return
			}
		}
	}()

	// Receive loop (this goroutine).
	for {
		select {
		case err := <-sendErr:
			return err
		default:
		}
		up, err := stream.Recv()
		if err != nil {
			return err // EOF / transport error -> disconnect (defer cleans up)
		}
		h.handleUp(serial, up)
	}
}

func (h *agentHub) handleUp(serial string, up *fleetpb.AgentUp) {
	switch m := up.GetMsg().(type) {
	case *fleetpb.AgentUp_Hello:
		hello := m.Hello
		host, ok := h.store.Get(serial)
		if !ok {
			return
		}
		_, _, _ = h.store.ReportState(serial,
			Version{ImageVersion: hello.GetImageVersion(), RootfsDigest: hello.GetRootfsDigest(), Kernel: hello.GetKernel()},
			host.Health, // unchanged until first snapshot
			Usage{KioskURL: hello.GetKioskUrl(), Display: host.Usage.Display},
			time.Now().UTC())
	case *fleetpb.AgentUp_Snapshot:
		snap := m.Snapshot
		at := time.Now().UTC()
		if snap.GetAtUnix() > 0 {
			at = time.Unix(snap.GetAtUnix(), 0).UTC()
		}
		host, ok := h.store.Get(serial)
		if !ok {
			return
		}
		h.store.ReportState(serial, host.Version, snapshotToHealth(snap), snapshotToUsage(snap, host.Usage), at)
		h.sink.WriteSnapshot(serial, host.Identity.Hostnqn, host.Identity.StarshipName, snap, at)
	case *fleetpb.AgentUp_Result:
		res := m.Result
		log.Printf("agent %s command %s result: ok=%v %s", serial, res.GetCommandId(), res.GetOk(), res.GetDetail())
	}
}

// ---- conversions ----------------------------------------------------------

func desiredMsg(d Desired) *fleetpb.ManagerDown {
	var rebootAt int64
	if d.RebootAt != nil {
		rebootAt = d.RebootAt.Unix()
	}
	return &fleetpb.ManagerDown{Msg: &fleetpb.ManagerDown_Desired{Desired: &fleetpb.DesiredState{
		TargetVersion: d.TargetVersion,
		KioskUrl:      d.KioskURL,
		RebootAtUnix:  rebootAt,
		Drain:         d.Drain,
	}}}
}

func snapshotToHealth(s *fleetpb.HealthSnapshot) Health {
	return Health{
		LoadAvg:       s.GetLoadAvg(),
		CPUPct:        s.GetCpuPct(),
		TempC:         s.GetTempC(),
		MemUsedPct:    s.GetMemUsedPct(),
		TmpfsUsedPct:  s.GetTmpfsUsedPct(),
		NVMeRoot:      s.GetNvmeRoot(),
		Auth:          s.GetAuth(),
		SELinux:       s.GetSelinux(),
		Chrony:        s.GetChrony(),
		ChronyOffset:  s.GetChronyOffsetSeconds(),
		ChronySource:  s.GetChronySource(),
		ChronyStratum: s.GetChronyStratum(),
		KioskActive:   s.GetKioskActive(),
		UptimeSeconds: s.GetUptimeSeconds(),
	}
}

func snapshotToUsage(s *fleetpb.HealthSnapshot, prev Usage) Usage {
	u := prev
	u.SessionUptime = s.GetSessionUptime()
	u.ChromiumCrashes = int(s.GetChromiumCrashes())
	if s.GetDisplay() != "" {
		u.Display = s.GetDisplay()
	}
	return u
}

// ---- mTLS identity --------------------------------------------------------

// peerSerial extracts the lowercase-hex serial of the client's mTLS leaf cert —
// the inventory key. It is trusted because RequireAndVerifyClientCert already
// verified the chain to the fleet CA.
func peerSerial(ctx context.Context) (string, error) {
	p, ok := peer.FromContext(ctx)
	if !ok {
		return "", fmt.Errorf("no peer in context")
	}
	tlsInfo, ok := p.AuthInfo.(credentials.TLSInfo)
	if !ok || len(tlsInfo.State.PeerCertificates) == 0 {
		return "", fmt.Errorf("no client certificate")
	}
	return serialHex(tlsInfo.State.PeerCertificates[0]), nil
}

// serialHex renders a cert serial as lowercase hex — the same form the broker
// stores (store.go) so the value matches the inventory key.
func serialHex(cert *x509.Certificate) string {
	return fmt.Sprintf("%x", cert.SerialNumber)
}

// ---- gRPC server wiring ---------------------------------------------------

// agentServer builds the configured gRPC server (mTLS against the fleet CA +
// keepalive) with the hub registered. Split from the listen/serve so tests can
// drive it on an ephemeral listener.
func agentServer(cfg config, hub *agentHub) (*grpc.Server, error) {
	cert, err := tls.LoadX509KeyPair(cfg.agentCert, cfg.agentKey)
	if err != nil {
		return nil, fmt.Errorf("agent server cert: %w", err)
	}
	caPEM, err := os.ReadFile(cfg.agentClientCA)
	if err != nil {
		return nil, fmt.Errorf("fleet CA: %w", err)
	}
	pool := x509.NewCertPool()
	if !pool.AppendCertsFromPEM(caPEM) {
		return nil, fmt.Errorf("no certificates parsed from FLEET_MGR_AGENT_CLIENT_CA")
	}
	creds := credentials.NewTLS(&tls.Config{
		MinVersion:   tls.VersionTLS12,
		Certificates: []tls.Certificate{cert},
		ClientAuth:   tls.RequireAndVerifyClientCert,
		ClientCAs:    pool,
	})
	gs := grpc.NewServer(
		grpc.Creds(creds),
		grpc.KeepaliveParams(keepalive.ServerParameters{Time: 30 * time.Second, Timeout: 10 * time.Second}),
		grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{MinTime: 10 * time.Second, PermitWithoutStream: true}),
	)
	fleetpb.RegisterFleetAgentServer(gs, hub)
	return gs, nil
}

// serveAgent starts the gRPC agent listener with mTLS against the fleet CA.
func serveAgent(cfg config, hub *agentHub) error {
	gs, err := agentServer(cfg, hub)
	if err != nil {
		return err
	}
	lis, err := net.Listen("tcp", cfg.agentAddr)
	if err != nil {
		return err
	}
	log.Printf("fleet-manager agent (gRPC/mTLS) listening on %s", cfg.agentAddr)
	return gs.Serve(lis)
}