package capsule_test

import (
	"bytes"
	"context"
	"encoding/hex"
	"errors"
	"fmt"
	"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
}

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 := capsule.Encrypt(&dkc, strings.NewReader("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 := capsule.Encrypt(&b, strings.NewReader("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)
	}
}