// enrol.go — the enrolment transaction: build a CSR, submit it to the broker,
// show the returned fingerprint, poll until an admin approves, return the cert.
// TPM-agnostic (takes a crypto.Signer) so it is fully unit-testable with a
// software key against a mock broker; the real signer is TPM-backed (tpm.go).
package main

import (
	"bytes"
	"context"
	"crypto"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/json"
	"encoding/pem"
	"fmt"
	"io"
	"net/http"
	"time"
)

// --- broker wire types (mirror serving/enrol-broker) ------------------------

type deviceContext struct {
	MAC       string `json:"mac,omitempty"`
	DMIUUID   string `json:"dmi_uuid,omitempty"`
	DMISerial string `json:"dmi_serial,omitempty"`
	EKHash    string `json:"ek_hash,omitempty"`
}

type enrolRequest struct {
	CSR     string        `json:"csr"`
	Context deviceContext `json:"context"`
}

type requestResp struct {
	EnrolID     string `json:"enrol_id"`
	Fingerprint string `json:"fingerprint"`
	State       string `json:"state"`
}

type statusResp struct {
	State      string `json:"state"`
	Cert       string `json:"cert,omitempty"`
	Hostnqn    string `json:"hostnqn,omitempty"`
	CertSerial string `json:"cert_serial,omitempty"`
}

// buildCSR produces a PEM CSR signed by signer (the TPM key). The CN is generic
// ("fleet-host") by convention — RADIUS derives real identity from the cert, and
// the per-host NVMe hostnqn is the cert serial; nothing per-host is embedded here.
func buildCSR(signer crypto.Signer, cn string) ([]byte, error) {
	tmpl := &x509.CertificateRequest{
		Subject: pkix.Name{CommonName: cn, Organization: []string{"fleet"}},
	}
	der, err := x509.CreateCertificateRequest(nil, tmpl, signer)
	if err != nil {
		return nil, fmt.Errorf("create CSR: %w", err)
	}
	return pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE REQUEST", Bytes: der}), nil
}

type brokerClient struct {
	http    *http.Client
	baseURL string // e.g. http://boot.fleet.lab/enrol
}

func (c *brokerClient) submit(csrPEM []byte, dctx deviceContext) (*requestResp, error) {
	body, _ := json.Marshal(enrolRequest{CSR: string(csrPEM), Context: dctx})
	resp, err := c.http.Post(c.baseURL+"/request", "application/json", bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK && resp.StatusCode != http.StatusAccepted {
		return nil, fmt.Errorf("broker /request: %s: %s", resp.Status, bytes.TrimSpace(b))
	}
	var r requestResp
	if err := json.Unmarshal(b, &r); err != nil {
		return nil, fmt.Errorf("broker /request: bad JSON: %w", err)
	}
	if r.EnrolID == "" {
		return nil, fmt.Errorf("broker /request: no enrol_id in response")
	}
	return &r, nil
}

func (c *brokerClient) status(id string) (*statusResp, error) {
	resp, err := c.http.Get(c.baseURL + "/status/" + id)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("broker /status: %s: %s", resp.Status, bytes.TrimSpace(b))
	}
	var s statusResp
	if err := json.Unmarshal(b, &s); err != nil {
		return nil, fmt.Errorf("broker /status: bad JSON: %w", err)
	}
	return &s, nil
}

// enrol runs the full transaction: submit -> show fingerprint -> poll to a
// terminal state. onFingerprint is called once with the code to display; the
// returned statusResp is APPROVED (with cert) or an error for denied/expired/timeout.
func enrol(ctx context.Context, c *brokerClient, csrPEM []byte, dctx deviceContext,
	pollEvery time.Duration, onFingerprint func(string)) (*statusResp, error) {

	req, err := c.submit(csrPEM, dctx)
	if err != nil {
		return nil, err
	}
	if onFingerprint != nil {
		onFingerprint(req.Fingerprint)
	}

	tick := time.NewTicker(pollEvery)
	defer tick.Stop()
	for {
		s, err := c.status(req.EnrolID)
		if err == nil {
			switch s.State {
			case "approved":
				if s.Cert == "" {
					return nil, fmt.Errorf("broker approved but returned no cert")
				}
				return s, nil
			case "denied":
				return nil, fmt.Errorf("enrolment denied by admin")
			case "expired":
				return nil, fmt.Errorf("enrolment request expired")
			}
		}
		// pending (or a transient status error) — keep waiting until the deadline.
		select {
		case <-ctx.Done():
			return nil, fmt.Errorf("enrolment not approved before timeout")
		case <-tick.C:
		}
	}
}

// leafDER returns the DER of the FIRST certificate in a (possibly chained) PEM
// bundle — the leaf, which is what wpa_supplicant presents and the hostnqn
// derives from. Also validates it parses.
func leafDER(bundlePEM string) ([]byte, error) {
	rest := []byte(bundlePEM)
	for {
		var blk *pem.Block
		blk, rest = pem.Decode(rest)
		if blk == nil {
			return nil, fmt.Errorf("no CERTIFICATE block in broker response")
		}
		if blk.Type == "CERTIFICATE" {
			if _, err := x509.ParseCertificate(blk.Bytes); err != nil {
				return nil, fmt.Errorf("issued cert does not parse: %w", err)
			}
			return blk.Bytes, nil
		}
	}
}