// store.go — the in-memory enrolment state machine + the OOB fingerprint.
//
// The store holds short-lived PENDING records keyed by a random enrol-id. State
// transitions are the load-bearing lifecycle from docs/tpm-enrolment-design.md:
//
//	PENDING --(admin approve)--> APPROVED   (cert minted, returned on poll)
//	PENDING --(admin deny)-----> DENIED
//	PENDING --(ttl elapsed)----> EXPIRED
//
// Records are time-boxed (a flooding bound) and evicted after a retention window.
// Nothing here is durable: a broker restart drops pending requests, which is
// acceptable — an un-enrolled host simply re-requests on its next poll/boot.
package main

import (
	"crypto/rand"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"math/bits"
	"strings"
	"sync"
	"time"
)

type State string

const (
	StatePending  State = "pending"
	StateApproved State = "approved"
	StateDenied   State = "denied"
	StateExpired  State = "expired"
)

// DeviceContext is what an un-enrolled host self-reports so the admin knows WHAT
// they are approving. It is advisory only — the fingerprint is the cryptographic
// anchor that binds the queue entry to the device on the console.
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"`
	Extra     map[string]string `json:"extra,omitempty"`
}

// Enrolment is one pending/decided enrolment. Unexported fields (nonce, csrPEM,
// certPEM) are deliberately NOT serialised into the admin queue view; the cert is
// only ever handed back to the requesting host via the status poll.
type Enrolment struct {
	ID          string        `json:"id"`
	State       State         `json:"state"`
	Fingerprint string        `json:"fingerprint"`
	Context     DeviceContext `json:"context"`
	PubKeyHash  string        `json:"pubkey_hash"`
	CN          string        `json:"cn"`
	CreatedAt   time.Time     `json:"created_at"`
	ExpiresAt   time.Time     `json:"expires_at"`
	DecidedAt   time.Time     `json:"decided_at,omitempty"`
	DecidedBy   string        `json:"decided_by,omitempty"`
	CertSerial  string        `json:"cert_serial,omitempty"`
	Hostnqn     string        `json:"hostnqn,omitempty"`

	nonce   []byte
	csrPEM  string
	certPEM string
}

// Age returns how long the record has existed (for the UI).
func (e *Enrolment) Age() time.Duration { return time.Since(e.CreatedAt) }

type Store struct {
	mu         sync.Mutex
	items      map[string]*Enrolment
	ttl        time.Duration
	retention  time.Duration
	maxPending int

	words    []string
	fpWords  int
	hostBase string
}

func NewStore(ttl, retention time.Duration, maxPending int, words []string, fpWords int, hostBase string) (*Store, error) {
	if !isPow2(len(words)) {
		return nil, fmt.Errorf("wordlist size %d is not a power of two (needed for whole bits/word)", len(words))
	}
	if fpWords < 4 {
		return nil, fmt.Errorf("fingerprint word count %d too small", fpWords)
	}
	return &Store{
		items:      make(map[string]*Enrolment),
		ttl:        ttl,
		retention:  retention,
		maxPending: maxPending,
		words:      words,
		fpWords:    fpWords,
		hostBase:   hostBase,
	}, nil
}

var errTooManyPending = errors.New("too many pending enrolments")

// Add creates a PENDING record for a parsed CSR's public-key DER + CN + context.
func (s *Store) Add(pubDER []byte, cn string, csrPEM string, ctx DeviceContext) (*Enrolment, error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.countPendingLocked() >= s.maxPending {
		return nil, errTooManyPending
	}

	id, err := randHex(16)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, 16)
	if _, err := rand.Read(nonce); err != nil {
		return nil, err
	}
	fp, err := fingerprint(pubDER, nonce, s.words, s.fpWords)
	if err != nil {
		return nil, err
	}
	sum := sha256.Sum256(pubDER)
	now := time.Now()
	e := &Enrolment{
		ID:          id,
		State:       StatePending,
		Fingerprint: fp,
		Context:     ctx,
		PubKeyHash:  hex.EncodeToString(sum[:]),
		CN:          cn,
		CreatedAt:   now,
		ExpiresAt:   now.Add(s.ttl),
		nonce:       nonce,
		csrPEM:      csrPEM,
	}
	s.items[id] = e
	return e.copy(), nil
}

// Get returns a copy of a record (or nil), refreshing EXPIRED lazily.
func (s *Store) Get(id string) *Enrolment {
	s.mu.Lock()
	defer s.mu.Unlock()
	e := s.items[id]
	if e == nil {
		return nil
	}
	s.expireLocked(e)
	return e.copy()
}

// csrFor returns the stored CSR PEM for a still-pending id (for the signer).
func (s *Store) csrFor(id string) (string, bool) {
	s.mu.Lock()
	defer s.mu.Unlock()
	e := s.items[id]
	if e == nil {
		return "", false
	}
	s.expireLocked(e)
	if e.State != StatePending {
		return "", false
	}
	return e.csrPEM, true
}

// Approve attaches the issued cert and flips the record to APPROVED. The signer
// runs OUTSIDE the lock (it shells out to step) so we re-check state here.
func (s *Store) Approve(id, certPEM, serial, admin string) (*Enrolment, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	e := s.items[id]
	if e == nil {
		return nil, errors.New("unknown enrol-id")
	}
	s.expireLocked(e)
	if e.State != StatePending {
		return nil, fmt.Errorf("enrol-id is %s, not pending", e.State)
	}
	e.State = StateApproved
	e.certPEM = certPEM
	e.CertSerial = serial
	e.Hostnqn = s.hostBase + serial
	e.DecidedAt = time.Now()
	e.DecidedBy = admin
	return e.copy(), nil
}

// Deny rejects a pending record.
func (s *Store) Deny(id, admin string) (*Enrolment, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	e := s.items[id]
	if e == nil {
		return nil, errors.New("unknown enrol-id")
	}
	s.expireLocked(e)
	if e.State != StatePending {
		return nil, fmt.Errorf("enrol-id is %s, not pending", e.State)
	}
	e.State = StateDenied
	e.DecidedAt = time.Now()
	e.DecidedBy = admin
	return e.copy(), nil
}

// Pending returns copies of the still-pending records, oldest first.
func (s *Store) Pending() []*Enrolment {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*Enrolment
	for _, e := range s.items {
		s.expireLocked(e)
		if e.State == StatePending {
			out = append(out, e.copy())
		}
	}
	sortByCreated(out)
	return out
}

// Reap expires timed-out pending records and evicts old decided ones. Returns the
// number expired (for audit logging).
func (s *Store) Reap() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	expired := 0
	for id, e := range s.items {
		before := e.State
		s.expireLocked(e)
		if before == StatePending && e.State == StateExpired {
			expired++
		}
		if e.State != StatePending && time.Since(e.decidedOrExpired()) > s.retention {
			delete(s.items, id)
		}
	}
	return expired
}

func (s *Store) countPendingLocked() int {
	n := 0
	for _, e := range s.items {
		if e.State == StatePending && time.Now().Before(e.ExpiresAt) {
			n++
		}
	}
	return n
}

func (s *Store) expireLocked(e *Enrolment) {
	if e.State == StatePending && !time.Now().Before(e.ExpiresAt) {
		e.State = StateExpired
		e.DecidedAt = e.ExpiresAt
	}
}

func (e *Enrolment) decidedOrExpired() time.Time {
	if !e.DecidedAt.IsZero() {
		return e.DecidedAt
	}
	return e.ExpiresAt
}

func (e *Enrolment) copy() *Enrolment {
	c := *e
	return &c
}

// --- fingerprint -----------------------------------------------------------

// fingerprint deterministically maps (CSR public key, server nonce) to a short
// word sequence. Both broker and host can compute it identically. The nonce is
// server-chosen per enrol-id so an attacker cannot pre-grind a key whose
// fingerprint collides with a target device's. Words (not hex) keep it short to
// read aloud yet resistant to casual collision. Swap wordlist.txt for the 2048-
// word BIP39 list (11 bits/word) if more entropy per word is wanted.
func fingerprint(pubDER, nonce []byte, words []string, n int) (string, error) {
	if !isPow2(len(words)) {
		return "", fmt.Errorf("wordlist size %d not a power of two", len(words))
	}
	bpw := bits.TrailingZeros(uint(len(words)))
	stream := digestStream(pubDER, nonce, (n*bpw+7)/8)
	out := make([]string, 0, n)
	bitpos := 0
	for range n {
		idx := 0
		for range bpw {
			byteIdx := bitpos / 8
			bit := 7 - uint(bitpos%8)
			b := (stream[byteIdx] >> bit) & 1
			idx = (idx << 1) | int(b)
			bitpos++
		}
		out = append(out, words[idx])
	}
	return strings.Join(out, "-"), nil
}

// digestStream is a SHA-256 counter-mode KDF producing >= n bytes deterministically.
func digestStream(a, b []byte, n int) []byte {
	var stream []byte
	var ctr uint32
	for len(stream) < n {
		h := sha256.New()
		h.Write(a)
		h.Write(b)
		var c [4]byte
		binary.BigEndian.PutUint32(c[:], ctr)
		h.Write(c[:])
		stream = h.Sum(stream)
		ctr++
	}
	return stream
}

func isPow2(n int) bool { return n > 0 && n&(n-1) == 0 }

func randHex(nbytes int) (string, error) {
	b := make([]byte, nbytes)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

// sortByCreated is a tiny insertion sort (oldest first); the queue is small.
func sortByCreated(xs []*Enrolment) {
	for i := 1; i < len(xs); i++ {
		for j := i; j > 0 && xs[j-1].CreatedAt.After(xs[j].CreatedAt); j-- {
			xs[j-1], xs[j] = xs[j], xs[j-1]
		}
	}
}