package main

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"math/big"
	"net"
	"os"
	"path/filepath"
	"testing"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	fleetpb "fleet-manager/proto"
)

// TestAgentSession drives the full agent plane over real mTLS: a host dials in
// with a fleet-CA client cert, the manager derives its serial from that cert,
// pushes the current desired-state on connect, ingests a Hello + HealthSnapshot
// into inventory, and delivers a live desired-state change over the same stream.
func TestAgentSession(t *testing.T) {
	dir := t.TempDir()
	caCert, caKey := makeCA(t)
	caPath := writeFile(t, dir, "ca.crt", pemCert(caCert.Raw))

	srvCertPEM, srvKeyPEM := makeLeaf(t, caCert, caKey, "fleet-manager", big.NewInt(1), true)
	srvCrt := writeFile(t, dir, "srv.crt", srvCertPEM)
	srvKey := writeFile(t, dir, "srv.key", srvKeyPEM)

	const serialHexStr = "1234abcd"
	cliCertPEM, cliKeyPEM := makeLeaf(t, caCert, caKey, "fleet-host", big.NewInt(0x1234abcd), false)
	cliCert, err := tls.X509KeyPair(cliCertPEM, cliKeyPEM)
	if err != nil {
		t.Fatalf("client keypair: %v", err)
	}

	store, err := newJSONStore("") // in-memory
	if err != nil {
		t.Fatalf("store: %v", err)
	}
	if err := store.Upsert(&Host{
		Identity: Identity{CertSerial: serialHexStr, Hostnqn: "nqn.test:host-" + serialHexStr, StarshipName: "test-ship"},
		Desired:  Desired{TargetVersion: "0.5.0", KioskURL: "http://k/"},
	}); err != nil {
		t.Fatalf("seed: %v", err)
	}

	hub := newAgentHub(store, nil)
	gs, err := agentServer(config{agentCert: srvCrt, agentKey: srvKey, agentClientCA: caPath}, hub)
	if err != nil {
		t.Fatalf("agentServer: %v", err)
	}
	lis, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatalf("listen: %v", err)
	}
	go gs.Serve(lis)
	defer gs.Stop()

	// ---- host dials in ----
	caPool := x509.NewCertPool()
	caPool.AddCert(caCert)
	creds := credentials.NewTLS(&tls.Config{
		Certificates: []tls.Certificate{cliCert},
		RootCAs:      caPool,
		ServerName:   "localhost",
	})
	conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(creds))
	if err != nil {
		t.Fatalf("dial: %v", err)
	}
	defer conn.Close()

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := fleetpb.NewFleetAgentClient(conn).Session(ctx)
	if err != nil {
		t.Fatalf("session: %v", err)
	}

	// 1. current desired-state arrives on connect.
	down, err := stream.Recv()
	if err != nil {
		t.Fatalf("recv desired: %v", err)
	}
	if got := down.GetDesired().GetTargetVersion(); got != "0.5.0" {
		t.Fatalf("initial desired target_version = %q, want 0.5.0", got)
	}

	// 2. host reports Hello (running an older version) + a health snapshot.
	if err := stream.Send(&fleetpb.AgentUp{Msg: &fleetpb.AgentUp_Hello{Hello: &fleetpb.Hello{ImageVersion: "0.4.2", Kernel: "5.14.0"}}}); err != nil {
		t.Fatalf("send hello: %v", err)
	}
	if err := stream.Send(&fleetpb.AgentUp{Msg: &fleetpb.AgentUp_Snapshot{Snapshot: &fleetpb.HealthSnapshot{
		TmpfsUsedPct: 63, NvmeRoot: "ok", Auth: "up", Selinux: "enforcing", KioskActive: true,
	}}}); err != nil {
		t.Fatalf("send snapshot: %v", err)
	}

	// 3. inventory reflects the report + online presence.
	waitFor(t, 3*time.Second, func() bool {
		h, ok := store.Get(serialHexStr)
		return ok && h.Online && h.Version.ImageVersion == "0.4.2" && h.Health.TmpfsUsedPct == 63 && h.Health.NVMeRoot == "ok"
	}, "inventory did not reflect the agent report")

	// 4. an admin changes desired-state -> live push over the same stream.
	if _, _, err := store.SetDesired(serialHexStr, Desired{TargetVersion: "0.6.0"}); err != nil {
		t.Fatalf("set desired: %v", err)
	}
	if !hub.pushDesired(serialHexStr, Desired{TargetVersion: "0.6.0"}) {
		t.Fatal("pushDesired reported the host as not live, but it is connected")
	}
	down2, err := stream.Recv()
	if err != nil {
		t.Fatalf("recv desired push: %v", err)
	}
	if got := down2.GetDesired().GetTargetVersion(); got != "0.6.0" {
		t.Fatalf("pushed desired target_version = %q, want 0.6.0", got)
	}
}

// TestAgentRejectsUnenrolled: a valid fleet-CA cert whose serial has no inventory
// record must be refused (enrolment gates the stream).
func TestAgentRejectsUnenrolled(t *testing.T) {
	dir := t.TempDir()
	caCert, caKey := makeCA(t)
	caPath := writeFile(t, dir, "ca.crt", pemCert(caCert.Raw))
	srvCertPEM, srvKeyPEM := makeLeaf(t, caCert, caKey, "fleet-manager", big.NewInt(1), true)
	srvCrt := writeFile(t, dir, "srv.crt", srvCertPEM)
	srvKey := writeFile(t, dir, "srv.key", srvKeyPEM)
	cliCertPEM, cliKeyPEM := makeLeaf(t, caCert, caKey, "fleet-host", big.NewInt(0xdead), false)
	cliCert, _ := tls.X509KeyPair(cliCertPEM, cliKeyPEM)

	store, _ := newJSONStore("") // empty — the host is NOT enrolled
	hub := newAgentHub(store, nil)
	gs, err := agentServer(config{agentCert: srvCrt, agentKey: srvKey, agentClientCA: caPath}, hub)
	if err != nil {
		t.Fatalf("agentServer: %v", err)
	}
	lis, _ := net.Listen("tcp", "127.0.0.1:0")
	go gs.Serve(lis)
	defer gs.Stop()

	caPool := x509.NewCertPool()
	caPool.AddCert(caCert)
	creds := credentials.NewTLS(&tls.Config{Certificates: []tls.Certificate{cliCert}, RootCAs: caPool, ServerName: "localhost"})
	conn, _ := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(creds))
	defer conn.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	stream, err := fleetpb.NewFleetAgentClient(conn).Session(ctx)
	if err != nil {
		t.Fatalf("session open: %v", err)
	}
	if _, err := stream.Recv(); err == nil {
		t.Fatal("expected the stream to be rejected for an unenrolled host, got nil error")
	}
}

// ---- test helpers ---------------------------------------------------------

func makeCA(t *testing.T) (*x509.Certificate, *ecdsa.PrivateKey) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("ca key: %v", err)
	}
	tmpl := &x509.Certificate{
		SerialNumber:          big.NewInt(1),
		Subject:               pkix.Name{CommonName: "test-fleet-ca"},
		NotBefore:             time.Now().Add(-time.Hour),
		NotAfter:              time.Now().Add(time.Hour),
		IsCA:                  true,
		KeyUsage:              x509.KeyUsageCertSign | x509.KeyUsageDigitalSignature,
		BasicConstraintsValid: true,
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
	if err != nil {
		t.Fatalf("ca cert: %v", err)
	}
	cert, _ := x509.ParseCertificate(der)
	return cert, key
}

func makeLeaf(t *testing.T, ca *x509.Certificate, caKey *ecdsa.PrivateKey, cn string, serial *big.Int, server bool) (certPEM, keyPEM []byte) {
	t.Helper()
	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		t.Fatalf("leaf key: %v", err)
	}
	tmpl := &x509.Certificate{
		SerialNumber: serial,
		Subject:      pkix.Name{CommonName: cn},
		NotBefore:    time.Now().Add(-time.Hour),
		NotAfter:     time.Now().Add(time.Hour),
		KeyUsage:     x509.KeyUsageDigitalSignature,
		ExtKeyUsage:  []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth},
	}
	if server {
		tmpl.ExtKeyUsage = append(tmpl.ExtKeyUsage, x509.ExtKeyUsageServerAuth)
		tmpl.DNSNames = []string{"localhost"}
		tmpl.IPAddresses = []net.IP{net.ParseIP("127.0.0.1")}
	}
	der, err := x509.CreateCertificate(rand.Reader, tmpl, ca, &key.PublicKey, caKey)
	if err != nil {
		t.Fatalf("leaf cert: %v", err)
	}
	kb, _ := x509.MarshalECPrivateKey(key)
	return pemCert(der), pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: kb})
}

func pemCert(der []byte) []byte {
	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
}

func writeFile(t *testing.T, dir, name string, b []byte) string {
	t.Helper()
	p := filepath.Join(dir, name)
	if err := os.WriteFile(p, b, 0o600); err != nil {
		t.Fatalf("write %s: %v", name, err)
	}
	return p
}

func waitFor(t *testing.T, d time.Duration, cond func() bool, msg string) {
	t.Helper()
	deadline := time.Now().Add(d)
	for time.Now().Before(deadline) {
		if cond() {
			return
		}
		time.Sleep(20 * time.Millisecond)
	}
	t.Fatal(msg)
}