// Package agewrap holds the age recipients and identities that DateKeys wraps
// around standard age files (spec §28-§37), plus the structural stanza rules
// every DateKeys age file must satisfy.
//
// The cryptography is age, tlock and drand's BLS verification. This package
// adds only the rules of the protocol:
//
//   - OUTER_TIME_AGE holds exactly one tlock stanza for the expected round and
//     the pinned chain hash (spec §32, §35, §63 step 11).
//   - PAYLOAD_AGE holds exactly one X25519 stanza, for R_PAYLOAD (spec §29,
//     §63 step 17).
//   - INNER_ACCESS_AGE holds stanzas of type X25519 only, one per recipient:
//     one or more in capsule format 1, exactly AccessSlots in format 2 (spec
//     §33, §39, §63 steps 12 and 13).
//
// The rules are enforced inside Identity.Unwrap, which age calls with the
// complete set of stanzas of the file, so that no file is accepted just
// because age managed to unwrap a file key (spec §27, §63). The same checks
// are exposed for the pre-unlock inspection, which reads the stanzas through a
// probe identity without decrypting anything or touching secrets.
//
// The errors of this package never copy the text of an error of age, tlock,
// kyber or drand: each failure has a fixed message and its normative error.
// That text can carry secrets: when the IBE check of a tlock stanza fails,
// kyber reports the candidate plaintext and r, from which whoever edited the
// stanza learns FK_TIME.
package agewrap

import (
	"bytes"
	"crypto/ecdh"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"strconv"
	"strings"

	"filippo.io/age"
	"github.com/drand/drand/v2/common"
	"github.com/drand/drand/v2/crypto"
	"github.com/drand/kyber"
	"github.com/drand/tlock"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/codec/bech32"
	"g.activething.com/go/DateKeys/profile"
	"g.activething.com/go/DateKeys/provider"
)

// Stanza types of V1.
const (
	StanzaTLock  = "tlock"
	StanzaX25519 = "X25519"
)

// FileKeySize is the size of every age file key: FK_PAYLOAD, FK_ACCESS and
// FK_TIME (spec §28).
const FileKeySize = 16

// ---------------------------------------------------------------------------
// Structural rules

// CheckTimeStanzas enforces the OUTER_TIME_AGE rule: exactly one stanza, of
// type tlock, whose round is the DateKey round and whose chain hash is the one
// of the pinned profile (spec §32, §35, §63 steps 5, 8 and 11).
func CheckTimeStanzas(stanzas []*age.Stanza, p *profile.Profile, round uint64) error {
	if len(stanzas) != 1 {
		return fmt.Errorf("agewrap: OUTER_TIME_AGE has %d stanzas, want exactly one tlock stanza: %w", len(stanzas), datekeys.ErrPolicyStructureMismatch)
	}
	s := stanzas[0]
	if s.Type != StanzaTLock {
		return fmt.Errorf("agewrap: OUTER_TIME_AGE stanza type %q, want %q: %w", s.Type, StanzaTLock, datekeys.ErrPolicyStructureMismatch)
	}
	if len(s.Args) != 2 {
		return fmt.Errorf("agewrap: tlock stanza has %d arguments, want 2: %w", len(s.Args), datekeys.ErrPolicyStructureMismatch)
	}
	if want := strconv.FormatUint(round, 10); s.Args[0] != want {
		return fmt.Errorf("agewrap: tlock stanza round %q, DateKey round %s: %w", s.Args[0], want, datekeys.ErrRoundMismatch)
	}
	if want := p.ChainHashHex(); s.Args[1] != want {
		return fmt.Errorf("agewrap: tlock stanza chain hash %q, pinned profile %s uses %s: %w", s.Args[1], p.ID, want, datekeys.ErrProfileMismatch)
	}
	return nil
}

// CheckPayloadStanzas enforces the PAYLOAD_AGE rule: exactly one stanza, of
// type X25519 (spec §29, §63 steps 6 and 17).
func CheckPayloadStanzas(stanzas []*age.Stanza) error {
	if len(stanzas) != 1 {
		return fmt.Errorf("agewrap: PAYLOAD_AGE has %d stanzas, want exactly one X25519 stanza: %w", len(stanzas), datekeys.ErrPolicyStructureMismatch)
	}
	if t := stanzas[0].Type; t != StanzaX25519 {
		return fmt.Errorf("agewrap: PAYLOAD_AGE stanza type %q, want %q: %w", t, StanzaX25519, datekeys.ErrPolicyStructureMismatch)
	}
	return nil
}

// AccessSlots is the number of stanzas of INNER_ACCESS_AGE in capsule format
// 2: one per credential, and a dummy in each slot left (spec §39).
const AccessSlots = 16

// CheckAccessStanzas enforces the INNER_ACCESS_AGE rule: stanzas of type
// X25519 only, one or more when slots is 0 (capsule format 1) and exactly
// slots otherwise, AccessSlots in format 2 (spec §33, §36, §39, §63 steps 12
// and 13). Two stanzas with the same ephemeral share would be two stanzas for
// one recipient and are rejected.
func CheckAccessStanzas(stanzas []*age.Stanza, slots int) error {
	if len(stanzas) == 0 {
		return fmt.Errorf("agewrap: INNER_ACCESS_AGE has no stanzas: %w", datekeys.ErrPolicyStructureMismatch)
	}
	if slots > 0 && len(stanzas) != slots {
		return fmt.Errorf("agewrap: INNER_ACCESS_AGE has %d stanzas, want exactly %d: %w", len(stanzas), slots, datekeys.ErrPolicyStructureMismatch)
	}
	seen := make(map[string]bool, len(stanzas))
	for i, s := range stanzas {
		if s.Type != StanzaX25519 {
			return fmt.Errorf("agewrap: INNER_ACCESS_AGE stanza %d has type %q, want %q: %w", i, s.Type, StanzaX25519, datekeys.ErrPolicyStructureMismatch)
		}
		if len(s.Args) == 1 {
			if seen[s.Args[0]] {
				return fmt.Errorf("agewrap: INNER_ACCESS_AGE stanza %d repeats an ephemeral share: %w", i, datekeys.ErrPolicyStructureMismatch)
			}
			seen[s.Args[0]] = true
		}
	}
	return nil
}

// ---------------------------------------------------------------------------
// Inspection probe

var errProbe = errors.New("agewrap: probe finished")

// probe records the stanzas age hands to Unwrap and stops decryption with an
// error that does not wrap age.ErrIncorrectIdentity, so age returns it as is.
type probe struct{ stanzas []*age.Stanza }

func (p *probe) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	p.stanzas = cloneStanzas(stanzas)
	return nil, errProbe
}

// Stanzas parses the age header at the start of r with age itself and returns
// its recipient stanzas. It decrypts nothing and uses no secret: the header is
// extracted with age.ExtractHeader and handed to age.DecryptHeader with a
// probe identity (spec §27, §63 steps 5 and 6).
//
// The result is structural. Its authenticity is only established when the
// header MAC is verified while opening the file (spec §27).
func Stanzas(r io.Reader) ([]*age.Stanza, error) {
	hdr, err := age.ExtractHeader(r)
	if err != nil {
		return nil, fmt.Errorf("agewrap: not an age v1 header: malformed, truncated or beyond the parser limits: %w", datekeys.ErrIntegrity)
	}
	var p probe
	if _, err := age.DecryptHeader(hdr, &p); !errors.Is(err, errProbe) {
		return nil, fmt.Errorf("agewrap: age did not hand the stanzas of the header to the probe: %w", datekeys.ErrIntegrity)
	}
	return p.stanzas, nil
}

func cloneStanzas(in []*age.Stanza) []*age.Stanza {
	out := make([]*age.Stanza, len(in))
	for i, s := range in {
		out[i] = &age.Stanza{Type: s.Type, Args: append([]string(nil), s.Args...), Body: bytes.Clone(s.Body)}
	}
	return out
}

// ---------------------------------------------------------------------------
// tlock recipient and identity (OUTER_TIME_AGE)

// TimeRecipient wraps the file key with tlock for one round of a pinned
// profile and emits the stanza "tlock <round> <chainhash>", byte-compatible
// with the stanza of the tlock library and the tle CLI (spec §32, §35). It
// uses only the exported core of tlock: TimeLock and CiphertextToBytes.
type TimeRecipient struct {
	chainHash string
	round     uint64
	scheme    *crypto.Scheme
	key       kyber.Point
}

var _ age.RecipientWithLabels = (*TimeRecipient)(nil)

// NewTimeRecipient returns the tlock recipient of round under p, using only
// the pinned parameters of p (strict mode, spec §35).
func NewTimeRecipient(p *profile.Profile, round uint64) (*TimeRecipient, error) {
	scheme, key, err := pinned(p)
	if err != nil {
		return nil, err
	}
	if round == 0 || round > p.MaxRound() {
		return nil, fmt.Errorf("agewrap: round %d outside the range of %s: %w", round, p.ID, datekeys.ErrDateKeyInvalid)
	}
	return &TimeRecipient{chainHash: p.ChainHashHex(), round: round, scheme: scheme, key: key}, nil
}

// Wrap implements age.Recipient.
func (r *TimeRecipient) Wrap(fileKey []byte) ([]*age.Stanza, error) {
	ct, err := tlock.TimeLock(*r.scheme, r.key, r.round, fileKey)
	if err != nil {
		return nil, errors.New("agewrap: tlock cannot wrap the file key")
	}
	body, err := tlock.CiphertextToBytes(*r.scheme, ct)
	if err != nil {
		return nil, errors.New("agewrap: tlock cannot encode its ciphertext")
	}
	return []*age.Stanza{{
		Type: StanzaTLock,
		Args: []string{strconv.FormatUint(r.round, 10), r.chainHash},
		Body: body,
	}}, nil
}

// WrapWithLabels implements age.RecipientWithLabels with a random label, so
// that age refuses to mix this recipient with any other one in the same file:
// OUTER_TIME_AGE must hold exactly one tlock stanza.
func (r *TimeRecipient) WrapWithLabels(fileKey []byte) ([]*age.Stanza, []string, error) {
	s, err := r.Wrap(fileKey)
	if err != nil {
		return nil, nil, err
	}
	var label [16]byte
	_, _ = rand.Read(label[:]) // never fails since Go 1.24
	return s, []string{"datekeys-tlock-" + hex.EncodeToString(label[:])}, nil
}

// TimeIdentity opens OUTER_TIME_AGE under the strict rules of spec §35 and
// §63 step 11. Unwrap validates the complete stanza set, verifies the release
// locally, checks the form of the stanza body and calls tlock.TimeUnlock,
// which verifies the beacon again before decrypting. Every failure keeps its
// own normative error: none is turned into "too early".
type TimeIdentity struct {
	profile *profile.Profile
	round   uint64
	release provider.Release
	scheme  *crypto.Scheme
	key     kyber.Point
}

var _ age.Identity = (*TimeIdentity)(nil)

// NewTimeIdentity returns the identity that opens OUTER_TIME_AGE for round
// with release.
func NewTimeIdentity(p *profile.Profile, round uint64, release provider.Release) (*TimeIdentity, error) {
	scheme, key, err := pinned(p)
	if err != nil {
		return nil, err
	}
	return &TimeIdentity{profile: p.Clone(), round: round, release: release, scheme: scheme, key: key}, nil
}

// tlockBlockLen is the size of V and of W in a tlock stanza body, fixed by
// tlock whatever the scheme (spec §63 step 11).
const tlockBlockLen = 16

// Unwrap implements age.Identity. The stanza body is U || V || W (spec §63
// step 11): |U| is the point size of the key group of the scheme, 96 bytes for
// Quicknet, and |V| = |W| = 16. U must be the canonical encoding of a point
// of that group other than the point at infinity (spec §12.2): the decoder of
// drand, which tlock.BytesToCiphertext runs, rejects every other encoding,
// and the point at infinity is rejected here. tlock.TimeUnlock then decrypts
// with the verified release and checks r·G == U. Every failure of the body is
// ErrIntegrity.
func (i *TimeIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	if err := CheckTimeStanzas(stanzas, i.profile, i.round); err != nil {
		return nil, err
	}
	if err := provider.Verify(i.profile, provider.Condition{Round: i.round}, i.release); err != nil {
		return nil, err
	}
	body := stanzas[0].Body
	if want := i.scheme.KeyGroup.PointLen() + 2*tlockBlockLen; len(body) != want {
		return nil, fmt.Errorf("agewrap: tlock stanza body of %d bytes, want %d: %w", len(body), want, datekeys.ErrIntegrity)
	}
	ct, err := tlock.BytesToCiphertext(*i.scheme, body)
	if err != nil {
		// With the length right, only the decoding of U fails.
		return nil, fmt.Errorf("agewrap: U of the tlock stanza is not the canonical encoding of a point of the key group: %w", datekeys.ErrIntegrity)
	}
	if ct.U.Equal(ct.U.Null()) {
		return nil, fmt.Errorf("agewrap: U of the tlock stanza is the point at infinity: %w", datekeys.ErrIntegrity)
	}
	beacon := common.Beacon{Round: i.release.Round, Signature: i.release.Signature}
	fileKey, err := tlock.TimeUnlock(*i.scheme, i.key, beacon, ct)
	if err != nil {
		// Not the error of tlock: for a failed IBE check it carries the
		// candidate plaintext and r (see the package documentation).
		return nil, fmt.Errorf("agewrap: the tlock stanza body does not decrypt under the verified release (IBE check r·G == U): %w", datekeys.ErrIntegrity)
	}
	if len(fileKey) != FileKeySize {
		return nil, fmt.Errorf("agewrap: tlock stanza wraps a %d-byte file key: %w", len(fileKey), datekeys.ErrIntegrity)
	}
	return fileKey, nil
}

func pinned(p *profile.Profile) (*crypto.Scheme, kyber.Point, error) {
	scheme, err := p.DrandScheme()
	if err != nil {
		return nil, nil, err
	}
	key := scheme.KeyGroup.Point()
	if err := key.UnmarshalBinary(p.PublicKey); err != nil {
		return nil, nil, fmt.Errorf("agewrap: pinned public key of %s is not the canonical encoding of a point of the key group: %w", p.ID, datekeys.ErrUnknownProfile)
	}
	if key.Equal(key.Null()) {
		return nil, nil, fmt.Errorf("agewrap: pinned public key of %s is the identity element: %w", p.ID, datekeys.ErrUnknownProfile)
	}
	return scheme, key, nil
}

// ---------------------------------------------------------------------------
// X25519 identities (PAYLOAD_AGE and INNER_ACCESS_AGE)

// x25519StanzaForm is the form of an X25519 stanza that the age
// specification requires (spec §63 step 13): age rejects any other before a
// key agreement, and its error is not copied.
const x25519StanzaForm = "one argument, a 32-byte ephemeral share not of low order, and a 32-byte body"

// PayloadIdentity opens PAYLOAD_AGE with I_PAYLOAD (spec §29, §30.1, §63 step
// 17). It rejects the file unless it holds exactly one X25519 stanza and that
// stanza is for R_PAYLOAD.
type PayloadIdentity struct {
	id *age.X25519Identity
}

var _ age.Identity = (*PayloadIdentity)(nil)

// NewPayloadIdentity returns the identity for the raw 32-byte I_PAYLOAD.
func NewPayloadIdentity(raw []byte) (*PayloadIdentity, error) {
	id, err := X25519IdentityFromRaw(raw)
	if err != nil {
		return nil, err
	}
	return &PayloadIdentity{id: id}, nil
}

// Unwrap implements age.Identity.
func (i *PayloadIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	if err := CheckPayloadStanzas(stanzas); err != nil {
		return nil, err
	}
	fileKey, err := i.id.Unwrap(stanzas)
	if errors.Is(err, age.ErrIncorrectIdentity) {
		// CONTROL_A + PAYLOAD_AGE_B: I_PAYLOAD_A cannot unwrap FK_PAYLOAD_B (spec §30.1).
		return nil, fmt.Errorf("agewrap: PAYLOAD_AGE is not encrypted to this control's R_PAYLOAD: %w", datekeys.ErrIntegrity)
	}
	if err != nil {
		return nil, fmt.Errorf("agewrap: malformed X25519 stanza in PAYLOAD_AGE (%s): %w", x25519StanzaForm, datekeys.ErrIntegrity)
	}
	return fileKey, nil
}

// AccessIdentity opens INNER_ACCESS_AGE with the caller's X25519 identities,
// including the one of a portable .dkk (spec §33, §38, §63 step 13). It
// validates the complete stanza set first, with the rule of the capsule
// format, and rejects the file if any identity unwraps more than one stanza,
// which would be two stanzas for the same recipient.
type AccessIdentity struct {
	slots int
	ids   []age.Identity
}

var _ age.Identity = (*AccessIdentity)(nil)

// NewAccessIdentity returns an AccessIdentity trying every identity of ids
// on a file that must hold one or more stanzas when slots is 0 (capsule
// format 1) and exactly slots otherwise, AccessSlots in format 2. Nil entries
// are dropped; at least one identity is required.
func NewAccessIdentity(slots int, ids ...age.Identity) (*AccessIdentity, error) {
	if slots < 0 {
		return nil, fmt.Errorf("agewrap: %d slots", slots)
	}
	var clean []age.Identity
	for _, id := range ids {
		if id != nil {
			clean = append(clean, id)
		}
	}
	if len(clean) == 0 {
		return nil, fmt.Errorf("agewrap: time_and_key needs an access identity: %w", datekeys.ErrAccessRequired)
	}
	return &AccessIdentity{slots: slots, ids: clean}, nil
}

// Unwrap implements age.Identity. The codes follow spec §63 step 13, in this
// order: the stanza rules (ErrPolicyStructureMismatch); a malformed X25519
// stanza, found by the first identity already, because age checks the form
// of a stanza before any key agreement (ErrIntegrity); an identity that
// unwraps more than one stanza, even if another one unwraps exactly one
// (ErrPolicyStructureMismatch); no identity that unwraps any
// (ErrAccessInvalid). Every identity is tried against every stanza, so that
// the result does not depend on the order of the identities (spec §69.1).
func (a *AccessIdentity) Unwrap(stanzas []*age.Stanza) ([]byte, error) {
	if err := CheckAccessStanzas(stanzas, a.slots); err != nil {
		return nil, err
	}
	var fileKey []byte
	for _, id := range a.ids {
		matches := 0
		for _, s := range stanzas {
			fk, err := id.Unwrap([]*age.Stanza{s})
			if errors.Is(err, age.ErrIncorrectIdentity) {
				continue
			}
			if err != nil {
				return nil, fmt.Errorf("agewrap: malformed X25519 stanza in INNER_ACCESS_AGE (%s): %w", x25519StanzaForm, datekeys.ErrIntegrity)
			}
			matches++
			if fileKey == nil {
				fileKey = fk
			}
		}
		if matches > 1 {
			return nil, fmt.Errorf("agewrap: one identity opens %d INNER_ACCESS_AGE stanzas, want one per recipient: %w", matches, datekeys.ErrPolicyStructureMismatch)
		}
	}
	if fileKey == nil {
		return nil, fmt.Errorf("agewrap: no supplied identity is a recipient of INNER_ACCESS_AGE: %w", datekeys.ErrAccessInvalid)
	}
	return fileKey, nil
}

// ---------------------------------------------------------------------------
// Raw X25519 keys

// X25519IdentityFromRaw converts 32 raw identity bytes, the canonical form
// inside CONTROL_CBOR and .dkk (spec §31, §38), to an age identity.
func X25519IdentityFromRaw(raw []byte) (*age.X25519Identity, error) {
	if len(raw) != 32 {
		return nil, fmt.Errorf("agewrap: X25519 identity is %d bytes, want 32: %w", len(raw), datekeys.ErrIntegrity)
	}
	s, err := bech32.Encode("AGE-SECRET-KEY-", raw)
	if err != nil {
		return nil, fmt.Errorf("agewrap: cannot encode the X25519 identity: %w", datekeys.ErrIntegrity)
	}
	id, err := age.ParseX25519Identity(strings.ToUpper(s))
	if err != nil {
		return nil, fmt.Errorf("agewrap: age rejects the encoded X25519 identity: %w", datekeys.ErrIntegrity)
	}
	return id, nil
}

// RawX25519Identity returns the 32 raw bytes of an age X25519 identity.
func RawX25519Identity(id *age.X25519Identity) ([]byte, error) {
	hrp, raw, err := bech32.Decode(id.String())
	if err != nil || hrp != "AGE-SECRET-KEY-" || len(raw) != 32 {
		return nil, fmt.Errorf("agewrap: unexpected age identity encoding")
	}
	return raw, nil
}

// RawX25519Recipient returns the 32 raw bytes of an age X25519 recipient.
func RawX25519Recipient(r *age.X25519Recipient) ([]byte, error) {
	hrp, raw, err := bech32.Decode(r.String())
	if err != nil || hrp != "age" || len(raw) != 32 {
		return nil, fmt.Errorf("agewrap: unexpected age recipient encoding")
	}
	return raw, nil
}

// lowOrderProbe is any X25519 scalar: with the clamping of RFC 7748, the
// result for a point of low order is the all-zero string whatever the scalar.
var lowOrderProbe = [32]byte{1}

// CheckX25519Recipient rejects the X25519 recipients a writer MUST NOT
// encrypt to (spec §37, §62.1 rule 3): a non-canonical one, with bit 255 set
// or with u >= p = 2^255 - 19, whose stanza no identity opens because age
// salts HKDF with the 32 bytes as given; and one of low order, for which the
// shared secret is zero, so that anyone could open its stanza. The rules of
// the reader (spec §63) detect neither case: the stanza does not contain the
// recipient.
func CheckX25519Recipient(r *age.X25519Recipient) error {
	raw, err := RawX25519Recipient(r)
	if err != nil {
		return err
	}
	if raw[31]&0x80 != 0 {
		return fmt.Errorf("agewrap: recipient %s is not canonical: bit 255 is set", r)
	}
	if raw[31] == 0x7f && raw[0] >= 0xed && bytes.Count(raw[1:31], []byte{0xff}) == 30 {
		return fmt.Errorf("agewrap: recipient %s is not canonical: u is not below 2^255 - 19", r)
	}
	pub, err := ecdh.X25519().NewPublicKey(raw)
	if err != nil {
		return fmt.Errorf("agewrap: recipient %s is not an X25519 public key", r)
	}
	probe, err := ecdh.X25519().NewPrivateKey(lowOrderProbe[:])
	if err != nil {
		return err
	}
	if _, err := probe.ECDH(pub); err != nil {
		return fmt.Errorf("agewrap: recipient %s is a point of low order: the shared secret would be zero", r)
	}
	return nil
}