// 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"
	"errors"
	"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"
	PreludeSize = 16

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

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

	CapsuleIDSize = 16
)

// Format is the format of a capsule: the VERSION byte of its PRELUDE. It
// also fixes the schema version of its CONTROL_CBOR, the number of stanzas of
// its INNER_ACCESS_AGE and whether its payload is padded (spec §22).
type Format uint8

// Formats of V1 (spec §22). A reader opens both.
const (
	// Format1 is the format of spec v0.8.2: one or more stanzas in
	// INNER_ACCESS_AGE and a payload without padding. Only a generator of
	// test vectors may write it (spec §62.1, §70): Encrypt never does.
	Format1 Format = 1
	// Format2 is the format of spec v0.9, the one Encrypt writes: exactly 16
	// stanzas in INNER_ACCESS_AGE and a padded payload (spec §29.1, §39).
	Format2 Format = 2
	// Format3 is the format of spec v0.10: the padded plaintext of
	// PAYLOAD_AGE is BODY, with the security area, the head and several
	// files (spec §29.2 to §29.7).
	Format3 Format = 3
)

func (f Format) valid() bool { return f >= Format1 && f <= Format3 }

// padded reports whether the payload of format f is padded, with L and the
// padding code in keys 6 and 7 of its control (spec §29.1, §31).
func (f Format) padded() bool { return f == Format2 || f == Format3 }

// 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 {
	Format           Format // VERSION
	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] = byte(p.Format)
	// 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, §23, §63 steps 1 and 2), in
// the order of spec §23: magic, a complete prelude, version (the format, 1 or
// 2), FLAGS == 0 and RESERVED == 0, and lengths from 1 up to the 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 !Format(b[4]).valid() {
		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{
		Format:           Format(b[4]),
		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.
// Its schema version is the format of its capsule, which is not part of
// Control: it is given to EncodeControl and DecodeControl.
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
	// PayloadLength is L, the length of the content (key 6), and Padding
	// the padding rule of PAYLOAD_AGE (key 7). Both exist in format 2 only
	// (spec §29.1, §31); in format 1 they are zero.
	PayloadLength uint64
	Padding       Padding
}

// payloadLengthSize is the fixed size of payload_length, so that the length
// of CONTROL_CBOR never depends on L (spec §31, §55.2).
const payloadLengthSize = 8

// controlWire is CONTROL_CBOR as it is encoded: keys 2 to 7, keys 0 and 1
// being the constants ControlTypeTag and the format.
type controlWire struct {
	Format          Format                // key 1, the schema version
	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
	PayloadLength   []byte                // key 6, format 2 only: 8 bytes, big-endian
	Padding         uint64                // key 7, format 2 only
}

func (w *controlWire) encode(e *codec.Encoder) {
	n := 4 + nonEmpty(w.Critical) + nonEmpty(w.Noncritical)
	if w.Format.padded() {
		n += 2
	}
	e.Map(n)
	e.Uint(0)
	e.Text(ControlTypeTag)
	e.Uint(1)
	e.Uint(uint64(w.Format))
	e.Uint(2)
	e.Bstr(w.HeaderBinding)
	e.Uint(3)
	e.Bstr(w.PayloadIdentity)
	encodeExtensions(e, 4, w.Critical, w.Noncritical)
	if w.Format.padded() {
		e.Uint(6)
		e.Bstr(w.PayloadLength)
		e.Uint(7)
		e.Uint(w.Padding)
	}
}

// decode reads CONTROL_CBOR with every CDDL rule of the schema version of
// w.Format: keys 6 and 7 are required in version 2 and not defined in
// version 1. The caller wipes PayloadIdentity, whatever the result.
func (w *controlWire) decode(d *codec.Decoder) error {
	maxPairs := 6
	if w.Format.padded() {
		maxPairs = 8
	}
	pairs, err := d.Map(maxPairs)
	if err != nil {
		return err
	}
	var seen uint
	for range pairs {
		k, err := d.Key()
		if err != nil {
			return err
		}
		if (k == 6 || k == 7) && !w.Format.padded() {
			return fmt.Errorf("key %d is not defined: %w", k, datekeys.ErrNonCanonicalCBOR)
		}
		switch k {
		case 0:
			_, err = d.Text(len(ControlTypeTag))
		case 1:
			_, err = d.Uint(uint64(w.Format))
		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)
		case 6:
			if w.PayloadLength, err = d.Bstr(payloadLengthSize, payloadLengthSize); err == nil {
				if l := binary.BigEndian.Uint64(w.PayloadLength); l > MaxPayloadLength {
					err = fmt.Errorf("payload_length %d exceeds L_MAX = %d: %w", l, uint64(MaxPayloadLength), datekeys.ErrNonCanonicalCBOR)
				}
			}
		case 7:
			// Compared as read, before any narrowing to Padding, which would
			// let 257 and the like pass as a defined code.
			if w.Padding, err = d.Uint(codec.MaxSafeUint); err == nil && w.Padding != uint64(Bloque256) && w.Padding != uint64(Reforzado) {
				err = fmt.Errorf("padding code %d is not defined: %w", w.Padding, datekeys.ErrNonCanonicalCBOR)
			}
		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
	}
	if w.Format.padded() {
		for _, k := range []uint{6, 7} {
			if seen&(1<<k) == 0 {
				return fmt.Errorf("key %d is missing: %w", k, datekeys.ErrNonCanonicalCBOR)
			}
		}
	}
	return d.EndMap()
}

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

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

// EncodeControl returns the Deterministic CBOR bytes of c as CONTROL_CBOR of
// a capsule of format f: schema version 1 without keys 6 and 7, or schema
// version 2 with them (spec §31). The caller must wipe the result: it
// contains I_PAYLOAD.
func EncodeControl(c *Control, f Format) ([]byte, error) {
	w := controlWire{
		Format:          f,
		HeaderBinding:   c.HeaderBinding[:],
		PayloadIdentity: c.PayloadIdentity[:],
	}
	switch f {
	case Format1:
		if c.PayloadLength != 0 || c.Padding != 0 {
			return nil, errors.New("capsule: CONTROL_CBOR of format 1 has no payload_length and no padding")
		}
	case Format2, Format3:
		if !c.Padding.valid() {
			return nil, fmt.Errorf("capsule: padding code %d is not defined", uint8(c.Padding))
		}
		if c.PayloadLength > MaxPayloadLength {
			return nil, fmt.Errorf("capsule: payload_length %d exceeds L_MAX = %d", c.PayloadLength, uint64(MaxPayloadLength))
		}
		w.PayloadLength = binary.BigEndian.AppendUint64(nil, c.PayloadLength)
		w.Padding = uint64(c.Padding)
	default:
		return nil, fmt.Errorf("capsule: format %d is not defined", uint8(f))
	}
	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 the CONTROL_CBOR of a capsule of format
// f (spec §31, §63 step 14). Its schema version must be f, and another is
// ErrUnsupportedVersion whatever follows (spec §69.1, layer 2); in version 2,
// payload_length is 8 bytes of at most MaxPayloadLength and padding is 1 or
// 2. A non-canonical encoding is rejected even though CONTROL_CBOR is not
// hashed.
func DecodeControl(b []byte, f Format) (*Control, error) {
	if !f.valid() {
		return nil, fmt.Errorf("capsule: format %d is not defined", uint8(f))
	}
	if err := codec.CheckSchema(b, ControlTypeTag, uint64(f)); err != nil {
		return nil, fmt.Errorf("capsule: CONTROL_CBOR: %w", err)
	}
	w := controlWire{Format: f}
	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)
	if f.padded() {
		// decode bounds both values, so the conversions are exact.
		c.PayloadLength = binary.BigEndian.Uint64(w.PayloadLength)
		c.Padding = Padding(w.Padding)
	}
	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"))
}