package capsule_test

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"testing"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/codec"
	"g.activething.com/go/DateKeys/datekey"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/internal/testkit"
	"g.activething.com/go/DateKeys/profile"
)

func TestPolicyNames(t *testing.T) {
	for _, p := range []capsule.Policy{capsule.TimeOnly, capsule.TimeAndKey} {
		got, err := capsule.ParsePolicy(p.String())
		if err != nil || got != p {
			t.Fatalf("%s: %v", p, err)
		}
	}
	if _, err := capsule.ParsePolicy("time_or_key"); err == nil {
		t.Fatal("unknown policy parsed")
	}
	if capsule.Policy(9).String() != "policy(9)" {
		t.Fatal("unknown policy name")
	}
}

func TestControlIsNotPrinted(t *testing.T) {
	c := capsule.Control{}
	for i := range c.PayloadIdentity {
		c.PayloadIdentity[i] = 0xab
	}
	for _, s := range []string{fmt.Sprint(c), fmt.Sprintf("%+v", &c), fmt.Sprintf("%#v", c)} {
		if strings.Contains(s, "abab") || !strings.Contains(s, "REDACTED") {
			t.Fatalf("I_PAYLOAD printed: %s", s)
		}
	}
}

func TestEncodeHeaderAndControlReject(t *testing.T) {
	dk := datekey.DateKey{ProfileID: profile.QuicknetID, Round: 1000}
	dup := []extension.Extension{{ID: "a", Version: 1}, {ID: "a", Version: 2}}
	for name, h := range map[string]capsule.Header{
		"invalid DateKey":      {},
		"unknown policy":       {DateKey: dk, Policy: 5},
		"repeated critical":    {DateKey: dk, Critical: dup},
		"repeated noncritical": {DateKey: dk, Noncritical: dup},
		"both arrays":          {DateKey: dk, Critical: dup[:1], Noncritical: dup[1:]},
	} {
		if _, err := capsule.EncodeHeader(&h); err == nil {
			t.Errorf("%s: accepted", name)
		}
	}
	for name, c := range map[string]capsule.Control{
		"repeated critical":    {Critical: dup},
		"repeated noncritical": {Noncritical: dup},
		"both arrays":          {Critical: dup[:1], Noncritical: dup[1:]},
	} {
		if _, err := capsule.EncodeControl(&c); err == nil {
			t.Errorf("control %s: accepted", name)
		}
	}
}

func marshal(t *testing.T, m map[uint64]any) []byte {
	t.Helper()
	b, err := codec.Marshal(m)
	if err != nil {
		t.Fatal(err)
	}
	return b
}

func ext(id string, v uint64) map[uint64]any { return map[uint64]any{0: id, 1: v} }

func extData(id string, data any) map[uint64]any { return map[uint64]any{0: id, 1: uint64(1), 2: data} }

// manyExts returns n distinct extensions in canonical order.
func manyExts(n int) []any {
	out := make([]any, n)
	for i := range out {
		out[i] = ext(fmt.Sprintf("org.example.%03d", i), 1)
	}
	return out
}

// Spec §54: extension rules shared by PUBLIC_HEADER and CONTROL_CBOR, as edits
// of the extension array at key.
func extensionRules(key uint64) map[string]func(m map[uint64]any) {
	return map[string]func(m map[uint64]any){
		"data of type text":         func(m map[uint64]any) { m[key] = []any{extData("a", "text")} },
		"data of type array":        func(m map[uint64]any) { m[key] = []any{extData("a", []any{uint64(1)})} },
		"data of type unsigned":     func(m map[uint64]any) { m[key] = []any{extData("a", uint64(7))} },
		"null data":                 func(m map[uint64]any) { m[key] = []any{extData("a", nil)} },
		"empty data h''":            func(m map[uint64]any) { m[key] = []any{extData("a", []byte{})} },
		"65 extensions":             func(m map[uint64]any) { m[key] = manyExts(65) },
		"extension_version 2^32":    func(m map[uint64]any) { m[key] = []any{ext("a", 1<<32)} },
		"extension_version 2^53-1":  func(m map[uint64]any) { m[key] = []any{ext("a", 1<<53-1)} },
		"extension map with key 3":  func(m map[uint64]any) { m[key] = []any{map[uint64]any{0: "a", 1: uint64(1), 3: []byte{1}}} },
		"extension without version": func(m map[uint64]any) { m[key] = []any{map[uint64]any{0: "a"}} },
	}
}

func TestDecodeHeaderRejects(t *testing.T) {
	dk := datekey.DateKey{ProfileID: profile.QuicknetID, Round: 1000}.Compact()
	base := func() map[uint64]any {
		return map[uint64]any{0: capsule.HeaderTypeTag, 1: uint64(1), 2: make([]byte, 16), 3: dk, 4: uint64(0)}
	}
	for _, tc := range []struct {
		name string
		edit func(m map[uint64]any)
		want error
	}{
		{"short capsule_id", func(m map[uint64]any) { m[2] = make([]byte, 15) }, datekeys.ErrNonCanonicalCBOR},
		{"invalid DateKey", func(m map[uint64]any) { m[3] = "dk1_x" }, datekeys.ErrDateKeyInvalid},
		{"type tag", func(m map[uint64]any) { m[0] = capsule.ControlTypeTag }, datekeys.ErrNonCanonicalCBOR},
		{"critical out of order", func(m map[uint64]any) { m[5] = []any{ext("b", 1), ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},
		{"noncritical repeated", func(m map[uint64]any) { m[6] = []any{ext("a", 1), ext("a", 2)} }, datekeys.ErrNonCanonicalCBOR},
		{"both arrays", func(m map[uint64]any) { m[5] = []any{ext("a", 1)}; m[6] = []any{ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},
	} {
		m := base()
		tc.edit(m)
		if _, err := capsule.DecodeHeader(marshal(t, m)); !errors.Is(err, tc.want) {
			t.Errorf("%s: got %v, want %v", tc.name, err, tc.want)
		}
	}
	for name, edit := range extensionRules(6) {
		m := base()
		edit(m)
		if _, err := capsule.DecodeHeader(marshal(t, m)); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
			t.Errorf("%s: got %v", name, err)
		}
	}
	valid := base()
	valid[6] = append(manyExts(63), extData("z", []byte{0}))
	if h, err := capsule.DecodeHeader(marshal(t, valid)); err != nil || len(h.Noncritical) != 64 {
		t.Fatalf("valid header with 64 extensions rejected: %v", err)
	}
}

func TestDecodeControlRejects(t *testing.T) {
	base := func() map[uint64]any {
		return map[uint64]any{0: capsule.ControlTypeTag, 1: uint64(1), 2: make([]byte, 32), 3: make([]byte, 32)}
	}
	for _, tc := range []struct {
		name string
		edit func(m map[uint64]any)
		want error
	}{
		{"schema version", func(m map[uint64]any) { m[1] = uint64(2) }, datekeys.ErrUnsupportedVersion},
		{"type tag", func(m map[uint64]any) { m[0] = capsule.HeaderTypeTag }, datekeys.ErrNonCanonicalCBOR},
		{"short binding", func(m map[uint64]any) { m[2] = make([]byte, 31) }, datekeys.ErrNonCanonicalCBOR},
		{"long identity", func(m map[uint64]any) { m[3] = make([]byte, 33) }, datekeys.ErrNonCanonicalCBOR},
		{"unknown key", func(m map[uint64]any) { m[6] = "x" }, datekeys.ErrNonCanonicalCBOR},
		{"critical repeated", func(m map[uint64]any) { m[4] = []any{ext("a", 1), ext("a", 2)} }, datekeys.ErrNonCanonicalCBOR},
		{"noncritical out of order", func(m map[uint64]any) { m[5] = []any{ext("b", 1), ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},
		{"both arrays", func(m map[uint64]any) { m[4] = []any{ext("a", 1)}; m[5] = []any{ext("a", 1)} }, datekeys.ErrNonCanonicalCBOR},
	} {
		m := base()
		tc.edit(m)
		if _, err := capsule.DecodeControl(marshal(t, m)); !errors.Is(err, tc.want) {
			t.Errorf("%s: got %v, want %v", tc.name, err, tc.want)
		}
	}
	for name, edit := range extensionRules(4) {
		m := base()
		edit(m)
		if _, err := capsule.DecodeControl(marshal(t, m)); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
			t.Errorf("%s: got %v", name, err)
		}
	}
	if _, err := capsule.DecodeControl([]byte("age-encryption.org/v1\n")); !errors.Is(err, datekeys.ErrNonCanonicalCBOR) {
		t.Fatalf("garbage control: %v", err)
	}
}

func TestHeaderLimit(t *testing.T) {
	big, err := extension.New("org.example.big", 1, bytes.Repeat([]byte{1}, capsule.MaxPublicHeaderLen))
	if err != nil {
		t.Fatal(err)
	}
	opts := past(t, 1000)
	opts.Noncritical = []extension.Extension{big}
	var dkc bytes.Buffer
	if _, err := capsule.Encrypt(&dkc, strings.NewReader("x"), opts); !errors.Is(err, datekeys.ErrIntegrity) || dkc.Len() != 0 {
		t.Fatalf("PUBLIC_HEADER above 1 MiB accepted: %v", err)
	}
	// Spec §57: the encoder refuses it with the code the decoder uses.
	if _, err := capsule.EncodeHeader(&capsule.Header{DateKey: datekey.DateKey{ProfileID: profile.QuicknetID, Round: 1000}, Noncritical: opts.Noncritical}); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("EncodeHeader above 1 MiB: %v", err)
	}
	// Spec §57: the decoder applies the limit too, whatever the framing says.
	if _, err := capsule.DecodeHeader(make([]byte, capsule.MaxPublicHeaderLen+1)); !errors.Is(err, datekeys.ErrIntegrity) {
		t.Fatalf("DecodeHeader above 1 MiB: %v", err)
	}
}

// A .dkk whose critical extension the application does not know is refused
// before any request.
func TestAccessKeyCriticalExtension(t *testing.T) {
	f := loadFixture(t, "time_and_key_portable")
	k := *f.dkk
	k.Critical = []extension.Extension{{ID: "org.example.must-understand", Version: 1}}
	o := f.openOptions(t)
	o.AccessKey = &k
	src := testkit.NewSource(f.release)
	o.Source = src
	if _, err := open(t, f.dkc, o); !errors.Is(err, datekeys.ErrExtensionCriticalUnknown) || src.Calls != 0 {
		t.Fatalf("got %v (requests %d)", err, src.Calls)
	}
	o.Extensions = extension.Set{"org.example.must-understand": {1}}
	if got, err := open(t, f.dkc, o); err != nil || !bytes.Equal(got, f.plaintext) {
		t.Fatalf("known .dkk extension rejected: %v", err)
	}
}

type brokenSeeker struct {
	*bytes.Reader
	seeks int
}

func (b *brokenSeeker) Seek(off int64, whence int) (int64, error) {
	b.seeks++
	if b.seeks > 1 {
		return 0, errors.New("seek failed")
	}
	return b.Reader.Seek(off, whence)
}

func TestCapsuleDigestSeekFailure(t *testing.T) {
	f := loadFixture(t, "time_and_key_portable")
	r := &brokenSeeker{Reader: bytes.NewReader(f.dkc)}
	if _, err := capsule.Open(context.Background(), io.Discard, r, f.openOptions(t)); err == nil {
		t.Fatal("seek failure ignored")
	}
}