// Package capsule implements the DateKeyCap .dkc container (spec §20-§39):
// framing, PUBLIC_HEADER, CONTROL_CBOR, header_binding, the time_only and
// time_and_key constructions, and the encryption (spec §61, §62) and
// decryption (spec §63) flows.
//
// A .dkc is PRELUDE || PUBLIC_HEADER || SEALED_CONTROL || PAYLOAD_AGE, where
// SEALED_CONTROL and PAYLOAD_AGE are complete standard age files and the
// payload runs to EOF (spec §22, §28-§34).
package capsule

import (
	"bytes"
	"crypto/sha256"
	"encoding/binary"
	"encoding/hex"
	"fmt"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/codec"
	"g.activething.com/go/DateKeys/datekey"
	"g.activething.com/go/DateKeys/extension"
)

// Framing constants (spec §22, §23) and parser limits (spec §57).
const (
	Magic          = "DKC1"
	FramingVersion = 1
	PreludeSize    = 16

	MaxPublicHeaderLen  = 1 << 20  // 1 MiB
	MaxSealedControlLen = 64 << 20 // 64 MiB

	HeaderTypeTag  = "datekeycap"
	HeaderVersion  = 1
	ControlTypeTag = "datekeys-control"
	ControlVersion = 1

	CapsuleIDSize = 16
)

// Policy is the declared access policy of PUBLIC_HEADER (spec §25).
type Policy uint8

// Access policies of V1.
const (
	TimeOnly   Policy = 0
	TimeAndKey Policy = 1
)

func (p Policy) String() string {
	switch p {
	case TimeOnly:
		return "time_only"
	case TimeAndKey:
		return "time_and_key"
	}
	return fmt.Sprintf("policy(%d)", uint8(p))
}

// ParsePolicy parses "time_only" or "time_and_key".
func ParsePolicy(s string) (Policy, error) {
	switch s {
	case "time_only":
		return TimeOnly, nil
	case "time_and_key":
		return TimeAndKey, nil
	}
	return 0, fmt.Errorf("capsule: unknown access policy %q", s)
}

func (p Policy) valid() bool { return p == TimeOnly || p == TimeAndKey }

// ---------------------------------------------------------------------------
// PRELUDE

// Prelude is the fixed 16-byte PRELUDE (spec §22, §23).
type Prelude struct {
	PublicHeaderLen  uint32
	SealedControlLen uint32
}

// Bytes returns the exact 16 prelude bytes, the ones covered by
// header_binding.
func (p Prelude) Bytes() [PreludeSize]byte {
	var b [PreludeSize]byte
	copy(b[0:4], Magic)
	b[4] = FramingVersion
	// FLAGS (b[5]) and RESERVED (b[6:8]) are zero in V1.
	binary.BigEndian.PutUint32(b[8:12], p.PublicHeaderLen)
	binary.BigEndian.PutUint32(b[12:16], p.SealedControlLen)
	return b
}

// PayloadOffset is where PAYLOAD_AGE starts:
// 16 + PUBLIC_HEADER_LEN + SEALED_CONTROL_LEN (spec §63).
func (p Prelude) PayloadOffset() int64 {
	return PreludeSize + int64(p.PublicHeaderLen) + int64(p.SealedControlLen)
}

// ParsePrelude validates the prelude (spec §22, §63 step 2): magic, version,
// FLAGS == 0, RESERVED == 0 and the length limits of spec §57.
func ParsePrelude(b []byte) (Prelude, error) {
	if len(b) < 4 || string(b[0:4]) != Magic {
		return Prelude{}, fmt.Errorf("capsule: %w", datekeys.ErrInvalidMagic)
	}
	if len(b) < PreludeSize {
		return Prelude{}, fmt.Errorf("capsule: truncated prelude: %w", datekeys.ErrIntegrity)
	}
	if b[4] != FramingVersion {
		return Prelude{}, fmt.Errorf("capsule: framing version %d: %w", b[4], datekeys.ErrUnsupportedVersion)
	}
	if b[5] != 0 || b[6] != 0 || b[7] != 0 {
		return Prelude{}, fmt.Errorf("capsule: flags %#x, reserved %#02x%02x: %w", b[5], b[6], b[7], datekeys.ErrInvalidFlags)
	}
	p := Prelude{
		PublicHeaderLen:  binary.BigEndian.Uint32(b[8:12]),
		SealedControlLen: binary.BigEndian.Uint32(b[12:16]),
	}
	if p.PublicHeaderLen == 0 || p.PublicHeaderLen > MaxPublicHeaderLen {
		return Prelude{}, fmt.Errorf("capsule: PUBLIC_HEADER_LEN %d outside 1..%d: %w", p.PublicHeaderLen, MaxPublicHeaderLen, datekeys.ErrIntegrity)
	}
	if p.SealedControlLen == 0 || p.SealedControlLen > MaxSealedControlLen {
		return Prelude{}, fmt.Errorf("capsule: SEALED_CONTROL_LEN %d outside 1..%d: %w", p.SealedControlLen, MaxSealedControlLen, datekeys.ErrIntegrity)
	}
	return p, nil
}

// HeaderBinding returns SHA-256(PRELUDE || PUBLIC_HEADER_BYTES) over the exact
// stored bytes; the header is never re-serialized for it (spec §26).
func HeaderBinding(prelude [PreludeSize]byte, publicHeader []byte) [32]byte {
	h := sha256.New()
	h.Write(prelude[:])
	h.Write(publicHeader)
	var out [32]byte
	h.Sum(out[:0])
	return out
}

// ---------------------------------------------------------------------------
// PUBLIC_HEADER

// Header is PUBLIC_HEADER (spec §24). There is no separate profile_id: the
// profile comes from the DateKey, the single source of truth.
type Header struct {
	CapsuleID   [CapsuleIDSize]byte   // key 2
	DateKey     datekey.DateKey       // key 3, canonical dk1_
	Policy      Policy                // key 4, access_policy
	Critical    []extension.Extension // key 5
	Noncritical []extension.Extension // key 6
}

// headerWire is PUBLIC_HEADER as it is encoded: keys 2 to 6, keys 0 and 1
// being the constants HeaderTypeTag and HeaderVersion.
type headerWire struct {
	CapsuleID   []byte                // key 2
	DateKey     string                // key 3
	Policy      uint64                // key 4
	Critical    []extension.Extension // key 5, omitted when empty
	Noncritical []extension.Extension // key 6, omitted when empty
}

func (w *headerWire) encode(e *codec.Encoder) {
	e.Map(5 + nonEmpty(w.Critical) + nonEmpty(w.Noncritical))
	e.Uint(0)
	e.Text(HeaderTypeTag)
	e.Uint(1)
	e.Uint(HeaderVersion)
	e.Uint(2)
	e.Bstr(w.CapsuleID)
	e.Uint(3)
	e.Text(w.DateKey)
	e.Uint(4)
	e.Uint(w.Policy)
	encodeExtensions(e, 5, w.Critical, w.Noncritical)
}

// decode reads PUBLIC_HEADER with every CDDL rule whose violation is
// ErrNonCanonicalCBOR, including the extension arrays; the DateKey, which
// has codes of its own (spec §57), is parsed afterwards.
func (w *headerWire) decode(d *codec.Decoder) error {
	pairs, err := d.Map(7)
	if err != nil {
		return err
	}
	var seen uint
	for range pairs {
		k, err := d.Key()
		if err != nil {
			return err
		}
		switch k {
		case 0:
			_, err = d.Text(len(HeaderTypeTag))
		case 1:
			_, err = d.Uint(HeaderVersion)
		case 2:
			w.CapsuleID, err = d.Bstr(CapsuleIDSize, CapsuleIDSize)
		case 3:
			w.DateKey, err = d.Text(MaxPublicHeaderLen)
		case 4:
			// Compared as read, before any narrowing to Policy, which would
			// let 256, 257, 2^32 and the like pass as a V1 policy.
			if w.Policy, err = d.Uint(codec.MaxSafeUint); err == nil && w.Policy > uint64(TimeAndKey) {
				err = fmt.Errorf("access_policy %d is not defined in V1: %w", w.Policy, datekeys.ErrNonCanonicalCBOR)
			}
		case 5:
			w.Critical, err = extension.DecodeArray(d)
		case 6:
			w.Noncritical, err = extension.DecodeArray(d)
		default:
			return fmt.Errorf("key %d is not defined: %w", k, datekeys.ErrNonCanonicalCBOR)
		}
		if err != nil {
			return fmt.Errorf("key %d: %w", k, err)
		}
		seen |= 1 << k
	}
	if err := required(seen, 5); err != nil {
		return err
	}
	return d.EndMap()
}

// nonEmpty is 1 for an extension array that is written and 0 for one that is
// omitted (spec §58.1).
func nonEmpty(exts []extension.Extension) int {
	if len(exts) == 0 {
		return 0
	}
	return 1
}

// encodeExtensions writes the critical and noncritical arrays at keys key and
// key+1, each only when it is not empty.
func encodeExtensions(e *codec.Encoder, key uint64, critical, noncritical []extension.Extension) {
	for i, exts := range [][]extension.Extension{critical, noncritical} {
		if len(exts) > 0 {
			e.Uint(key + uint64(i))
			extension.EncodeArray(e, exts)
		}
	}
}

// required checks that seen holds the keys 0 to n-1, which are required.
func required(seen uint, n int) error {
	for k := range n {
		if seen&(1<<k) == 0 {
			return fmt.Errorf("key %d is missing: %w", k, datekeys.ErrNonCanonicalCBOR)
		}
	}
	return nil
}

// CapsuleIDHex returns the capsule_id in hexadecimal.
func (h *Header) CapsuleIDHex() string { return hex.EncodeToString(h.CapsuleID[:]) }

// EncodeHeader returns the Deterministic CBOR bytes of h, at most
// MaxPublicHeaderLen of them (spec §57).
func EncodeHeader(h *Header) ([]byte, error) {
	compact := h.DateKey.Compact()
	if compact == "" {
		return nil, fmt.Errorf("capsule: invalid DateKey: %w", datekeys.ErrDateKeyInvalid)
	}
	if !h.Policy.valid() {
		return nil, fmt.Errorf("capsule: unknown access policy %d", h.Policy)
	}
	w := headerWire{
		CapsuleID: h.CapsuleID[:],
		DateKey:   compact,
		Policy:    uint64(h.Policy),
	}
	var err error
	if w.Critical, err = extension.Canonical(h.Critical); err != nil {
		return nil, err
	}
	if w.Noncritical, err = extension.Canonical(h.Noncritical); err != nil {
		return nil, err
	}
	if err := extension.CheckDisjoint(w.Critical, w.Noncritical); err != nil {
		return nil, err
	}
	var e codec.Encoder
	w.encode(&e)
	b, err := e.Out()
	if err != nil {
		return nil, err
	}
	if len(b) > MaxPublicHeaderLen {
		return nil, fmt.Errorf("capsule: PUBLIC_HEADER of %d bytes exceeds %d: %w", len(b), MaxPublicHeaderLen, datekeys.ErrIntegrity)
	}
	return b, nil
}

// DecodeHeader validates and decodes PUBLIC_HEADER bytes (spec §24, §27,
// §63 step 4): the §57 limit, the schema version, canonical CBOR, the schema
// with a 16-byte capsule_id, a V1 access policy and well-formed extension
// arrays, and then a canonical DateKey, so that a header that also breaks the
// CDDL reports ErrNonCanonicalCBOR. Whether the profile is pinned and the
// critical extensions known is decided by the caller.
func DecodeHeader(b []byte) (*Header, error) {
	if len(b) > MaxPublicHeaderLen {
		return nil, fmt.Errorf("capsule: PUBLIC_HEADER of %d bytes exceeds %d: %w", len(b), MaxPublicHeaderLen, datekeys.ErrIntegrity)
	}
	if err := codec.CheckSchema(b, HeaderTypeTag, HeaderVersion); err != nil {
		return nil, fmt.Errorf("capsule: PUBLIC_HEADER: %w", err)
	}
	var w headerWire
	if err := codec.Unmarshal(b, w.decode, w.encode); err != nil {
		return nil, fmt.Errorf("capsule: PUBLIC_HEADER: %w", err)
	}
	if err := extension.CheckDisjoint(w.Critical, w.Noncritical); err != nil {
		return nil, fmt.Errorf("capsule: PUBLIC_HEADER: %w", err)
	}
	dk, err := datekey.Parse(w.DateKey)
	if err != nil {
		return nil, fmt.Errorf("capsule: PUBLIC_HEADER: %w", err)
	}
	// decode bounds w.Policy to 0 or 1, so the conversion is exact.
	h := &Header{DateKey: dk, Policy: Policy(w.Policy), Critical: w.Critical, Noncritical: w.Noncritical}
	copy(h.CapsuleID[:], w.CapsuleID)
	return h, nil
}

// ---------------------------------------------------------------------------
// CONTROL_CBOR

// Control is CONTROL_CBOR (spec §31). PayloadIdentity is I_PAYLOAD, a secret.
type Control struct {
	HeaderBinding   [32]byte              // key 2
	PayloadIdentity [32]byte              // key 3, raw X25519 identity bytes. SECRET.
	Critical        []extension.Extension // key 4
	Noncritical     []extension.Extension // key 5
}

// controlWire is CONTROL_CBOR as it is encoded: keys 2 to 5, keys 0 and 1
// being the constants ControlTypeTag and ControlVersion.
type controlWire struct {
	HeaderBinding   []byte                // key 2
	PayloadIdentity []byte                // key 3, SECRET
	Critical        []extension.Extension // key 4, omitted when empty
	Noncritical     []extension.Extension // key 5, omitted when empty
}

func (w *controlWire) encode(e *codec.Encoder) {
	e.Map(4 + nonEmpty(w.Critical) + nonEmpty(w.Noncritical))
	e.Uint(0)
	e.Text(ControlTypeTag)
	e.Uint(1)
	e.Uint(ControlVersion)
	e.Uint(2)
	e.Bstr(w.HeaderBinding)
	e.Uint(3)
	e.Bstr(w.PayloadIdentity)
	encodeExtensions(e, 4, w.Critical, w.Noncritical)
}

// decode reads CONTROL_CBOR with every CDDL rule. The caller wipes
// PayloadIdentity, whatever the result.
func (w *controlWire) decode(d *codec.Decoder) error {
	pairs, err := d.Map(6)
	if err != nil {
		return err
	}
	var seen uint
	for range pairs {
		k, err := d.Key()
		if err != nil {
			return err
		}
		switch k {
		case 0:
			_, err = d.Text(len(ControlTypeTag))
		case 1:
			_, err = d.Uint(ControlVersion)
		case 2:
			w.HeaderBinding, err = d.Bstr(32, 32)
		case 3:
			w.PayloadIdentity, err = d.Bstr(32, 32)
		case 4:
			w.Critical, err = extension.DecodeArray(d)
		case 5:
			w.Noncritical, err = extension.DecodeArray(d)
		default:
			return fmt.Errorf("key %d is not defined: %w", k, datekeys.ErrNonCanonicalCBOR)
		}
		if err != nil {
			return fmt.Errorf("key %d: %w", k, err)
		}
		seen |= 1 << k
	}
	if err := required(seen, 4); err != nil {
		return err
	}
	return d.EndMap()
}

// String describes c without I_PAYLOAD.
func (c Control) String() string {
	return fmt.Sprintf("Control{header_binding=%x payload_identity=REDACTED}", c.HeaderBinding)
}

// GoString describes c without I_PAYLOAD.
func (c Control) GoString() string { return c.String() }

// EncodeControl returns the Deterministic CBOR bytes of c. The caller must
// wipe the result: it contains I_PAYLOAD.
func EncodeControl(c *Control) ([]byte, error) {
	w := controlWire{
		HeaderBinding:   c.HeaderBinding[:],
		PayloadIdentity: c.PayloadIdentity[:],
	}
	var err error
	if w.Critical, err = extension.Canonical(c.Critical); err != nil {
		return nil, err
	}
	if w.Noncritical, err = extension.Canonical(c.Noncritical); err != nil {
		return nil, err
	}
	if err := extension.CheckDisjoint(w.Critical, w.Noncritical); err != nil {
		return nil, err
	}
	var e codec.Encoder
	w.encode(&e)
	return e.Out()
}

// DecodeControl validates and decodes CONTROL_CBOR (spec §31, §63 step 14).
// A non-canonical encoding is rejected even though CONTROL_CBOR is not hashed.
func DecodeControl(b []byte) (*Control, error) {
	if err := codec.CheckSchema(b, ControlTypeTag, ControlVersion); err != nil {
		return nil, fmt.Errorf("capsule: CONTROL_CBOR: %w", err)
	}
	var w controlWire
	defer func() { clear(w.PayloadIdentity) }()
	if err := codec.Unmarshal(b, w.decode, w.encode); err != nil {
		return nil, fmt.Errorf("capsule: CONTROL_CBOR: %w", err)
	}
	if err := extension.CheckDisjoint(w.Critical, w.Noncritical); err != nil {
		return nil, fmt.Errorf("capsule: CONTROL_CBOR: %w", err)
	}
	c := &Control{Critical: w.Critical, Noncritical: w.Noncritical}
	copy(c.HeaderBinding[:], w.HeaderBinding)
	copy(c.PayloadIdentity[:], w.PayloadIdentity)
	return c, nil
}

// looksLikeAge reports whether b starts with the age v1 intro line.
func looksLikeAge(b []byte) bool {
	return bytes.HasPrefix(b, []byte("age-encryption.org/v1\n"))
}