package capsule_test

import (
	"bytes"
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	"slices"
	"strings"
	"testing"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/internal/cbortest"
	"g.activething.com/go/DateKeys/internal/testkit"
)

// strictRegistry knows every org.example.* extension at version 1 and accepts
// only the data "ok" (spec §54, §72).
type strictRegistry struct{}

func (strictRegistry) Known(id string, v uint64) bool {
	return v == 1 && strings.HasPrefix(id, "org.example.")
}

func (strictRegistry) ValidateData(e extension.Extension) error {
	if string(e.Data) != "ok" {
		return fmt.Errorf("data %x is not \"ok\"", e.Data)
	}
	return nil
}

// place is one extension array of one object.
type place struct {
	obj extension.Object
	arr extension.Array
}

// placedRegistry is strictRegistry with the placement of its registrations
// (spec §72): each org.example.* extension is registered only in the places
// listed for it.
type placedRegistry struct {
	strictRegistry
	places map[string][]place
}

func (r placedRegistry) RegisteredIn(id string, _ uint64, obj extension.Object, arr extension.Array) bool {
	return slices.Contains(r.places[id], place{obj, arr})
}

func mustExt(t *testing.T, id string, data []byte) extension.Extension {
	t.Helper()
	e, err := extension.New(id, 1, data)
	if err != nil {
		t.Fatal(err)
	}
	return e
}

func mustUnhex(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

// Spec §76, case 5: control data made of 14 or 15 nested CBOR arrays was
// sealed by Encrypt and rejected by Open at step 14, after the unlock, which
// made the capsule unrecoverable. Data is opaque now: it seals and opens.
func TestNestedDataSealsAndOpens(t *testing.T) {
	for _, depth := range []int{14, 15, 40} {
		data := mustUnhex(t, strings.Repeat("81", depth)+"00")
		opts := past(t, 1000)
		opts.Noncritical = []extension.Extension{mustExt(t, "org.example.nested", data)}
		opts.ControlNoncritical = []extension.Extension{mustExt(t, "org.example.nested", data)}
		var dkc bytes.Buffer
		if _, err := encrypt(t, &dkc, "nested", opts); err != nil {
			t.Fatalf("depth %d: %v", depth, err)
		}
		var out bytes.Buffer
		opened, err := capsule.Open(context.Background(), &out, bytes.NewReader(dkc.Bytes()), defaultOpen(1000))
		if err != nil || out.String() != "nested" {
			t.Fatalf("depth %d: sealed but does not open: %v", depth, err)
		}
		if !bytes.Equal(opened.ControlNoncritical[0].Data, data) || !bytes.Equal(opened.Inspection.Header.Noncritical[0].Data, data) {
			t.Fatalf("depth %d: data changed", depth)
		}
	}
}

// Spec §76, case 3: header data {NaN: 0, NaN: 1} (a2f97e0000f97e0001) made
// the verdict on one capsule depend on map iteration order. The base protocol
// no longer decodes data, so every run gives the same verdict.
func TestNaNKeyedDataHasOneVerdict(t *testing.T) {
	data := mustUnhex(t, "a2f97e0000f97e0001")
	opts := past(t, 1000)
	opts.Noncritical = []extension.Extension{mustExt(t, "org.example.nan", data)}
	var b bytes.Buffer
	if _, err := encrypt(t, &b, "nan", opts); err != nil {
		t.Fatal(err)
	}
	dkc := b.Bytes()
	for i := range 500 {
		in, err := capsule.Inspect(bytes.NewReader(dkc), capsule.InspectOptions{Registry: testkit.Registry()})
		if err != nil || !bytes.Equal(in.Header.Noncritical[0].Data, data) {
			t.Fatalf("Inspect run %d: %v", i, err)
		}
	}
	for i := range 200 {
		if _, err := capsule.Open(context.Background(), &bytes.Buffer{}, bytes.NewReader(dkc), defaultOpen(1000)); err != nil {
			t.Fatalf("Open run %d: %v", i, err)
		}
	}
}

// Spec §76, case 6: a PUBLIC_HEADER of about 880 KB with 40 000 + 40 000
// extensions took 8.3 s in the pairwise disjointness check. The 64-extension
// limit rejects it first.
func TestHugeExtensionArraysAreRejected(t *testing.T) {
	f := loadFixture(t, "time_only")
	parts, _ := testkit.Split(f.dkc)
	m, err := cbortest.UnmarshalMap(parts.Header)
	if err != nil {
		t.Fatal(err)
	}
	const n = 40_000
	crit, non := make([]any, n), make([]any, n)
	for i := range n {
		crit[i] = map[uint64]any{0: fmt.Sprintf("c%04x", i), 1: uint64(1)}
		non[i] = map[uint64]any{0: fmt.Sprintf("n%04x", i), 1: uint64(1)}
	}
	m[5], m[6] = crit, non
	h, err := cbortest.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}
	if len(h) > capsule.MaxPublicHeaderLen {
		t.Fatalf("header of %d bytes", len(h))
	}
	dkc := testkit.Reframe(parts.Prelude, h, parts.Sealed, parts.Payload)
	in, err := capsule.Inspect(bytes.NewReader(dkc), capsule.InspectOptions{Registry: testkit.Registry()})
	if last := in.Checks[len(in.Checks)-1]; !errors.Is(err, datekeys.ErrNonCanonicalCBOR) || last.Step != 4 {
		t.Fatalf("Inspect: %v at %+v", err, last)
	}
}

// Spec §54: a known noncritical extension with invalid data leaves the
// capsule valid; it is reported as unusable, in the object where it is.
func TestUnusableNoncriticalExtensions(t *testing.T) {
	bad, good := []byte("ko"), []byte("ok")
	dkc, o := build(t, testkit.Build{
		HeaderNoncritical:  []extension.Extension{mustExt(t, "org.example.header", bad), mustExt(t, "org.example.fine", good)},
		ControlNoncritical: []extension.Extension{mustExt(t, "org.example.control", bad), {ID: "org.example.nodata", Version: 1}, mustExt(t, "org.other", bad)},
	})
	o.Extensions = strictRegistry{}
	var out bytes.Buffer
	opened, err := capsule.Open(context.Background(), &out, bytes.NewReader(dkc), o)
	if err != nil || out.String() != "malicious creator" {
		t.Fatalf("a noncritical extension with invalid data failed the capsule: %v", err)
	}
	in := opened.Inspection
	if u := in.UnusableExtensions; len(u) != 1 || u[0].ID != "org.example.header" || !errors.Is(u[0].Err, datekeys.ErrExtensionDataInvalid) {
		t.Fatalf("header: %+v", u)
	}
	// An extension without data is known too: the registry rejects its
	// absent data. org.other is unknown and ignored.
	if u := opened.UnusableControlExtensions; len(u) != 2 || u[0].ID != "org.example.control" || u[1].ID != "org.example.nodata" {
		t.Fatalf("control: %+v", u)
	}
	for _, c := range in.Checks {
		if (c.Step == 4 || c.Step == 14) && !strings.Contains(c.Detail, "unusable noncritical") {
			t.Fatalf("step %d does not report the unusable extensions: %s", c.Step, c.Detail)
		}
	}
	// Without a validating registry nothing is unusable.
	o.Extensions = nil
	opened, err = capsule.Open(context.Background(), &bytes.Buffer{}, bytes.NewReader(dkc), o)
	if err != nil || opened.Inspection.UnusableExtensions != nil || opened.UnusableControlExtensions != nil {
		t.Fatalf("base protocol: %v", err)
	}
}

// The official .dkk with an extension (spec §68) opens its capsule; a
// registry that rejects its data reports it without refusing the credential.
func TestAccessKeyFixtureWithExtension(t *testing.T) {
	f := loadFixture(t, "time_and_key_portable")
	k := loadAccessKey(t, "time_and_key_portable_extension")
	if len(k.Noncritical) != 1 || k.Noncritical[0].ID != "org.example.delivery" {
		t.Fatalf("extensions %+v", k.Noncritical)
	}
	o := f.openOptions(t)
	o.AccessKey = k
	var out bytes.Buffer
	opened, err := capsule.Open(context.Background(), &out, bytes.NewReader(f.dkc), o)
	if err != nil || !bytes.Equal(out.Bytes(), f.plaintext) || opened.UnusableAccessKeyExtensions != nil {
		t.Fatalf("%v", err)
	}
	o.Extensions = strictRegistry{}
	out.Reset()
	opened, err = capsule.Open(context.Background(), &out, bytes.NewReader(f.dkc), o)
	if err != nil || !bytes.Equal(out.Bytes(), f.plaintext) {
		t.Fatalf("%v", err)
	}
	if u := opened.UnusableAccessKeyExtensions; len(u) != 1 || u[0].ID != "org.example.delivery" || !errors.Is(u[0].Err, datekeys.ErrExtensionDataInvalid) {
		t.Fatalf("unusable: %+v", u)
	}
}

// Spec §54, §72: a known extension that appears in an object or array it is
// not registered for is treated there as unknown. An extension registered for
// the critical_extensions of CONTROL_CBOR only and copied into PUBLIC_HEADER,
// which anyone can write (§55.1), or into a .dkk, is rejected at step 4 or
// 9; one registered as noncritical only is rejected in a critical array, and
// a noncritical copy outside its registration, in PUBLIC_HEADER or in a
// .dkk, is ignored, its data unchecked. A registry that is not an
// extension.Placement keeps today's behaviour: it knows its extensions
// everywhere.
func TestExtensionPlacement(t *testing.T) {
	sealed := []extension.Extension{mustExt(t, "org.example.sealed", []byte("ok"))}
	note := []extension.Extension{mustExt(t, "org.example.note", []byte("ok"))}
	badNote := []extension.Extension{mustExt(t, "org.example.note", []byte("ko"))}
	reg := placedRegistry{places: map[string][]place{
		"org.example.sealed":   {{extension.Control, extension.Critical}},
		"org.example.note":     {{extension.Control, extension.Noncritical}},
		"org.example.delivery": {{extension.AccessKey, extension.Noncritical}},
	}}
	both := []struct {
		name   string
		reg    extension.Registry
		placed bool
	}{{"placement", reg, true}, {"no placement", strictRegistry{}, false}}

	// Where it is registered, the capsule opens.
	dkc, _ := build(t, testkit.Build{ControlCritical: sealed})
	if step, _, err := openStep(t, dkc, capsule.OpenOptions{Extensions: reg}); err != nil {
		t.Fatalf("CONTROL_CBOR extension in CONTROL_CBOR: %v at step %d", err, step)
	}

	f := loadFixture(t, "time_and_key_portable")
	k := *f.dkk
	k.Critical = sealed
	for _, tc := range []struct {
		name string
		dkc  []byte
		o    capsule.OpenOptions
		step int
	}{
		{"CONTROL_CBOR extension copied into the critical_extensions of PUBLIC_HEADER", mustBuild(t, testkit.Build{HeaderCritical: sealed, ControlCritical: sealed}), capsule.OpenOptions{}, 4},
		{"CONTROL_CBOR extension in the critical_extensions of a .dkk", f.dkc, capsule.OpenOptions{AccessKey: &k, Now: testkit.Fixed(f.unlock(t))}, 9},
		{"noncritical-only extension in the critical_extensions of CONTROL_CBOR", mustBuild(t, testkit.Build{ControlCritical: note}), capsule.OpenOptions{}, 14},
	} {
		for _, r := range both {
			tc.o.Extensions = r.reg
			step, calls, err := openStep(t, tc.dkc, tc.o)
			if !r.placed {
				if err != nil {
					t.Errorf("%s, %s: %v at step %d", tc.name, r.name, err, step)
				}
				continue
			}
			expectStep(t, tc.name, step, err, datekeys.ErrExtensionCriticalUnknown, tc.step)
			if tc.step < 10 && calls != 0 {
				t.Errorf("%s: %d release requests", tc.name, calls)
			}
		}
	}

	// A noncritical copy outside its registration is ignored: its invalid
	// data is reported only in CONTROL_CBOR, where it is registered.
	dkc, o := build(t, testkit.Build{HeaderNoncritical: badNote, ControlNoncritical: badNote})
	for _, r := range both {
		o.Extensions = r.reg
		opened, err := capsule.Open(context.Background(), &bytes.Buffer{}, bytes.NewReader(dkc), o)
		if err != nil {
			t.Fatalf("%s: %v", r.name, err)
		}
		header, control := opened.Inspection.UnusableExtensions, opened.UnusableControlExtensions
		if len(control) != 1 || control[0].ID != "org.example.note" || r.placed != (header == nil) {
			t.Errorf("%s: unusable in PUBLIC_HEADER %+v, in CONTROL_CBOR %+v", r.name, header, control)
		}
	}

	// So is one in the noncritical_extensions of a .dkk: of two extensions
	// with invalid data there, only the one registered for the .dkk is
	// reported.
	kn := *f.dkk
	kn.Noncritical = []extension.Extension{mustExt(t, "org.example.delivery", []byte("ko")), badNote[0]}
	o = f.openOptions(t)
	o.AccessKey = &kn
	for _, r := range both {
		o.Extensions = r.reg
		opened, err := capsule.Open(context.Background(), &bytes.Buffer{}, bytes.NewReader(f.dkc), o)
		if err != nil {
			t.Fatalf("%s, .dkk: %v", r.name, err)
		}
		want := []string{"org.example.delivery"}
		if !r.placed {
			want = append(want, "org.example.note")
		}
		var got []string
		for _, u := range opened.UnusableAccessKeyExtensions {
			got = append(got, u.ID)
		}
		if !slices.Equal(got, want) {
			t.Errorf("%s: unusable in the .dkk %v, want %v", r.name, got, want)
		}
	}
}

func mustBuild(t *testing.T, b testkit.Build) []byte {
	t.Helper()
	dkc, _ := build(t, b)
	return dkc
}