// 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 is given; vectors are always regenerated, and the tests fail
// if the implementation stops reproducing the committed ones.
//
//	go run ./internal/testkit/genfixtures -out testdata
package main

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

	"filippo.io/age"

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

func main() {
	out := flag.String("out", "testdata", "output directory")
	force := flag.Bool("force", false, "overwrite existing fixtures")
	flag.Parse()
	if err := vectors(filepath.Join(*out, "vectors")); err != nil {
		log.Fatal(err)
	}
	if err := fixtures(filepath.Join(*out, "fixtures"), *force); err != nil {
		log.Fatal(err)
	}
}

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
	}
	return testkit.WriteJSON(filepath.Join(dir, "dk1.json"), testkit.DK1Vectors())
}

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

func fixtures(dir string, force bool) error {
	if err := os.MkdirAll(dir, 0o755); err != nil {
		return err
	}
	large := []byte(strings.Repeat("DateKeys fixture: this plaintext spans more than one 64 KiB age STREAM chunk.\n", 1000))
	hExt, err := extension.New("org.example.label", 1, "public label")
	if err != nil {
		return err
	}
	cExt, err := extension.New("org.example.note", 2, map[string]any{"sealed": true, "n": 7})
	if err != nil {
		return err
	}
	specs := []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{}},
	}
	for _, s := range specs {
		path := filepath.Join(dir, s.name+".dkc")
		if _, err := os.Stat(path); err == nil && !force {
			log.Printf("keeping existing %s", path)
			continue
		}
		if err := generate(dir, s); err != nil {
			return fmt.Errorf("%s: %w", s.name, err)
		}
		log.Printf("generated %s", path)
	}
	return nil
}

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
}