package main

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"math/big"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

func bigOne() *big.Int           { return big.NewInt(1) }
func bigN(n int64) *big.Int      { return big.NewInt(n) }
func pkixCN(cn string) pkix.Name { return pkix.Name{CommonName: cn} }

// mockBroker mimics serving/enrol-broker: /request returns a fingerprint +
// enrol-id; /status returns pending until approveAfter polls, then approved with
// a cert signed by an in-process CA over the submitted CSR's public key.
func mockBroker(t *testing.T, approveAfterPolls int) *httptest.Server {
	t.Helper()
	caKey, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	caTmpl := &x509.Certificate{SerialNumber: bigOne(), Subject: pkixCN("mock CA"), IsCA: true, BasicConstraintsValid: true}
	caDER, _ := x509.CreateCertificate(rand.Reader, caTmpl, caTmpl, &caKey.PublicKey, caKey)
	caCert, _ := x509.ParseCertificate(caDER)

	var polls int
	var certPEM string
	mux := http.NewServeMux()
	mux.HandleFunc("/request", func(w http.ResponseWriter, r *http.Request) {
		var req enrolRequest
		json.NewDecoder(r.Body).Decode(&req)
		blk, _ := pem.Decode([]byte(req.CSR))
		if blk == nil {
			t.Errorf("submitted CSR is not PEM")
			http.Error(w, "bad csr", 400)
			return
		}
		csr, err := x509.ParseCertificateRequest(blk.Bytes)
		if err != nil || csr.CheckSignature() != nil {
			t.Errorf("CSR signature invalid: %v", err)
			http.Error(w, "bad csr sig", 400)
			return
		}
		// sign a leaf over the CSR's public key
		leaf := &x509.Certificate{SerialNumber: bigN(0xabc123), Subject: csr.Subject}
		der, _ := x509.CreateCertificate(rand.Reader, leaf, caCert, csr.PublicKey, caKey)
		certPEM = string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})) +
			string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: caDER}))
		json.NewEncoder(w).Encode(requestResp{EnrolID: "e1", Fingerprint: "cedar-onyx-meadow", State: "pending"})
	})
	mux.HandleFunc("/status/", func(w http.ResponseWriter, r *http.Request) {
		polls++
		if polls <= approveAfterPolls {
			json.NewEncoder(w).Encode(statusResp{State: "pending"})
			return
		}
		json.NewEncoder(w).Encode(statusResp{
			State: "approved", Cert: certPEM,
			Hostnqn: "nqn.2026-06.org.example.fleet:host-abc123", CertSerial: "abc123",
		})
	})
	return httptest.NewServer(mux)
}

func TestEnrolHappyPath(t *testing.T) {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader) // stands in for the TPM signer
	csrPEM, err := buildCSR(key, "fleet-host")
	if err != nil {
		t.Fatalf("buildCSR: %v", err)
	}
	srv := mockBroker(t, 2) // pending twice, then approved
	defer srv.Close()

	c := &brokerClient{http: srv.Client(), baseURL: srv.URL}
	var gotFP string
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	st, err := enrol(ctx, c, csrPEM, deviceContext{MAC: "de:ad:be:ef:00:01"}, 10*time.Millisecond,
		func(fp string) { gotFP = fp })
	if err != nil {
		t.Fatalf("enrol: %v", err)
	}
	if gotFP != "cedar-onyx-meadow" {
		t.Errorf("fingerprint not surfaced: %q", gotFP)
	}
	der, err := leafDER(st.Cert)
	if err != nil {
		t.Fatalf("leafDER: %v", err)
	}
	crt, err := x509.ParseCertificate(der)
	if err != nil {
		t.Fatalf("parse leaf: %v", err)
	}
	if crt.Subject.CommonName != "fleet-host" {
		t.Errorf("leaf CN = %q, want fleet-host", crt.Subject.CommonName)
	}
}

func TestEnrolDenied(t *testing.T) {
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	csrPEM, _ := buildCSR(key, "fleet-host")
	mux := http.NewServeMux()
	mux.HandleFunc("/request", func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(requestResp{EnrolID: "e1", Fingerprint: "x", State: "pending"})
	})
	mux.HandleFunc("/status/", func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(statusResp{State: "denied"})
	})
	srv := httptest.NewServer(mux)
	defer srv.Close()
	c := &brokerClient{http: srv.Client(), baseURL: srv.URL}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	_, err := enrol(ctx, c, csrPEM, deviceContext{}, time.Millisecond, nil)
	if err == nil || !strings.Contains(err.Error(), "denied") {
		t.Fatalf("expected denied error, got %v", err)
	}
}

func TestLeafDERChain(t *testing.T) {
	// two-cert bundle -> leaf is the first
	key, _ := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	mk := func(cn string) string {
		tmpl := &x509.Certificate{SerialNumber: bigOne(), Subject: pkixCN(cn)}
		der, _ := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
		return string(pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der}))
	}
	der, err := leafDER(mk("leaf") + mk("inter"))
	if err != nil {
		t.Fatalf("leafDER: %v", err)
	}
	crt, _ := x509.ParseCertificate(der)
	if crt.Subject.CommonName != "leaf" {
		t.Errorf("got %q, want leaf", crt.Subject.CommonName)
	}
}