// Package registry resolves DSSE keyids to Ed25519 public keys so the
// chain verifier can check approval and transform signatures without
// trusting the storage layer that produced them.
//
// Phase 1 ships a Static keyring loaded from the environment / config.
// A later subtask swaps in a real Agent-Registry HTTP client that pulls
// the active public key for an agent + period (`trust-api/agent-<slug>
// -<yyyy>-q<n>`). The Keyring interface is the swap point.
package registry

import (
	"context"
	"crypto/ed25519"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"errors"
	"fmt"
	"strings"
	"sync"
)

// Keyring is the lookup surface the chain verifier depends on. Lookup
// is per-request so a future HTTP-backed Registry can refresh keys
// without a restart.
type Keyring interface {
	Lookup(ctx context.Context, keyID string) (ed25519.PublicKey, error)
}

// ErrUnknownKey means the keyring has no public key registered for the
// requested keyid. The HTTP layer maps this to a chain-rejection with
// reason `signature_unknown_key`.
var ErrUnknownKey = errors.New("registry: unknown keyid")

// Static is an in-memory Keyring backed by a mutex-protected map. Suitable
// for Phase 1 where the two signer keys (approval, transform) are known
// at startup. The map is cloned on every Register call so a running
// lookup never races a mutation.
type Static struct {
	mu   sync.RWMutex
	keys map[string]ed25519.PublicKey
}

// NewStatic returns an empty Static keyring.
func NewStatic() *Static {
	return &Static{keys: map[string]ed25519.PublicKey{}}
}

// Register adds (or overwrites) keyID -> pub.
func (s *Static) Register(keyID string, pub ed25519.PublicKey) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.keys[keyID] = pub
}

// Lookup returns the registered key or ErrUnknownKey.
func (s *Static) Lookup(_ context.Context, keyID string) (ed25519.PublicKey, error) {
	s.mu.RLock()
	defer s.mu.RUnlock()
	pub, ok := s.keys[keyID]
	if !ok {
		return nil, fmt.Errorf("%w: %q", ErrUnknownKey, keyID)
	}
	return pub, nil
}

// Snapshot returns a copy of the keyring as a `dsse.Keyring`-shaped map so
// callers can hand it to a verifier that wants the full map at once.
// This avoids repeated map lookups during a single chain verification.
func (s *Static) Snapshot() map[string]ed25519.PublicKey {
	s.mu.RLock()
	defer s.mu.RUnlock()
	out := make(map[string]ed25519.PublicKey, len(s.keys))
	for k, v := range s.keys {
		out[k] = v
	}
	return out
}

// LoadFromSpec populates s from a `;`-separated spec of entries of the
// form "<keyid>=<encoding>:<data>". Supported encodings:
//
//	pem   — PKIX-wrapped Ed25519 public key (PEM-encoded)
//	hex   — raw 32-byte Ed25519 public key, hex
//	b64   — raw 32-byte Ed25519 public key, standard base64
//	b64u  — raw 32-byte Ed25519 public key, base64url (no padding)
//
// Entries may also be separated by a comma, provided the entry's data
// does not itself contain a comma (safe for pem/hex/b64/b64u). Newlines
// inside a PEM payload are preserved. Blank entries are ignored. Keys
// with duplicate ids overwrite earlier entries.
func (s *Static) LoadFromSpec(spec string) error {
	for _, raw := range splitEntries(spec) {
		entry := strings.TrimSpace(raw)
		if entry == "" {
			continue
		}
		eq := strings.IndexByte(entry, '=')
		if eq <= 0 || eq == len(entry)-1 {
			return fmt.Errorf("registry: malformed key spec entry %q (want keyid=<encoding>:<data>)", entry)
		}
		keyID := strings.TrimSpace(entry[:eq])
		enc, data, ok := strings.Cut(entry[eq+1:], ":")
		if !ok {
			return fmt.Errorf("registry: malformed key spec entry %q (missing <encoding>: prefix)", entry)
		}
		pub, err := decodePublicKey(enc, data)
		if err != nil {
			return fmt.Errorf("registry: decode %q: %w", keyID, err)
		}
		s.Register(keyID, pub)
	}
	return nil
}

// splitEntries splits a spec on `;` or `,`. Newlines are NOT separators
// so a PEM block (which contains newlines internally) survives
// round-tripping through this function.
func splitEntries(s string) []string {
	return strings.FieldsFunc(s, func(r rune) bool { return r == ';' || r == ',' })
}

func decodePublicKey(encoding, data string) (ed25519.PublicKey, error) {
	switch strings.ToLower(strings.TrimSpace(encoding)) {
	case "pem":
		return decodePEM(data)
	case "hex":
		return decodeFixedRaw(data, decodeHex)
	case "b64":
		return decodeFixedRaw(data, decodeB64)
	case "b64u":
		return decodeFixedRaw(data, decodeB64URL)
	default:
		return nil, fmt.Errorf("unsupported encoding %q (want pem|hex|b64|b64u)", encoding)
	}
}

type rawDecoder func(string) ([]byte, error)

func decodeFixedRaw(data string, dec rawDecoder) (ed25519.PublicKey, error) {
	raw, err := dec(data)
	if err != nil {
		return nil, err
	}
	if len(raw) != ed25519.PublicKeySize {
		return nil, fmt.Errorf("expected %d-byte Ed25519 public key, got %d", ed25519.PublicKeySize, len(raw))
	}
	return ed25519.PublicKey(raw), nil
}

func decodePEM(data string) (ed25519.PublicKey, error) {
	// PEM headers may have been shell-escaped as "\n"; accept either.
	normalised := strings.ReplaceAll(data, `\n`, "\n")
	block, _ := pem.Decode([]byte(normalised))
	if block == nil {
		return nil, errors.New("not a PEM block")
	}
	pub, err := x509.ParsePKIXPublicKey(block.Bytes)
	if err != nil {
		return nil, fmt.Errorf("parse PKIX: %w", err)
	}
	ed, ok := pub.(ed25519.PublicKey)
	if !ok {
		return nil, fmt.Errorf("not an Ed25519 public key (got %T)", pub)
	}
	return ed, nil
}

func decodeHex(data string) ([]byte, error) {
	return hexDecode(strings.TrimSpace(data))
}

func decodeB64(data string) ([]byte, error) {
	return base64.StdEncoding.DecodeString(strings.TrimSpace(data))
}

func decodeB64URL(data string) ([]byte, error) {
	return base64.RawURLEncoding.DecodeString(strings.TrimSpace(data))
}

// hexDecode isolated so decodeFixedRaw can share the shape of its callers.
func hexDecode(s string) ([]byte, error) {
	n := len(s)
	if n%2 != 0 {
		return nil, errors.New("odd-length hex string")
	}
	out := make([]byte, n/2)
	for i := 0; i < n; i += 2 {
		hi, err := hexNibble(s[i])
		if err != nil {
			return nil, err
		}
		lo, err := hexNibble(s[i+1])
		if err != nil {
			return nil, err
		}
		out[i/2] = hi<<4 | lo
	}
	return out, nil
}

func hexNibble(c byte) (byte, error) {
	switch {
	case '0' <= c && c <= '9':
		return c - '0', nil
	case 'a' <= c && c <= 'f':
		return c - 'a' + 10, nil
	case 'A' <= c && c <= 'F':
		return c - 'A' + 10, nil
	}
	return 0, fmt.Errorf("invalid hex char %q", c)
}