package capsule

import (
	"bytes"
	"crypto/rand"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"time"

	"filippo.io/age"

	datekeys "github.com/datekeys/datekeys-go"
	"github.com/datekeys/datekeys-go/accesskey"
	"github.com/datekeys/datekeys-go/agewrap"
	"github.com/datekeys/datekeys-go/datekey"
	"github.com/datekeys/datekeys-go/extension"
	"github.com/datekeys/datekeys-go/profile"
)

// EncryptOptions configures Encrypt.
type EncryptOptions struct {
	// Profile is the pinned Provider Profile. Required.
	Profile *profile.Profile
	// UnlockAt is the requested instant. It resolves locally to the first
	// round at or after it (spec §15) and must be after Now.
	UnlockAt time.Time
	// Policy is time_only or time_and_key (spec §25).
	Policy Policy
	// Recipients are the X25519 recipients of known holders, for
	// time_and_key (spec §33, §37, §39). Only *age.X25519Recipient is
	// accepted: INNER_ACCESS_AGE must hold X25519 stanzas only.
	Recipients []age.Recipient
	// NewPortableKey generates a fresh I_ACCESS for this capsule only and
	// returns it as a .dkk (spec §38). An I_ACCESS is never reused: Encrypt
	// accepts no existing one.
	NewPortableKey bool
	// Critical and Noncritical are the PUBLIC_HEADER extensions (visible to
	// anyone holding the .dkc).
	Critical, Noncritical []extension.Extension
	// ControlCritical and ControlNoncritical are the CONTROL_CBOR extensions,
	// sealed with the control.
	ControlCritical, ControlNoncritical []extension.Extension
	// Now is the clock. Required: no package of this module reads the wall
	// clock on its own.
	Now func() time.Time
}

// Result describes a capsule written by Encrypt.
type Result struct {
	DateKey   datekey.DateKey
	UnlockAt  time.Time // effective round time, never before the requested instant
	CapsuleID [CapsuleIDSize]byte
	// PortableKey is the .dkk generated when NewPortableKey is set. Encode it
	// with accesskey.Encode and treat it as a sensitive capability.
	PortableKey *accesskey.AccessKey
}

// Encrypt writes a .dkc for the payload read from src (spec §61 for
// time_only, §62 for time_and_key). It needs no network: the round is
// resolved locally and tlock uses only the pinned public key.
//
// PAYLOAD_AGE is streamed after the small, in-memory SEALED_CONTROL, so the
// payload is never held in memory. On error dst may hold a partial capsule
// that must be discarded.
func Encrypt(dst io.Writer, src io.Reader, opts EncryptOptions) (*Result, error) {
	p := opts.Profile
	if p == nil {
		return nil, errors.New("capsule: EncryptOptions.Profile is required")
	}
	if opts.Now == nil {
		return nil, errors.New("capsule: EncryptOptions.Now is required")
	}
	if err := p.Validate(); err != nil {
		return nil, err
	}
	if !opts.UnlockAt.After(opts.Now()) {
		return nil, fmt.Errorf("capsule: unlock time %s is not in the future", opts.UnlockAt.UTC().Format(time.RFC3339Nano))
	}

	// Step 1: resolve the DateKey locally.
	dk, err := datekey.Resolve(p, opts.UnlockAt)
	if err != nil {
		return nil, err
	}
	unlock := dk.UnlockAt(p)
	// Spec §17: round_time(round) >= requested_unlock_at, never earlier.
	if unlock.Before(opts.UnlockAt) {
		return nil, fmt.Errorf("capsule: resolved round %d opens before the requested time: %w", dk.Round, datekeys.ErrRoundMismatch)
	}

	access, portable, err := accessRecipients(opts)
	if err != nil {
		return nil, err
	}
	var portableRaw []byte
	if portable != nil {
		if portableRaw, err = agewrap.RawX25519Identity(portable); err != nil {
			return nil, err
		}
		defer clear(portableRaw)
	}

	// Step 2: capsule_id, 16 random bytes (spec §21).
	var capsuleID [CapsuleIDSize]byte
	_, _ = rand.Read(capsuleID[:]) // never fails since Go 1.24

	// Step 3: I_PAYLOAD, a fresh X25519 identity (spec §29).
	payloadID, err := age.GenerateX25519Identity()
	if err != nil {
		return nil, err
	}
	payloadRaw, err := agewrap.RawX25519Identity(payloadID)
	if err != nil {
		return nil, err
	}
	defer clear(payloadRaw)

	// Step 5: PUBLIC_HEADER.
	header := &Header{CapsuleID: capsuleID, DateKey: dk, Policy: opts.Policy, Critical: opts.Critical, Noncritical: opts.Noncritical}
	headerBytes, err := EncodeHeader(header)
	if err != nil {
		return nil, err
	}
	if len(headerBytes) > MaxPublicHeaderLen {
		return nil, fmt.Errorf("capsule: PUBLIC_HEADER of %d bytes exceeds %d", len(headerBytes), MaxPublicHeaderLen)
	}

	timeRecipient, err := agewrap.NewTimeRecipient(p, dk.Round)
	if err != nil {
		return nil, err
	}
	seal := func(control []byte) ([]byte, error) {
		plaintext := control
		if opts.Policy == TimeAndKey {
			// INNER_ACCESS_AGE: FK_ACCESS wrapped for every access recipient.
			innerAge, err := encryptAll(control, access...)
			if err != nil {
				return nil, err
			}
			stanzas, err := agewrap.Stanzas(bytes.NewReader(innerAge))
			if err != nil {
				return nil, err
			}
			if err := agewrap.CheckAccessStanzas(stanzas); err != nil || len(stanzas) != len(access) {
				return nil, fmt.Errorf("capsule: INNER_ACCESS_AGE self-check failed: %w", datekeys.ErrPolicyStructureMismatch)
			}
			plaintext = innerAge
		}
		// OUTER_TIME_AGE: FK_TIME wrapped with tlock for the DateKey round.
		return encryptAll(plaintext, timeRecipient)
	}

	// Steps 6 to 10. PRELUDE carries SEALED_CONTROL_LEN and header_binding
	// covers PRELUDE, so the length is measured first by sealing a control of
	// identical size with a zero binding and a zero identity. age output
	// lengths depend only on plaintext length and stanza shapes; the real
	// seal is checked to have the same length.
	ctrl := &Control{Critical: opts.ControlCritical, Noncritical: opts.ControlNoncritical}
	draft, err := EncodeControl(ctrl)
	if err != nil {
		return nil, err
	}
	draftSealed, err := seal(draft)
	if err != nil {
		return nil, err
	}
	if len(draftSealed) > MaxSealedControlLen {
		return nil, fmt.Errorf("capsule: SEALED_CONTROL of %d bytes exceeds %d", len(draftSealed), MaxSealedControlLen)
	}
	prelude := Prelude{PublicHeaderLen: uint32(len(headerBytes)), SealedControlLen: uint32(len(draftSealed))}
	preludeBytes := prelude.Bytes()

	// Step 7: header_binding = SHA-256(PRELUDE || PUBLIC_HEADER_BYTES).
	ctrl.HeaderBinding = HeaderBinding(preludeBytes, headerBytes)
	copy(ctrl.PayloadIdentity[:], payloadRaw)
	defer clear(ctrl.PayloadIdentity[:])

	// Step 8: CONTROL_CBOR.
	controlBytes, err := EncodeControl(ctrl)
	if err != nil {
		return nil, err
	}
	defer clear(controlBytes)

	// Steps 9 and 10: SEALED_CONTROL = OUTER_TIME_AGE.
	sealed, err := seal(controlBytes)
	if err != nil {
		return nil, err
	}
	if len(sealed) != len(draftSealed) {
		return nil, fmt.Errorf("capsule: internal error: SEALED_CONTROL is %d bytes, measured %d", len(sealed), len(draftSealed))
	}

	// Step 11: PRELUDE || PUBLIC_HEADER || SEALED_CONTROL || PAYLOAD_AGE.
	digest := sha256.New()
	w := io.MultiWriter(dst, digest)
	for _, b := range [][]byte{preludeBytes[:], headerBytes, sealed} {
		if _, err := w.Write(b); err != nil {
			return nil, err
		}
	}
	// Step 4: PAYLOAD_AGE, a standard age file for R_PAYLOAD (FK_PAYLOAD is
	// generated by age). It is written last and streamed.
	aw, err := age.Encrypt(w, payloadID.Recipient())
	if err != nil {
		return nil, err
	}
	if _, err := io.Copy(aw, src); err != nil {
		return nil, err
	}
	if err := aw.Close(); err != nil {
		return nil, err
	}

	res := &Result{DateKey: dk, UnlockAt: unlock, CapsuleID: capsuleID}
	if portable != nil {
		// Step 13: the portable identity as 32 raw bytes in a .dkk.
		k := &accesskey.AccessKey{
			CapsuleID:    capsuleID,
			Type:         accesskey.TypeX25519,
			Material:     bytes.Clone(portableRaw),
			Verification: &accesskey.Verification{CapsuleDigest: digest.Sum(nil)},
		}
		_, _ = rand.Read(k.CredentialID[:]) // spec §42; never fails since Go 1.24
		res.PortableKey = k
	}
	return res, nil
}

// accessRecipients validates the policy options and returns the recipients of
// INNER_ACCESS_AGE, including R_ACCESS when a portable key is requested.
func accessRecipients(opts EncryptOptions) ([]age.Recipient, *age.X25519Identity, error) {
	switch opts.Policy {
	case TimeOnly:
		if len(opts.Recipients) != 0 || opts.NewPortableKey {
			return nil, nil, errors.New("capsule: time_only takes no recipients and no portable key")
		}
		return nil, nil, nil
	case TimeAndKey:
	default:
		return nil, nil, fmt.Errorf("capsule: unknown access policy %d", opts.Policy)
	}
	var out []age.Recipient
	seen := make(map[string]bool)
	for i, r := range opts.Recipients {
		x, ok := r.(*age.X25519Recipient)
		if !ok || x == nil {
			return nil, nil, fmt.Errorf("capsule: recipient %d is %T; time_and_key accepts X25519 recipients only", i, r)
		}
		if seen[x.String()] {
			return nil, nil, fmt.Errorf("capsule: recipient %s listed twice; INNER_ACCESS_AGE holds one stanza per recipient", x)
		}
		seen[x.String()] = true
		out = append(out, x)
	}
	var portable *age.X25519Identity
	if opts.NewPortableKey {
		var err error
		if portable, err = age.GenerateX25519Identity(); err != nil {
			return nil, nil, err
		}
		out = append(out, portable.Recipient())
	}
	if len(out) == 0 {
		return nil, nil, errors.New("capsule: time_and_key needs at least one recipient or a portable key")
	}
	return out, portable, nil
}

// encryptAll produces a complete in-memory age file.
func encryptAll(plaintext []byte, recipients ...age.Recipient) ([]byte, error) {
	var buf bytes.Buffer
	w, err := age.Encrypt(&buf, recipients...)
	if err != nil {
		return nil, err
	}
	_, writeErr := w.Write(plaintext)
	if err := errors.Join(writeErr, w.Close()); err != nil {
		return nil, err
	}
	return buf.Bytes(), nil
}