package capsule

import (
	"bytes"
	"context"
	"crypto/hmac"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"time"

	"filippo.io/age"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/accesskey"
	"g.activething.com/go/DateKeys/agewrap"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/profile"
	"g.activething.com/go/DateKeys/provider"
)

// OpenOptions configures Open.
type OpenOptions struct {
	// Registry holds the locally pinned profiles. Required.
	Registry profile.Registry
	// Extensions lists the critical extensions the application implements.
	Extensions extension.Registry
	// Source fetches the release. Required. Its answer is always verified
	// locally.
	Source provider.ReleaseSource
	// Identities are the caller's own X25519 identities, for time_and_key
	// capsules encrypted to known recipients.
	Identities []age.Identity
	// AccessKey is a portable .dkk, for time_and_key capsules.
	AccessKey *accesskey.AccessKey
	// Now is the clock. Required: no package of this module reads the wall
	// clock on its own. Open does not ask Source for a round whose time has
	// not been reached.
	Now func() time.Time
}

// Opened describes a capsule that Open decrypted completely.
type Opened struct {
	Inspection *Inspection
	Release    provider.Release
	// ControlCritical and ControlNoncritical are the extensions of the sealed
	// CONTROL_CBOR, only visible after opening.
	ControlCritical    []extension.Extension
	ControlNoncritical []extension.Extension
}

// Open runs the complete decryption flow of spec §63 and streams the
// plaintext to dst.
//
// Everything verifiable locally is checked before a release is requested or a
// secret is used (steps 1 to 8, the presence and capsule binding of access
// credentials, and the .dkk capsule_digest when r is seekable). The stanza
// rules are enforced again, as a MUST, by the identities that open each age
// file (steps 11, 13 and 17).
//
// Open returns nil only after age has authenticated the whole payload
// (step 18). On error, dst may hold a partial plaintext that MUST be
// discarded: write to a temporary file and publish it only on success
// (spec §56), as the datekeys CLI does.
func Open(ctx context.Context, dst io.Writer, r io.Reader, opts OpenOptions) (*Opened, error) {
	if opts.Source == nil {
		return nil, errors.New("capsule: OpenOptions.Source is required")
	}
	if opts.Now == nil {
		return nil, errors.New("capsule: OpenOptions.Now is required")
	}
	start, seekable := int64(0), false
	if s, ok := r.(io.Seeker); ok {
		if pos, err := s.Seek(0, io.SeekCurrent); err == nil {
			start, seekable = pos, true
		}
	}

	// Steps 1 to 8.
	in, st, err := inspect(r, InspectOptions{Registry: opts.Registry, Extensions: opts.Extensions})
	out := &Opened{Inspection: in}
	if err != nil {
		return out, err
	}
	h, p := in.Header, in.Profile

	// Access credentials are checked before any network request.
	var ids []age.Identity
	if h.Policy == TimeAndKey {
		ids = append(ids, opts.Identities...)
		if k := opts.AccessKey; k != nil {
			if err := checkAccessKey(k, h, opts.Extensions); err != nil {
				return out, in.fail(9, "access credential", err)
			}
			if k.Verification != nil && seekable {
				payload, err := checkCapsuleDigest(r.(io.ReadSeeker), start, in.PayloadOffset, k.Verification.CapsuleDigest)
				if err != nil {
					return out, in.fail(9, "access credential", err)
				}
				st.payload = payload
			}
			id, err := k.Identity()
			if err != nil {
				return out, in.fail(9, "access credential", err)
			}
			ids = append(ids, id)
		}
		if len(ids) == 0 {
			return out, in.fail(9, "access credential", fmt.Errorf("capsule: time_and_key capsule and no identity or .dkk supplied: %w", datekeys.ErrAccessRequired))
		}
		in.pass(9, "access credential", fmt.Sprintf("%d identities to try", len(ids)))
	}

	// Step 9: obtain the release, never before its round time.
	cond := provider.Condition{Round: h.DateKey.Round}
	if now := opts.Now(); now.Before(in.UnlockAt) {
		err := fmt.Errorf("capsule: round %d is published at %s, it is %s: %w", cond.Round,
			in.UnlockAt.Format(time.RFC3339), now.UTC().Format(time.RFC3339), datekeys.ErrReleaseUnavailable)
		return out, in.fail(9, "release", err)
	}
	release, err := opts.Source.Fetch(ctx, p, cond)
	if err != nil {
		if datekeys.Code(err) == "" {
			err = fmt.Errorf("capsule: %v: %w", err, datekeys.ErrReleaseUnavailable)
		}
		return out, in.fail(9, "release", err)
	}
	in.pass(9, "release", fmt.Sprintf("round %d obtained", release.Round))

	// Step 10: verify the release locally.
	if err := provider.Verify(p, cond, release); err != nil {
		return out, in.fail(10, "release verification", err)
	}
	out.Release = release
	in.pass(10, "release verification", "BLS signature valid under the pinned key")

	// Step 11: open OUTER_TIME_AGE with the strict tlock identity.
	timeID, err := agewrap.NewTimeIdentity(p, cond.Round, release)
	if err != nil {
		return out, in.fail(11, "open sealed control", err)
	}
	inner, err := decryptAll(st.sealed, timeID)
	if err != nil {
		return out, in.fail(11, "open sealed control", err)
	}
	defer clear(inner)
	in.pass(11, "open sealed control", "one tlock stanza, header MAC valid")

	// Step 12: the structure must match access_policy (spec §36).
	var controlBytes []byte
	switch h.Policy {
	case TimeOnly:
		if looksLikeAge(inner) {
			return out, in.fail(12, "policy structure", fmt.Errorf("capsule: time_only capsule seals an age file: %w", datekeys.ErrPolicyStructureMismatch))
		}
		controlBytes = inner
		in.pass(12, "policy structure", "time_only: CONTROL_CBOR sealed directly")
	case TimeAndKey:
		stanzas, err := agewrap.Stanzas(bytes.NewReader(inner))
		if err != nil {
			return out, in.fail(12, "policy structure", fmt.Errorf("capsule: time_and_key capsule does not seal an age file: %w", datekeys.ErrPolicyStructureMismatch))
		}
		if err := agewrap.CheckAccessStanzas(stanzas); err != nil {
			return out, in.fail(12, "policy structure", err)
		}
		in.pass(12, "policy structure", fmt.Sprintf("time_and_key: INNER_ACCESS_AGE with %d X25519 stanzas", len(stanzas)))

		// Step 13: open INNER_ACCESS_AGE with the caller's identities.
		accessID, err := agewrap.NewAccessIdentity(ids...)
		if err != nil {
			return out, in.fail(13, "open access layer", err)
		}
		if controlBytes, err = decryptAll(inner, accessID); err != nil {
			return out, in.fail(13, "open access layer", err)
		}
		defer clear(controlBytes)
		in.pass(13, "open access layer", "identity matched exactly one stanza")
	}

	// Step 14: parse the canonical CONTROL_CBOR.
	control, err := DecodeControl(controlBytes)
	if err != nil {
		return out, in.fail(14, "control", err)
	}
	defer clear(control.PayloadIdentity[:])
	if err := extension.CheckCritical(control.Critical, opts.Extensions); err != nil {
		return out, in.fail(14, "control", fmt.Errorf("capsule: CONTROL_CBOR: %w", err))
	}
	out.ControlCritical, out.ControlNoncritical = control.Critical, control.Noncritical
	in.pass(14, "control", "canonical CONTROL_CBOR")

	// Step 15: verify header_binding over the exact stored bytes.
	binding := HeaderBinding(st.prelude, in.PublicHeader)
	if !hmac.Equal(binding[:], control.HeaderBinding[:]) {
		return out, in.fail(15, "header binding", fmt.Errorf("capsule: header_binding does not match PRELUDE || PUBLIC_HEADER: %w", datekeys.ErrHeaderBinding))
	}
	in.pass(15, "header binding", "matches")

	// Steps 16 and 17: recover I_PAYLOAD and open PAYLOAD_AGE with it.
	payloadID, err := agewrap.NewPayloadIdentity(control.PayloadIdentity[:])
	if err != nil {
		return out, in.fail(16, "payload identity", err)
	}
	in.pass(16, "payload identity", "I_PAYLOAD recovered")
	pr, err := age.Decrypt(st.payload, payloadID)
	if err != nil {
		return out, in.fail(17, "open payload", classify("PAYLOAD_AGE", err))
	}
	if _, err := io.Copy(dst, pr); err != nil {
		return out, in.fail(17, "open payload", classify("PAYLOAD_AGE", err))
	}

	// Step 18: age completed without error.
	in.pass(18, "commit", "payload authenticated completely")
	return out, nil
}

// checkAccessKey validates a .dkk against the capsule before it is used.
func checkAccessKey(k *accesskey.AccessKey, h *Header, reg extension.Registry) error {
	if k.CapsuleID != h.CapsuleID {
		return fmt.Errorf("capsule: the .dkk is for capsule %x, this is %x: %w", k.CapsuleID, h.CapsuleID, datekeys.ErrAccessInvalid)
	}
	if err := extension.CheckCritical(k.Critical, reg); err != nil {
		return fmt.Errorf("capsule: .dkk: %w", err)
	}
	return nil
}

// checkCapsuleDigest compares the .dkk capsule_digest with SHA-256 of the
// .dkc (spec §43), then repositions r at the payload. The digest is a fast
// failure for a wrong file, not a security property.
func checkCapsuleDigest(r io.ReadSeeker, start, payloadOffset int64, want []byte) (io.Reader, error) {
	if _, err := r.Seek(start, io.SeekStart); err != nil {
		return nil, fmt.Errorf("capsule: %w", err)
	}
	sum := sha256.New()
	if _, err := io.Copy(sum, r); err != nil {
		return nil, fmt.Errorf("capsule: %w", err)
	}
	if !hmac.Equal(sum.Sum(nil), want) {
		return nil, fmt.Errorf("capsule: the .dkk capsule_digest does not match this .dkc: %w", datekeys.ErrAccessInvalid)
	}
	if _, err := r.Seek(start+payloadOffset, io.SeekStart); err != nil {
		return nil, fmt.Errorf("capsule: %w", err)
	}
	return r, nil
}

// decryptAll opens a bounded, in-memory age file. The plaintext is never
// longer than the ciphertext.
func decryptAll(ciphertext []byte, id age.Identity) ([]byte, error) {
	r, err := age.Decrypt(bytes.NewReader(ciphertext), id)
	if err != nil {
		return nil, classify("age", err)
	}
	out, err := io.ReadAll(io.LimitReader(r, int64(len(ciphertext))))
	if err != nil {
		clear(out)
		return nil, classify("age", err)
	}
	return out, nil
}

// classify keeps the normative error an identity returned from Unwrap, and
// maps every other age failure (malformed header, bad header MAC, STREAM
// authentication, truncation, trailing data) to ErrIntegrity.
func classify(what string, err error) error {
	if datekeys.Code(err) != "" {
		return fmt.Errorf("capsule: %s: %w", what, err)
	}
	return fmt.Errorf("capsule: %s: %v: %w", what, err, datekeys.ErrIntegrity)
}