// Package accesskey implements the DateKeys Access Key, the portable .dkk
// credential (spec §38, §40-§44).
//
// A .dkk is a sensitive capability (spec §7.4). Its X25519 identity is stored
// as 32 raw bytes; the Bech32 AGE-SECRET-KEY-1... form is only an export
// format for humans (spec §38). No type in this package prints the material.
package accesskey

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"io"

	"filippo.io/age"

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

// Framing and schema constants (spec §40, §41).
const (
	Magic          = "DKK1"
	FramingVersion = 1
	PreludeSize    = 12
	// MaxBodyLen is the parser limit of spec §57, checked before allocating.
	MaxBodyLen    = 16 << 20
	TypeTag       = "datekeys-access-key"
	SchemaVersion = 1
	// TypeX25519 is the only access_type of V1 (spec §41).
	TypeX25519 = "x25519"

	idSize     = 16
	digestSize = 32
	x25519Size = 32
)

// AccessKey is a decoded .dkk.
type AccessKey struct {
	CredentialID [16]byte // key 2, random and opaque (spec §42)
	CapsuleID    [16]byte // key 3, the only capsule this credential is for (spec §38)
	Type         string   // key 4, access_type
	Material     []byte   // key 5, access_material: 32 raw X25519 identity bytes. SECRET.
	// Verification is key 6, optional; nil when absent (spec §43, §58.1).
	Verification *Verification
	Critical     []extension.Extension // key 7
	Noncritical  []extension.Extension // key 8
}

// Verification is verification_metadata (spec §43). It supports fast failure
// and UX only; it is not a security property.
type Verification struct {
	CapsuleDigest []byte // key 0, SHA-256 of the exact .dkc bytes
}

// bodyWire is BODY_CBOR as it is encoded: keys 2 to 8, keys 0 and 1 being
// the constants TypeTag and SchemaVersion.
type bodyWire struct {
	CredentialID []byte // key 2
	CapsuleID    []byte // key 3
	AccessType   string // key 4
	Material     []byte // key 5, SECRET
	// Digest is capsule_digest, the only key of verification_metadata
	// (key 6); nil when key 6 is omitted.
	Digest      []byte
	Critical    []extension.Extension // key 7, omitted when empty
	Noncritical []extension.Extension // key 8, omitted when empty
}

func (w *bodyWire) encode(e *codec.Encoder) {
	pairs := 6
	for _, present := range []bool{w.Digest != nil, len(w.Critical) > 0, len(w.Noncritical) > 0} {
		if present {
			pairs++
		}
	}
	e.Map(pairs)
	e.Uint(0)
	e.Text(TypeTag)
	e.Uint(1)
	e.Uint(SchemaVersion)
	e.Uint(2)
	e.Bstr(w.CredentialID)
	e.Uint(3)
	e.Bstr(w.CapsuleID)
	e.Uint(4)
	e.Text(w.AccessType)
	e.Uint(5)
	e.Bstr(w.Material)
	if w.Digest != nil {
		e.Uint(6)
		e.Map(1)
		e.Uint(0)
		e.Bstr(w.Digest)
	}
	if len(w.Critical) > 0 {
		e.Uint(7)
		extension.EncodeArray(e, w.Critical)
	}
	if len(w.Noncritical) > 0 {
		e.Uint(8)
		extension.EncodeArray(e, w.Noncritical)
	}
}

// decode reads BODY_CBOR with every CDDL rule whose violation is
// ErrNonCanonicalCBOR; access_type and access_material, which have a code of
// their own (spec §57), are checked afterwards. The caller wipes Material,
// whatever the result.
func (w *bodyWire) decode(d *codec.Decoder) error {
	pairs, err := d.Map(9)
	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(TypeTag))
		case 1:
			_, err = d.Uint(SchemaVersion)
		case 2:
			w.CredentialID, err = d.Bstr(idSize, idSize)
		case 3:
			w.CapsuleID, err = d.Bstr(idSize, idSize)
		case 4:
			w.AccessType, err = d.Text(MaxBodyLen)
		case 5:
			w.Material, err = d.Bstr(0, MaxBodyLen)
		case 6:
			w.Digest, err = decodeVerification(d)
		case 7:
			w.Critical, err = extension.DecodeArray(d)
		case 8:
			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
	}
	for k := range 6 {
		if seen&(1<<k) == 0 {
			return fmt.Errorf("key %d is missing: %w", k, datekeys.ErrNonCanonicalCBOR)
		}
	}
	return d.EndMap()
}

// decodeVerification reads verification_metadata, {0: capsule_digest}. It is
// present only when it holds a digest: an empty map is not a representation
// of absence (spec §43, §58.1).
func decodeVerification(d *codec.Decoder) ([]byte, error) {
	pairs, err := d.Map(1)
	if err != nil {
		return nil, err
	}
	if pairs == 0 {
		return nil, fmt.Errorf("empty verification_metadata; an absent one omits key 6: %w", datekeys.ErrNonCanonicalCBOR)
	}
	k, err := d.Key()
	if err != nil {
		return nil, err
	}
	if k != 0 {
		return nil, fmt.Errorf("verification_metadata key %d is not defined: %w", k, datekeys.ErrNonCanonicalCBOR)
	}
	digest, err := d.Bstr(digestSize, digestSize)
	if err != nil {
		return nil, fmt.Errorf("capsule_digest: %w", err)
	}
	return digest, d.EndMap()
}

// String describes k without its material.
func (k AccessKey) String() string {
	return fmt.Sprintf("AccessKey{credential_id=%x capsule_id=%x type=%s material=REDACTED}", k.CredentialID, k.CapsuleID, k.Type)
}

// GoString describes k without its material.
func (k AccessKey) GoString() string { return k.String() }

// Identity returns the age identity of an x25519 access key.
func (k *AccessKey) Identity() (age.Identity, error) {
	if err := k.validateMaterial(); err != nil {
		return nil, err
	}
	id, err := agewrap.X25519IdentityFromRaw(k.Material)
	if err != nil {
		return nil, fmt.Errorf("accesskey: %v: %w", err, datekeys.ErrAccessInvalid)
	}
	return id, nil
}

// Wipe overwrites the material in place. It is best effort: Go may have made
// copies that cannot be reached.
func (k *AccessKey) Wipe() { clear(k.Material) }

func (k *AccessKey) validateMaterial() error {
	if k.Type != TypeX25519 {
		return fmt.Errorf("accesskey: access_type %q is not supported by V1: %w", k.Type, datekeys.ErrAccessInvalid)
	}
	if len(k.Material) != x25519Size {
		return fmt.Errorf("accesskey: x25519 access_material is %d bytes, want %d: %w", len(k.Material), x25519Size, datekeys.ErrAccessInvalid)
	}
	return nil
}

// MarshalBody returns BODY_CBOR, the Deterministic CBOR body of k (spec §41),
// after checking that DecodeBody accepts it. The extensions of k are written
// as given, once they pass the rules of spec §54: MarshalBody, and Encode,
// take no extension.Registry, so the application writes a registered
// extension only in the arrays of a .dkk it is registered for (spec §72).
func (k *AccessKey) MarshalBody() ([]byte, error) {
	if err := k.validateMaterial(); err != nil {
		return nil, err
	}
	w := bodyWire{
		CredentialID: k.CredentialID[:],
		CapsuleID:    k.CapsuleID[:],
		AccessType:   k.Type,
		Material:     k.Material,
	}
	if k.Verification != nil {
		if len(k.Verification.CapsuleDigest) != digestSize {
			// An empty map is not a canonical representation of absence (spec §43).
			return nil, fmt.Errorf("accesskey: capsule_digest must be %d bytes: %w", digestSize, datekeys.ErrNonCanonicalCBOR)
		}
		w.Digest = k.Verification.CapsuleDigest
	}
	var err error
	if w.Critical, err = extension.Canonical(k.Critical); err != nil {
		return nil, err
	}
	if w.Noncritical, err = extension.Canonical(k.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) > MaxBodyLen {
		clear(b)
		return nil, fmt.Errorf("accesskey: BODY_CBOR of %d bytes exceeds %d: %w", len(b), MaxBodyLen, datekeys.ErrIntegrity)
	}
	// Self-check (spec §72): the reader must accept what is written.
	back, err := DecodeBody(b)
	if err != nil {
		clear(b)
		return nil, fmt.Errorf("accesskey: self-check: the reader rejects this body: %w", err)
	}
	back.Wipe()
	return b, nil
}

// Encode writes k as a complete .dkk: prelude and BODY_CBOR (spec §40). The
// body it encodes, which holds the material, is wiped once written.
func Encode(w io.Writer, k *AccessKey) error {
	body, err := k.MarshalBody()
	if err != nil {
		return err
	}
	defer clear(body)
	var pre [PreludeSize]byte
	copy(pre[0:4], Magic)
	pre[4] = FramingVersion
	binary.BigEndian.PutUint32(pre[8:12], uint32(len(body)))
	if _, err := w.Write(pre[:]); err != nil {
		return err
	}
	_, err = w.Write(body)
	return err
}

// Decode reads exactly one .dkk from r and validates its framing, its
// canonical body and its fields. Bytes after BODY_CBOR are rejected.
//
// Decode does not decide whether critical extensions are known; the consumer
// checks them against its extension.Registry (capsule.Open does).
func Decode(r io.Reader) (*AccessKey, error) {
	var pre [PreludeSize]byte
	n, err := io.ReadFull(r, pre[:])
	if n < 4 || string(pre[0:4]) != Magic {
		return nil, fmt.Errorf("accesskey: %w", datekeys.ErrInvalidMagic)
	}
	if err != nil {
		return nil, fmt.Errorf("accesskey: truncated prelude: %w", datekeys.ErrIntegrity)
	}
	if pre[4] != FramingVersion {
		return nil, fmt.Errorf("accesskey: framing version %d: %w", pre[4], datekeys.ErrUnsupportedVersion)
	}
	if pre[5] != 0 || pre[6] != 0 || pre[7] != 0 {
		return nil, fmt.Errorf("accesskey: flags %#x, reserved %#x%02x: %w", pre[5], pre[6], pre[7], datekeys.ErrInvalidFlags)
	}
	// Spec §40, §57: BODY_LEN in 1..16 MiB. No empty frame holds a valid
	// body, so 0 is a framing error, like a PUBLIC_HEADER_LEN of 0 (§22).
	bodyLen := binary.BigEndian.Uint32(pre[8:12])
	if bodyLen == 0 || bodyLen > MaxBodyLen {
		return nil, fmt.Errorf("accesskey: BODY_LEN %d outside 1..%d: %w", bodyLen, MaxBodyLen, datekeys.ErrIntegrity)
	}
	body, err := readBody(r, int(bodyLen))
	if err != nil {
		return nil, fmt.Errorf("accesskey: truncated body: %w", datekeys.ErrIntegrity)
	}
	defer clear(body)
	var extra [1]byte
	switch n, err := io.ReadFull(r, extra[:]); {
	case n != 0:
		return nil, fmt.Errorf("accesskey: data after BODY_CBOR: %w", datekeys.ErrIntegrity)
	case !errors.Is(err, io.EOF):
		return nil, fmt.Errorf("accesskey: reading after BODY_CBOR: %w", err)
	}
	return DecodeBody(body)
}

// readBody reads the n bytes of BODY_CBOR from r. The buffer grows with the
// data actually read, so a short file that declares a large BODY_LEN does not
// force an allocation of that size. The body holds access_material: every
// buffer it outgrows is wiped, and so is the partial body on error.
func readBody(r io.Reader, n int) ([]byte, error) {
	buf := make([]byte, 0, min(n, bytes.MinRead))
	for len(buf) < n {
		if len(buf) == cap(buf) {
			grown := make([]byte, len(buf), min(n, 2*cap(buf)))
			copy(grown, buf)
			clear(buf)
			buf = grown
		}
		m, err := r.Read(buf[len(buf):cap(buf)])
		buf = buf[:len(buf)+m]
		if err != nil && len(buf) < n {
			// Read may use the whole of its argument as scratch space.
			clear(buf[:cap(buf)])
			return nil, err
		}
	}
	return buf, nil
}

// DecodeBody validates and decodes BODY_CBOR, which the DKK BODY limit of
// spec §57 bounds whatever the caller read it from.
func DecodeBody(body []byte) (*AccessKey, error) {
	if len(body) > MaxBodyLen {
		return nil, fmt.Errorf("accesskey: BODY_CBOR of %d bytes exceeds %d: %w", len(body), MaxBodyLen, datekeys.ErrIntegrity)
	}
	if err := codec.CheckSchema(body, TypeTag, SchemaVersion); err != nil {
		return nil, fmt.Errorf("accesskey: %w", err)
	}
	var w bodyWire
	defer func() { clear(w.Material) }()
	if err := codec.Unmarshal(body, w.decode, w.encode); err != nil {
		return nil, fmt.Errorf("accesskey: %w", err)
	}
	if err := extension.CheckDisjoint(w.Critical, w.Noncritical); err != nil {
		return nil, fmt.Errorf("accesskey: %w", err)
	}
	k := &AccessKey{Type: w.AccessType, Critical: w.Critical, Noncritical: w.Noncritical}
	copy(k.CredentialID[:], w.CredentialID)
	copy(k.CapsuleID[:], w.CapsuleID)
	if w.Digest != nil {
		k.Verification = &Verification{CapsuleDigest: w.Digest}
	}
	k.Material = bytes.Clone(w.Material)
	if err := k.validateMaterial(); err != nil {
		k.Wipe()
		return nil, err
	}
	return k, nil
}