package accesskey_test

import (
	"bytes"
	"encoding/binary"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"reflect"
	"slices"
	"strings"
	"testing"

	"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/internal/cbortest"
	"g.activething.com/go/DateKeys/internal/testkit"
	"g.activething.com/go/DateKeys/profile"
)

const fixtures = "../testdata/fixtures"

func loadDKK(t *testing.T, name string) ([]byte, testkit.DKKFixture) {
	t.Helper()
	var f testkit.DKKFixture
	if err := testkit.ReadJSON(filepath.Join(fixtures, name+".dkk.json"), &f); err != nil {
		t.Fatal(err)
	}
	b, err := os.ReadFile(filepath.Join(fixtures, f.File))
	if err != nil {
		t.Fatal(err)
	}
	return b, f
}

var fixtureNames = []string{"time_and_key_portable", "time_and_key_recipients", "time_and_key_portable_extension"}

// fixtureExts lists extensions as the fixtures record them, with the exact
// data bytes.
func fixtureExts(k *accesskey.AccessKey) []testkit.FixtureExt {
	var out []testkit.FixtureExt
	for i, e := range append(slices.Clone(k.Critical), k.Noncritical...) {
		out = append(out, testkit.FixtureExt{Critical: i < len(k.Critical), ID: e.ID, Version: e.Version, Data: hex.EncodeToString(e.Data)})
	}
	return out
}

// Spec §68: parse, validate and use the official .dkk fixtures.
func TestFixtures(t *testing.T) {
	withData := 0
	for _, name := range fixtureNames {
		t.Run(name, func(t *testing.T) {
			b, f := loadDKK(t, name)
			k, err := accesskey.Decode(bytes.NewReader(b))
			if err != nil {
				t.Fatal(err)
			}
			if hex.EncodeToString(k.CredentialID[:]) != f.CredentialID || hex.EncodeToString(k.CapsuleID[:]) != f.CapsuleID ||
				k.Type != f.AccessType || hex.EncodeToString(k.Material) != f.Material ||
				k.Verification == nil || hex.EncodeToString(k.Verification.CapsuleDigest) != f.CapsuleDigest {
				t.Fatalf("decoded values differ from the fixture: %v", k)
			}
			if got := fixtureExts(k); !reflect.DeepEqual(got, f.Extensions) {
				t.Fatalf("extensions differ:\n got %+v\nwant %+v", got, f.Extensions)
			}
			for _, e := range f.Extensions {
				if e.Data != "" {
					withData++
				}
			}
			// Encode(Decode(x)) == x.
			var out bytes.Buffer
			if err := accesskey.Encode(&out, k); err != nil || !bytes.Equal(out.Bytes(), b) {
				t.Fatal("re-encoding differs from the fixture bytes")
			}
			// Expected result: the identity opens the INNER_ACCESS_AGE of its capsule.
			var cf testkit.DKCFixture
			if err := testkit.ReadJSON(filepath.Join(fixtures, strings.TrimSuffix(f.Capsule, ".dkc")+".json"), &cf); err != nil {
				t.Fatal(err)
			}
			dkc, _ := os.ReadFile(filepath.Join(fixtures, f.Capsule))
			parts, _ := testkit.Split(dkc)
			timeID, _ := agewrap.NewTimeIdentity(testkitProfile(), cf.Release.Round, testkit.Release(cf.Release.Round))
			inner := mustDecrypt(t, parts.Sealed, timeID)
			id, err := k.Identity()
			if err != nil {
				t.Fatal(err)
			}
			acc, _ := agewrap.NewAccessIdentity(id)
			if control := mustDecrypt(t, inner, acc); hex.EncodeToString(control) != cf.ControlCBOR {
				t.Fatal("the .dkk does not yield the expected CONTROL_CBOR")
			}
		})
	}
	// Spec §68: at least one official .dkk carries an extension with data.
	if withData == 0 {
		t.Fatal("no .dkk fixture carries extension data")
	}
}

// The .dkk with an extension is the credential of time_and_key_portable.dkk
// plus one noncritical extension, whose data the base protocol never decodes.
func TestFixtureWithExtension(t *testing.T) {
	src, _ := loadDKK(t, "time_and_key_portable")
	b, _ := loadDKK(t, "time_and_key_portable_extension")
	k, err := accesskey.Decode(bytes.NewReader(b))
	if err != nil {
		t.Fatal(err)
	}
	if len(k.Critical) != 0 || len(k.Noncritical) != 1 || k.Noncritical[0].ID != "org.example.delivery" || k.Noncritical[0].Version != 1 ||
		hex.EncodeToString(k.Noncritical[0].Data) != "a1006468616e64" {
		t.Fatalf("extension %+v", k.Noncritical)
	}
	k.Noncritical = nil
	var out bytes.Buffer
	if err := accesskey.Encode(&out, k); err != nil || !bytes.Equal(out.Bytes(), src) {
		t.Fatal("without its extension the .dkk differs from time_and_key_portable.dkk")
	}
}

func TestSecretsAreNotPrinted(t *testing.T) {
	b, f := loadDKK(t, "time_and_key_portable")
	k, _ := accesskey.Decode(bytes.NewReader(b))
	for _, format := range []string{"%v", "%+v", "%#v", "%s"} {
		for _, v := range []any{k, *k} {
			if s := fmt.Sprintf(format, v); strings.Contains(s, f.Material) || !strings.Contains(s, "REDACTED") {
				t.Fatalf("%s leaks or hides nothing: %s", format, s)
			}
		}
	}
}

func frame(body []byte) []byte {
	pre := make([]byte, accesskey.PreludeSize)
	copy(pre, accesskey.Magic)
	pre[4] = accesskey.FramingVersion
	binary.BigEndian.PutUint32(pre[8:], uint32(len(body)))
	return append(pre, body...)
}

func TestDecodeRejects(t *testing.T) {
	good, _ := loadDKK(t, "time_and_key_portable")
	body := good[accesskey.PreludeSize:]
	w, err := cbortest.UnmarshalMap(body)
	if err != nil {
		t.Fatal(err)
	}
	with := func(edit func(m map[uint64]any)) []byte {
		m := map[uint64]any{}
		for k, v := range w {
			m[k] = v
		}
		edit(m)
		b, err := cbortest.Marshal(m)
		if err != nil {
			t.Fatal(err)
		}
		return frame(b)
	}
	set := func(b []byte, i int, v byte) []byte { c := bytes.Clone(b); c[i] = v; return c }
	bigLen := bytes.Clone(good)
	binary.BigEndian.PutUint32(bigLen[8:], accesskey.MaxBodyLen+1)
	for _, tc := range []struct {
		name string
		in   []byte
		want error
	}{
		{"magic", set(good, 3, '2'), datekeys.ErrInvalidMagic},
		{"a .dkc", append([]byte("DKC1"), good[4:]...), datekeys.ErrInvalidMagic},
		{"empty", nil, datekeys.ErrInvalidMagic},
		{"framing version", set(good, 4, 2), datekeys.ErrUnsupportedVersion},
		{"flags", set(good, 5, 1), datekeys.ErrInvalidFlags},
		{"reserved", set(good, 7, 1), datekeys.ErrInvalidFlags},
		{"body length above the limit", bigLen, datekeys.ErrIntegrity},
		{"truncated prelude", good[:10], datekeys.ErrIntegrity},
		{"truncated body", good[:len(good)-1], datekeys.ErrIntegrity},
		{"trailing data", append(bytes.Clone(good), 0), datekeys.ErrIntegrity},
		{"schema version", with(func(m map[uint64]any) { m[1] = uint64(2) }), datekeys.ErrUnsupportedVersion},
		{"type tag", with(func(m map[uint64]any) { m[0] = "datekeycap" }), datekeys.ErrNonCanonicalCBOR},
		{"empty verification map", with(func(m map[uint64]any) { m[6] = map[uint64]any{} }), datekeys.ErrNonCanonicalCBOR},
		{"empty extension array", with(func(m map[uint64]any) { m[8] = []any{} }), datekeys.ErrNonCanonicalCBOR},
		{"null verification", with(func(m map[uint64]any) { m[6] = nil }), datekeys.ErrNonCanonicalCBOR},
		{"unknown key", with(func(m map[uint64]any) { m[9] = "x" }), datekeys.ErrNonCanonicalCBOR},
		{"short capsule_id", with(func(m map[uint64]any) { m[3] = make([]byte, 15) }), datekeys.ErrNonCanonicalCBOR},
		{"short digest", with(func(m map[uint64]any) { m[6] = map[uint64]any{0: make([]byte, 31)} }), datekeys.ErrNonCanonicalCBOR},
		{"unknown access type", with(func(m map[uint64]any) { m[4] = "mlkem768" }), datekeys.ErrAccessInvalid},
		{"short material", with(func(m map[uint64]any) { m[5] = make([]byte, 31) }), datekeys.ErrAccessInvalid},
		{"null material", with(func(m map[uint64]any) { m[5] = nil }), datekeys.ErrNonCanonicalCBOR},
		{"non-canonical body", frame(append([]byte{0xb9, 0x00, 0x07}, body[1:]...)), datekeys.ErrNonCanonicalCBOR},
	} {
		t.Run(tc.name, func(t *testing.T) {
			if _, err := accesskey.Decode(bytes.NewReader(tc.in)); !errors.Is(err, tc.want) {
				t.Fatalf("got %v, want %v", err, tc.want)
			}
		})
	}
}

func TestEncodeRejectsAbsenceAsEmptyMap(t *testing.T) {
	id, _ := age.GenerateX25519Identity()
	raw, _ := agewrap.RawX25519Identity(id)
	k := &accesskey.AccessKey{Type: accesskey.TypeX25519, Material: raw, Verification: &accesskey.Verification{}}
	if err := accesskey.Encode(io.Discard, k); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
		t.Fatalf("empty verification map encoded: %v", err)
	}
	k.Verification = nil
	k.Noncritical = []extension.Extension{{ID: "org.example.delivery", Version: 1}}
	var b bytes.Buffer
	if err := accesskey.Encode(&b, k); err != nil {
		t.Fatal(err)
	}
	back, err := accesskey.Decode(&b)
	if err != nil || back.Verification != nil || len(back.Noncritical) != 1 {
		t.Fatalf("%v %v", back, err)
	}
}

func TestIdentityWipeAndEncodeErrors(t *testing.T) {
	b, _ := loadDKK(t, "time_and_key_portable")
	k, _ := accesskey.Decode(bytes.NewReader(b))
	other := *k
	other.Type = "mlkem768"
	if _, err := other.Identity(); !errors.Is(err, datekeys.ErrAccessInvalid) {
		t.Fatalf("unsupported type: %v", err)
	}
	dup := []extension.Extension{{ID: "a", Version: 1}, {ID: "a", Version: 2}}
	for name, edit := range map[string]func(k *accesskey.AccessKey){
		"short material":       func(k *accesskey.AccessKey) { k.Material = k.Material[:31] },
		"repeated critical":    func(k *accesskey.AccessKey) { k.Critical = dup },
		"repeated noncritical": func(k *accesskey.AccessKey) { k.Noncritical = dup },
		"both arrays":          func(k *accesskey.AccessKey) { k.Critical, k.Noncritical = dup[:1], dup[1:] },
	} {
		c := *k
		c.Material = bytes.Clone(k.Material)
		edit(&c)
		if err := accesskey.Encode(io.Discard, &c); err == nil {
			t.Errorf("%s: encoded", name)
		}
	}
	if err := accesskey.Encode(failingWriter{}, k); err == nil {
		t.Fatal("write error ignored")
	}
	k.Wipe()
	if !bytes.Equal(k.Material, make([]byte, 32)) {
		t.Fatal("material not wiped")
	}
}

// retainingWriter keeps the slices it is given, which io.Writer forbids, so
// that a test can see what the writer's caller leaves in them.
type retainingWriter struct{ writes [][]byte }

func (w *retainingWriter) Write(p []byte) (int, error) {
	w.writes = append(w.writes, p)
	return len(p), nil
}

// retainingReader delivers data a few bytes at a time and keeps every slice
// it is given, up to its capacity, so that a test can see what the reader's
// caller leaves in them.
type retainingReader struct {
	data  []byte
	chunk int
	seen  [][]byte
}

func (r *retainingReader) Read(p []byte) (int, error) {
	r.seen = append(r.seen, p[:cap(p)])
	if len(r.data) == 0 {
		return 0, io.EOF
	}
	n := copy(p[:min(len(p), r.chunk)], r.data)
	r.data = r.data[n:]
	return n, nil
}

// Encode and Decode leave no copy of access_material in the buffers they
// allocate, including the ones Decode outgrows while it reads the body.
func TestEncodeAndDecodeLeaveNoStaleMaterial(t *testing.T) {
	material := bytes.Repeat([]byte{0xab}, 32)
	big, err := extension.New("org.example.big", 1, make([]byte, 4096))
	if err != nil {
		t.Fatal(err)
	}
	k := &accesskey.AccessKey{Type: accesskey.TypeX25519, Material: bytes.Clone(material), Noncritical: []extension.Extension{big}}
	w := &retainingWriter{}
	var dkk bytes.Buffer
	if err := accesskey.Encode(io.MultiWriter(&dkk, w), k); err != nil {
		t.Fatal(err)
	}
	for i, p := range w.writes {
		if bytes.Contains(p, material) {
			t.Errorf("Encode: slice %d written still holds the material", i)
		}
	}
	if !bytes.Equal(k.Material, material) {
		t.Fatal("Encode wiped the material of its argument")
	}
	if !bytes.Contains(dkk.Bytes(), material) {
		t.Fatal("the .dkk does not hold the material")
	}

	r := &retainingReader{data: dkk.Bytes(), chunk: 100}
	back, err := accesskey.Decode(r)
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(back.Material, material) {
		t.Fatal("material not decoded")
	}
	back.Wipe()
	for i, p := range r.seen {
		if bytes.Contains(p, material[:8]) {
			t.Errorf("Decode: read buffer %d of %d bytes still holds the material", i, len(p))
		}
	}

	// A truncated body is wiped too.
	r = &retainingReader{data: dkk.Bytes()[:dkk.Len()-10], chunk: 100}
	if _, err := accesskey.Decode(r); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("truncated body: %v", err)
	}
	for i, p := range r.seen {
		if bytes.Contains(p, material[:8]) {
			t.Errorf("Decode of a truncated body: read buffer %d still holds the material", i)
		}
	}
}

// A short .dkk that declares the largest BODY_LEN is rejected without an
// allocation of that size: the body buffer grows with the data read.
func TestDecodeShortBodyAllocatesLittle(t *testing.T) {
	var pre [accesskey.PreludeSize]byte
	copy(pre[:], accesskey.Magic)
	pre[4] = accesskey.FramingVersion
	binary.BigEndian.PutUint32(pre[8:12], accesskey.MaxBodyLen)
	r := &retainingReader{data: append(pre[:], make([]byte, 5000)...), chunk: 1000}
	if _, err := accesskey.Decode(r); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("short body: %v", err)
	}
	for _, p := range r.seen {
		if len(p) > 16<<10 {
			t.Fatalf("a read buffer of %d bytes for 5000 bytes of body", len(p))
		}
	}
}

// Spec §57: the DKK BODY limit binds the encoder and every decoder entry
// point, whatever the framing says, with the code of a frame violation.
func TestBodyLimit(t *testing.T) {
	if _, err := accesskey.DecodeBody(make([]byte, accesskey.MaxBodyLen+1)); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("DecodeBody above 16 MiB: %v", err)
	}
	b, _ := loadDKK(t, "time_and_key_portable")
	k, err := accesskey.Decode(bytes.NewReader(b))
	if err != nil {
		t.Fatal(err)
	}
	big, err := extension.New("org.example.big", 1, make([]byte, accesskey.MaxBodyLen))
	if err != nil {
		t.Fatal(err)
	}
	k.Noncritical = []extension.Extension{big}
	if err := accesskey.Encode(io.Discard, k); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("body above 16 MiB encoded: %v", err)
	}
}

func TestDecodeBodyExtensionRules(t *testing.T) {
	good, _ := loadDKK(t, "time_and_key_portable")
	m, err := cbortest.UnmarshalMap(good[accesskey.PreludeSize:])
	if err != nil {
		t.Fatal(err)
	}
	ext := func(id string, v uint64) map[uint64]any { return map[uint64]any{0: id, 1: v} }
	for name, edit := range map[string]func(m map[uint64]any){
		"critical out of order": func(m map[uint64]any) { m[7] = []any{ext("b", 1), ext("a", 1)} },
		"noncritical repeated":  func(m map[uint64]any) { m[8] = []any{ext("a", 1), ext("a", 2)} },
		"both arrays":           func(m map[uint64]any) { m[7] = []any{ext("a", 1)}; m[8] = []any{ext("a", 1)} },
		// Spec §54: data is absent or a non-empty byte string in its
		// shortest encoding; the extension map rejects anything else.
		"data length not in shortest form": func(m map[uint64]any) {
			m[8] = []any{map[uint64]any{0: "a", 1: uint64(1), 2: cbortest.Raw{0x58, 0x01, 0x00}}}
		},
		"data of type text": func(m map[uint64]any) { m[8] = []any{map[uint64]any{0: "a", 1: uint64(1), 2: "x"}} },
		"data of type unsigned": func(m map[uint64]any) {
			m[8] = []any{map[uint64]any{0: "a", 1: uint64(1), 2: cbortest.Raw{0x18, 0x01}}}
		},
		"empty data":   func(m map[uint64]any) { m[7] = []any{map[uint64]any{0: "a", 1: uint64(1), 2: []byte{}}} },
		"null data":    func(m map[uint64]any) { m[8] = []any{map[uint64]any{0: "a", 1: uint64(1), 2: nil}} },
		"version 2^32": func(m map[uint64]any) { m[8] = []any{ext("a", 1<<32)} },
		"65 extensions": func(m map[uint64]any) {
			var exts []any
			for i := range 65 {
				exts = append(exts, ext(fmt.Sprintf("x.%02d", i), 1))
			}
			m[8] = exts
		},
		"empty critical array":    func(m map[uint64]any) { m[7] = []any{} },
		"extension id not string": func(m map[uint64]any) { m[8] = []any{map[uint64]any{0: uint64(1), 1: uint64(1)}} },
	} {
		c := map[uint64]any{}
		for k, v := range m {
			c[k] = v
		}
		edit(c)
		b, err := cbortest.Marshal(c)
		if err != nil {
			t.Fatal(err)
		}
		if _, err := accesskey.DecodeBody(b); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
			t.Errorf("%s: %v", name, err)
		}
	}
}

type failingWriter struct{}

func (failingWriter) Write([]byte) (int, error) { return 0, errors.New("disk full") }

func FuzzDecode(f *testing.F) {
	for _, name := range fixtureNames {
		b, err := os.ReadFile(filepath.Join(fixtures, name+".dkk"))
		if err == nil {
			f.Add(b)
		}
	}
	f.Add([]byte("DKK1\x01\x00\x00\x00\x00\x00\x00\x01\xa0"))
	f.Fuzz(func(t *testing.T, in []byte) {
		k, err := accesskey.Decode(bytes.NewReader(in))
		if err != nil {
			if datekeys.Code(err) == "" {
				t.Fatalf("error without a normative code: %v", err)
			}
			return
		}
		var out bytes.Buffer
		if err := accesskey.Encode(&out, k); err != nil || !bytes.Equal(out.Bytes(), in) {
			t.Fatal("accepted a .dkk that does not re-encode to its input")
		}
	})
}

func testkitProfile() *profile.Profile { return profile.Quicknet() }

func mustDecrypt(t *testing.T, file []byte, id age.Identity) []byte {
	t.Helper()
	r, err := age.Decrypt(bytes.NewReader(file), id)
	if err != nil {
		t.Fatal(err)
	}
	b, err := io.ReadAll(r)
	if err != nil {
		t.Fatal(err)
	}
	return b
}