// Command genfixtures generates the official DateKeys vectors and fixtures
// (spec §65-§68) into testdata/.
//
// Fixtures are generated once, over rounds that are already published, and
// then committed: age randomness cannot be injected through its public API,
// so they are decryption and validation fixtures, not byte-reproducible
// encryption outputs (spec §67). Existing fixtures are never overwritten
// unless -force (every fixture) or -only (the named ones) is given; vectors
// are always regenerated, and the tests fail if the implementation stops
// reproducing the committed ones.
//
// The .dkk with an extension (spec §68) is derived from the portable .dkk of
// time_and_key_portable, so it is regenerated whenever its source is.
//
// Derived from the fixtures and always regenerated, like the vectors: the
// frozen "datekeys inspect -json" output of each .dkc (<name>.inspect.json),
// the exported mutation corpus (vectors/mutations.json) and the differential
// corpus of the pre-unlock checks (vectors/inspect_differential.json). The
// mutations whose capsule is built with age randomness keep the bytes
// recorded in the committed mutations.json; -force, or -only mutations,
// builds them afresh.
//
//	go run ./internal/testkit/genfixtures -out testdata
//	go run ./internal/testkit/genfixtures -out testdata -only time_only_extensions
package main

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/fs"
	"log"
	"os"
	"path/filepath"
	"strings"
	"time"

	"filippo.io/age"

	"g.activething.com/go/DateKeys/accesskey"
	"g.activething.com/go/DateKeys/agewrap"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/datekey"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/internal/inspectview"
	"g.activething.com/go/DateKeys/internal/testkit"
	"g.activething.com/go/DateKeys/profile"
)

func main() {
	out := flag.String("out", "testdata", "output directory")
	force := flag.Bool("force", false, "overwrite every existing fixture")
	only := flag.String("only", "", "comma-separated fixture names to regenerate, overwriting them, and \"mutations\" to rebuild the randomly built mutations; every other fixture is left untouched")
	flag.Parse()
	sel, err := selection(*force, *only)
	if err != nil {
		log.Fatal(err)
	}
	if err := vectors(filepath.Join(*out, "vectors")); err != nil {
		log.Fatal(err)
	}
	if err := fixtures(filepath.Join(*out, "fixtures"), sel); err != nil {
		log.Fatal(err)
	}
	if err := derived(*out, sel); err != nil {
		log.Fatal(err)
	}
}

// selector decides which fixtures are (re)generated.
type selector struct {
	force bool            // overwrite every fixture
	only  map[string]bool // when non-empty, regenerate exactly these
}

func selection(force bool, only string) (selector, error) {
	sel := selector{force: force, only: map[string]bool{}}
	if only == "" {
		return sel, nil
	}
	known := map[string]bool{extDKK: true, mutationsName: true}
	for _, s := range specs() {
		known[s.name] = true
	}
	for _, name := range strings.Split(only, ",") {
		if !known[name] {
			return sel, fmt.Errorf("-only: unknown fixture %q", name)
		}
		sel.only[name] = true
	}
	return sel, nil
}

// generate reports whether the fixture name, whose main file is path, is
// written.
func (s selector) generate(name, path string) bool {
	if len(s.only) > 0 {
		return s.only[name]
	}
	_, err := os.Stat(path)
	return s.force || err != nil
}

func vectors(dir string) error {
	pv, err := testkit.QuicknetProfileVector()
	if err != nil {
		return err
	}
	if err := testkit.WriteJSON(filepath.Join(dir, "profile_quicknet.json"), pv); err != nil {
		return err
	}
	if err := testkit.WriteJSON(filepath.Join(dir, "quicknet_rounds.json"), testkit.RoundVectors()); err != nil {
		return err
	}
	if err := testkit.WriteJSON(filepath.Join(dir, "dk1.json"), testkit.DK1Vectors()); err != nil {
		return err
	}
	cv, err := testkit.CBORVectors()
	if err != nil {
		return err
	}
	if err := testkit.WriteJSON(filepath.Join(dir, "cbor.json"), cv); err != nil {
		return err
	}
	iv, err := testkit.IBEVectors()
	if err != nil {
		return err
	}
	return testkit.WriteJSON(filepath.Join(dir, "tlock_ibe.json"), iv)
}

// mutationsName is the -only name that rebuilds the capsules of the
// mutations built with age randomness.
const mutationsName = "mutations"

// derived writes what is computed from the fixtures: the inspect outputs,
// the mutation corpus and the differential corpus.
func derived(out string, sel selector) error {
	fixtureDir := filepath.Join(out, "fixtures")
	var names []string
	for _, s := range specs() {
		names = append(names, s.name)
		if err := inspectOutput(fixtureDir, s.name); err != nil {
			return fmt.Errorf("%s: %w", s.name, err)
		}
	}

	path := filepath.Join(out, "vectors", "mutations.json")
	var frozen *testkit.MutationFile
	if !sel.force && !sel.only[mutationsName] {
		var f testkit.MutationFile
		switch err := testkit.ReadJSON(path, &f); {
		case err == nil:
			frozen = &f
		case !errors.Is(err, fs.ErrNotExist):
			return err
		}
	}
	m, err := testkit.MutationCorpus(fixtureDir, frozen)
	if err != nil {
		return err
	}
	if err := testkit.WriteJSONEdits(path, m); err != nil {
		return err
	}

	d, err := testkit.InspectDifferential(fixtureDir, names)
	if err != nil {
		return err
	}
	return testkit.WriteDifferential(filepath.Join(out, "vectors", "inspect_differential.json"), d)
}

// inspectOutput writes <name>.inspect.json: the output of
// "datekeys inspect -json -in <name>.dkc" run in the fixture directory.
func inspectOutput(dir, name string) error {
	file := name + ".dkc"
	dkc, err := os.ReadFile(filepath.Join(dir, file))
	if err != nil {
		return err
	}
	in, err := capsule.Inspect(bytes.NewReader(dkc), capsule.InspectOptions{Registry: testkit.Registry()})
	var b bytes.Buffer
	if err := inspectview.New(file, in, err).WriteJSON(&b); err != nil {
		return err
	}
	return os.WriteFile(filepath.Join(dir, name+".inspect.json"), b.Bytes(), 0o644)
}

type spec struct {
	name, description string
	round             uint64
	policy            capsule.Policy
	recipients        int
	portable          bool
	plaintext         []byte
	headerExt         []extension.Extension
	controlExt        []extension.Extension
}

// Extension data of the fixtures (spec §54, §72): the header carries the raw
// UTF-8 bytes of a label, which are not CBOR; the control carries
// {0: 7, 1: "sealed"} in the CBOR profile of spec §58; the .dkk carries
// {0: "hand"}. The base protocol decodes none of them.
var (
	headerExtData  = []byte("public label")
	controlExtData = mustHex("a2000701667365616c6564")
	dkkExtData     = mustHex("a1006468616e64")
)

func mustHex(s string) []byte {
	b, err := hex.DecodeString(s)
	if err != nil {
		panic(err)
	}
	return b
}

func mustExt(id string, version uint64, data []byte) extension.Extension {
	e, err := extension.New(id, version, data)
	if err != nil {
		panic(err)
	}
	return e
}

func specs() []spec {
	large := []byte(strings.Repeat("DateKeys fixture: this plaintext spans more than one 64 KiB age STREAM chunk.\n", 1000))
	hExt := mustExt("org.example.label", 1, headerExtData)
	cExt := mustExt("org.example.note", 2, controlExtData)
	return []spec{
		{name: "time_only", description: "time_only capsule, two STREAM chunks, no extensions", round: 1000, policy: capsule.TimeOnly, plaintext: large},
		{name: "time_only_extensions", description: "time_only capsule with a noncritical PUBLIC_HEADER extension and a noncritical CONTROL_CBOR extension", round: 2000, policy: capsule.TimeOnly, plaintext: []byte("DateKeys fixture with extensions.\n"), headerExt: []extension.Extension{hExt}, controlExt: []extension.Extension{cExt}},
		{name: "time_and_key_portable", description: "time_and_key capsule whose only recipient is a portable .dkk", round: 1000, policy: capsule.TimeAndKey, portable: true, plaintext: []byte("DateKeys fixture opened with a portable .dkk.\n")},
		{name: "time_and_key_recipients", description: "time_and_key capsule for two known X25519 recipients and a portable .dkk", round: 1001, policy: capsule.TimeAndKey, recipients: 2, portable: true, plaintext: []byte("DateKeys fixture for several recipients.\n")},
		{name: "empty_payload", description: "time_only capsule with an empty payload", round: 1001, policy: capsule.TimeOnly, plaintext: []byte{}},
	}
}

func fixtures(dir string, sel selector) error {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	regenerated := map[string]bool{}
	for _, s := range specs() {
		path := filepath.Join(dir, s.name+".dkc")
		if !sel.generate(s.name, path) {
			log.Printf("leaving %s untouched", path)
			continue
		}
		if err := generate(dir, s); err != nil {
			return fmt.Errorf("%s: %w", s.name, err)
		}
		regenerated[s.name] = true
		log.Printf("generated %s", path)
	}
	path := filepath.Join(dir, extDKK+".dkk")
	if !sel.generate(extDKK, path) && !regenerated[extDKKSource] {
		log.Printf("leaving %s untouched", path)
		return nil
	}
	if err := deriveDKK(dir); err != nil {
		return fmt.Errorf("%s: %w", extDKK, err)
	}
	log.Printf("generated %s", path)
	return nil
}

// extDKK is the .dkk vector with an extension (spec §68): the portable
// credential of extDKKSource re-issued with a noncritical extension. It keeps
// the credential_id, the key and the capsule_digest, so its bytes are a
// function of the source .dkk.
const (
	extDKK       = "time_and_key_portable_extension"
	extDKKSource = "time_and_key_portable"
)

func deriveDKK(dir string) error {
	src, err := os.ReadFile(filepath.Join(dir, extDKKSource+".dkk"))
	if err != nil {
		return err
	}
	k, err := accesskey.Decode(bytes.NewReader(src))
	if err != nil {
		return err
	}
	k.Noncritical = []extension.Extension{mustExt("org.example.delivery", 1, dkkExtData)}
	var kb bytes.Buffer
	if err := accesskey.Encode(&kb, k); err != nil {
		return err
	}
	back, err := accesskey.Decode(bytes.NewReader(kb.Bytes()))
	if err != nil {
		return err
	}

	// The credential must open its capsule through the public API.
	dkcFile := extDKKSource + ".dkc"
	dkc, err := os.ReadFile(filepath.Join(dir, dkcFile))
	if err != nil {
		return err
	}
	reg := testkit.Registry()
	in, err := capsule.Inspect(bytes.NewReader(dkc), capsule.InspectOptions{Registry: reg})
	if err != nil {
		return err
	}
	oo := capsule.OpenOptions{Registry: reg, Source: testkit.NewSource(testkit.Release(in.Header.DateKey.Round)),
		AccessKey: back, Now: testkit.Fixed(in.UnlockAt)}
	if _, err := capsule.Open(context.Background(), io.Discard, bytes.NewReader(dkc), oo); err != nil {
		return fmt.Errorf("the .dkk does not open %s: %w", dkcFile, err)
	}

	ksum := sha256.Sum256(kb.Bytes())
	kf := testkit.DKKFixture{
		Description:    "portable X25519 .dkk of " + dkcFile + " with a noncritical extension: the credential of " + extDKKSource + ".dkk re-issued with org.example.delivery",
		Spec:           testkit.SpecVersion,
		File:           extDKK + ".dkk",
		SHA256:         hex.EncodeToString(ksum[:]),
		CredentialID:   hex.EncodeToString(back.CredentialID[:]),
		CapsuleID:      hex.EncodeToString(back.CapsuleID[:]),
		AccessType:     back.Type,
		Material:       hex.EncodeToString(back.Material),
		CapsuleDigest:  hex.EncodeToString(back.Verification.CapsuleDigest),
		Extensions:     exts(false, back.Noncritical),
		Capsule:        dkcFile,
		ExpectedResult: "opens INNER_ACCESS_AGE of " + dkcFile + " and yields its CONTROL_CBOR",
	}
	if err := os.WriteFile(filepath.Join(dir, kf.File), kb.Bytes(), 0o644); err != nil {
		return err
	}
	return testkit.WriteJSON(filepath.Join(dir, extDKK+".dkk.json"), kf)
}

func generate(dir string, s spec) error {
	p := profile.Quicknet()
	reg := testkit.Registry()
	unlock, err := datekey.RoundTime(p, s.round)
	if err != nil {
		return err
	}
	opts := capsule.EncryptOptions{
		Profile: p, UnlockAt: unlock, Policy: s.policy, NewPortableKey: s.portable,
		Noncritical: s.headerExt, ControlNoncritical: s.controlExt, Now: testkit.Fixed(testkit.Genesis()),
	}
	var ids []*age.X25519Identity
	for range s.recipients {
		id, err := age.GenerateX25519Identity()
		if err != nil {
			return err
		}
		ids = append(ids, id)
		opts.Recipients = append(opts.Recipients, id.Recipient())
	}
	var dkc bytes.Buffer
	res, err := capsule.Encrypt(&dkc, bytes.NewReader(s.plaintext), opts)
	if err != nil {
		return err
	}
	release := testkit.Release(s.round)

	// Recover every intermediate value by opening the fixture step by step.
	parts, err := testkit.Split(dkc.Bytes())
	if err != nil {
		return err
	}
	timeID, err := agewrap.NewTimeIdentity(p, s.round, release)
	if err != nil {
		return err
	}
	inner, err := decrypt(parts.Sealed, timeID)
	if err != nil {
		return err
	}
	control := inner
	var innerStanzas []testkit.FixtureStanza
	var identities []string
	var dkkFile string
	var dkkBytes []byte
	if s.policy == capsule.TimeAndKey {
		st, err := agewrap.Stanzas(bytes.NewReader(inner))
		if err != nil {
			return err
		}
		innerStanzas = stanzas(st)
		var tryIDs []age.Identity
		for _, id := range ids {
			identities = append(identities, id.String())
			tryIDs = append(tryIDs, id)
		}
		if res.PortableKey != nil {
			var kb bytes.Buffer
			if err := accesskey.Encode(&kb, res.PortableKey); err != nil {
				return err
			}
			dkkBytes = kb.Bytes()
			dkkFile = s.name + ".dkk"
			kid, err := res.PortableKey.Identity()
			if err != nil {
				return err
			}
			tryIDs = append(tryIDs, kid)
		}
		accessID, err := agewrap.NewAccessIdentity(tryIDs...)
		if err != nil {
			return err
		}
		if control, err = decrypt(inner, accessID); err != nil {
			return err
		}
	}
	ctrl, err := capsule.DecodeControl(control)
	if err != nil {
		return err
	}
	outer, err := agewrap.Stanzas(bytes.NewReader(parts.Sealed))
	if err != nil {
		return err
	}
	payload, err := agewrap.Stanzas(bytes.NewReader(parts.Payload))
	if err != nil {
		return err
	}

	// The fixture must open through the public API with the embedded release.
	oo := capsule.OpenOptions{Registry: reg, Source: testkit.NewSource(release), Now: testkit.Fixed(unlock)}
	for _, id := range ids {
		oo.Identities = append(oo.Identities, id)
	}
	if s.policy == capsule.TimeAndKey && len(ids) == 0 {
		oo.AccessKey = res.PortableKey
	}
	var plain bytes.Buffer
	opened, err := capsule.Open(context.Background(), &plain, bytes.NewReader(dkc.Bytes()), oo)
	if err != nil {
		return fmt.Errorf("fixture does not open: %w", err)
	}
	if !bytes.Equal(plain.Bytes(), s.plaintext) {
		return fmt.Errorf("fixture plaintext mismatch")
	}

	sum := sha256.Sum256(dkc.Bytes())
	psum := sha256.Sum256(s.plaintext)
	f := testkit.DKCFixture{
		Description:      s.description,
		Spec:             testkit.SpecVersion,
		File:             s.name + ".dkc",
		SHA256:           hex.EncodeToString(sum[:]),
		Release:          testkit.FixtureRelease{Round: release.Round, Signature: hex.EncodeToString(release.Signature)},
		Prelude:          hex.EncodeToString(parts.Prelude),
		PublicHeader:     hex.EncodeToString(parts.Header),
		DateKey:          res.DateKey.Compact(),
		CapsuleID:        hex.EncodeToString(res.CapsuleID[:]),
		AccessPolicy:     s.policy.String(),
		Structure:        s.policy.String(),
		UnlockAt:         unlock.Format(time.RFC3339),
		HeaderBinding:    hex.EncodeToString(ctrl.HeaderBinding[:]),
		OuterStanzas:     stanzas(outer),
		PayloadStanzas:   stanzas(payload),
		InnerStanzas:     innerStanzas,
		AccessKeyFile:    dkkFile,
		Identities:       identities,
		ControlCBOR:      hex.EncodeToString(control),
		PayloadIdentity:  hex.EncodeToString(ctrl.PayloadIdentity[:]),
		PlaintextFile:    s.name + ".plaintext",
		PlaintextSHA256:  hex.EncodeToString(psum[:]),
		HeaderExtensions: exts(false, s.headerExt),
		ControlExt:       exts(false, s.controlExt),
	}
	for _, c := range opened.Inspection.Checks {
		f.Stages = append(f.Stages, testkit.FixtureStage{Step: c.Step, Name: c.Name, OK: c.OK, Error: c.Error})
	}
	if err := os.WriteFile(filepath.Join(dir, f.File), dkc.Bytes(), 0o644); err != nil {
		return err
	}
	if err := os.WriteFile(filepath.Join(dir, f.PlaintextFile), s.plaintext, 0o644); err != nil {
		return err
	}
	if err := testkit.WriteJSON(filepath.Join(dir, s.name+".json"), f); err != nil {
		return err
	}
	if dkkBytes == nil {
		return nil
	}
	k := res.PortableKey
	ksum := sha256.Sum256(dkkBytes)
	kf := testkit.DKKFixture{
		Description:    "portable X25519 .dkk of " + f.File,
		Spec:           testkit.SpecVersion,
		File:           dkkFile,
		SHA256:         hex.EncodeToString(ksum[:]),
		CredentialID:   hex.EncodeToString(k.CredentialID[:]),
		CapsuleID:      hex.EncodeToString(k.CapsuleID[:]),
		AccessType:     k.Type,
		Material:       hex.EncodeToString(k.Material),
		CapsuleDigest:  hex.EncodeToString(k.Verification.CapsuleDigest),
		Capsule:        f.File,
		ExpectedResult: "opens INNER_ACCESS_AGE of " + f.File + " and yields its CONTROL_CBOR",
	}
	if err := os.WriteFile(filepath.Join(dir, dkkFile), dkkBytes, 0o644); err != nil {
		return err
	}
	return testkit.WriteJSON(filepath.Join(dir, s.name+".dkk.json"), kf)
}

func decrypt(file []byte, id age.Identity) ([]byte, error) {
	r, err := age.Decrypt(bytes.NewReader(file), id)
	if err != nil {
		return nil, err
	}
	var b bytes.Buffer
	if _, err := b.ReadFrom(r); err != nil {
		return nil, err
	}
	return b.Bytes(), nil
}

func stanzas(in []*age.Stanza) []testkit.FixtureStanza {
	out := make([]testkit.FixtureStanza, len(in))
	for i, s := range in {
		out[i] = testkit.FixtureStanza{Type: s.Type, Args: s.Args}
	}
	return out
}

func exts(critical bool, in []extension.Extension) []testkit.FixtureExt {
	var out []testkit.FixtureExt
	for _, e := range in {
		out = append(out, testkit.FixtureExt{Critical: critical, ID: e.ID, Version: e.Version, Data: hex.EncodeToString(e.Data)})
	}
	return out
}