package capsule_test

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

	"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/capsule"
	"g.activething.com/go/DateKeys/codec/bech32"
	"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 past(t *testing.T, round uint64) capsule.EncryptOptions {
	t.Helper()
	p := profile.Quicknet()
	unlock, err := datekey.RoundTime(p, round)
	if err != nil {
		t.Fatal(err)
	}
	return capsule.EncryptOptions{Profile: p, UnlockAt: unlock, Now: testkit.Fixed(testkit.Genesis()), TestVectors: true}
}

// encrypt runs capsule.Encrypt with content as its source, of the length it
// declares.
func encrypt(t *testing.T, dst io.Writer, content string, opts capsule.EncryptOptions) (*capsule.Result, error) {
	t.Helper()
	opts.Length = int64(len(content))
	return capsule.Encrypt(dst, strings.NewReader(content), opts)
}

// recipientFromRaw returns the age X25519 recipient of 32 raw bytes, as age
// parses it: without checking that they are a canonical public key.
func recipientFromRaw(t *testing.T, raw []byte) *age.X25519Recipient {
	t.Helper()
	s, err := bech32.Encode("age", raw)
	if err != nil {
		t.Fatal(err)
	}
	r, err := age.ParseX25519Recipient(s)
	if err != nil {
		t.Fatal(err)
	}
	return r
}

func open(t *testing.T, dkc []byte, o capsule.OpenOptions) ([]byte, error) {
	t.Helper()
	var out bytes.Buffer
	_, err := capsule.Open(context.Background(), &out, bytes.NewReader(dkc), o)
	return out.Bytes(), err
}

func defaultOpen(round uint64) capsule.OpenOptions {
	return capsule.OpenOptions{Registry: testkit.Registry(), Source: testkit.NewSource(testkit.Release(round)), Now: time.Now}
}

func TestEncryptRoundTripBothPolicies(t *testing.T) {
	msg := strings.Repeat("0123456789abcdef", 20000) // several STREAM chunks
	for _, tc := range []struct {
		name  string
		setup func(o *capsule.EncryptOptions) []age.Identity
	}{
		{"time_only", func(o *capsule.EncryptOptions) []age.Identity { return nil }},
		{"time_and_key, portable", func(o *capsule.EncryptOptions) []age.Identity {
			o.Policy, o.NewPortableKey = capsule.TimeAndKey, true
			return nil
		}},
		{"time_and_key, three recipients", func(o *capsule.EncryptOptions) []age.Identity {
			o.Policy = capsule.TimeAndKey
			var ids []age.Identity
			for range 3 {
				id, _ := age.GenerateX25519Identity()
				o.Recipients = append(o.Recipients, id.Recipient())
				ids = append(ids, id)
			}
			return ids
		}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			opts := past(t, 1000)
			ids := tc.setup(&opts)
			var dkc bytes.Buffer
			res, err := encrypt(t, &dkc, msg, opts)
			if err != nil {
				t.Fatal(err)
			}
			if res.DateKey.Round != 1000 || !res.UnlockAt.Equal(opts.UnlockAt) {
				t.Fatalf("result %+v", res)
			}
			o := defaultOpen(1000)
			o.Identities = ids
			if res.PortableKey != nil {
				o.AccessKey = res.PortableKey
			}
			got, err := open(t, dkc.Bytes(), o)
			if err != nil || string(got) != msg {
				t.Fatalf("open: %v", err)
			}
			for i, id := range ids {
				o := defaultOpen(1000)
				o.Identities = []age.Identity{id}
				if got, err := open(t, dkc.Bytes(), o); err != nil || string(got) != msg {
					t.Fatalf("recipient %d: %v", i, err)
				}
			}
		})
	}
}

// Spec §37, §39, §62.1: the writer rejects options that would make a capsule
// nobody opens, anybody opens, or one that reveals more than it should,
// before it writes anything.
func TestEncryptRejectsInvalidOptions(t *testing.T) {
	x, _ := age.GenerateX25519Identity()
	scrypt, _ := age.NewScryptRecipient("pw")
	seventeen := make([]age.Recipient, 17)
	for i := range seventeen {
		id, _ := age.GenerateX25519Identity()
		seventeen[i] = id.Recipient()
	}
	// p = 2^255 - 19 in little-endian, and a public key with bit 255 set.
	p := append([]byte{0xed}, bytes.Repeat([]byte{0xff}, 30)...)
	p = append(p, 0x7f)
	bit255, _ := agewrap.RawX25519Recipient(x.Recipient())
	bit255[31] |= 0x80
	for _, tc := range []struct {
		name string
		edit func(o *capsule.EncryptOptions)
	}{
		{"no profile", func(o *capsule.EncryptOptions) { o.Profile = nil }},
		{"no clock", func(o *capsule.EncryptOptions) { o.Now = nil }},
		{"unlock time in the past", func(o *capsule.EncryptOptions) { o.Now = testkit.Fixed(o.UnlockAt.Add(time.Second)) }},
		{"unlock time equal to now", func(o *capsule.EncryptOptions) { o.Now = testkit.Fixed(o.UnlockAt) }},
		{"time_only with recipients", func(o *capsule.EncryptOptions) { o.Recipients = []age.Recipient{x.Recipient()} }},
		{"time_only with a portable key", func(o *capsule.EncryptOptions) { o.NewPortableKey = true }},
		{"time_and_key without recipients", func(o *capsule.EncryptOptions) { o.Policy = capsule.TimeAndKey }},
		{"non-X25519 recipient", func(o *capsule.EncryptOptions) {
			o.Policy, o.Recipients = capsule.TimeAndKey, []age.Recipient{scrypt}
		}},
		{"recipient listed twice", func(o *capsule.EncryptOptions) {
			o.Policy, o.Recipients = capsule.TimeAndKey, []age.Recipient{x.Recipient(), x.Recipient()}
		}},
		{"unknown policy", func(o *capsule.EncryptOptions) { o.Policy = 7 }},
		{"invalid profile", func(o *capsule.EncryptOptions) { o.Profile.ChainHash[0] ^= 1 }},
		{"duplicate header extension", func(o *capsule.EncryptOptions) {
			o.Noncritical = []extension.Extension{{ID: "a", Version: 1}, {ID: "a", Version: 2}}
		}},
		{"extension both critical and noncritical", func(o *capsule.EncryptOptions) {
			o.ControlCritical = []extension.Extension{{ID: "a", Version: 1}}
			o.ControlNoncritical = []extension.Extension{{ID: "a", Version: 1}}
		}},
		// Format 2 (spec §29.1, §39, §62.1).
		{"negative length", func(o *capsule.EncryptOptions) { o.Length = -1 }},
		{"length above L_MAX", func(o *capsule.EncryptOptions) { o.Length = capsule.MaxPayloadLength + 1 }},
		{"padding code 3", func(o *capsule.EncryptOptions) { o.Padding = 3 }},
		{"17 recipients", func(o *capsule.EncryptOptions) { o.Policy, o.Recipients = capsule.TimeAndKey, seventeen }},
		{"16 recipients and a portable key", func(o *capsule.EncryptOptions) {
			o.Policy, o.Recipients, o.NewPortableKey = capsule.TimeAndKey, seventeen[:16], true
		}},
		{"recipient with bit 255 set", func(o *capsule.EncryptOptions) {
			o.Policy, o.Recipients = capsule.TimeAndKey, []age.Recipient{recipientFromRaw(t, bit255)}
		}},
		{"recipient u = p", func(o *capsule.EncryptOptions) {
			o.Policy, o.Recipients = capsule.TimeAndKey, []age.Recipient{recipientFromRaw(t, p)}
		}},
		{"recipient of low order, u = 0", func(o *capsule.EncryptOptions) {
			o.Policy, o.Recipients = capsule.TimeAndKey, []age.Recipient{recipientFromRaw(t, make([]byte, 32))}
		}},
		{"recipient of low order, u = 1", func(o *capsule.EncryptOptions) {
			o.Policy, o.Recipients = capsule.TimeAndKey, []age.Recipient{recipientFromRaw(t, append([]byte{1}, make([]byte, 31)...))}
		}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			opts := past(t, 1000)
			opts.Length = 1
			tc.edit(&opts)
			var dkc bytes.Buffer
			if _, err := capsule.Encrypt(&dkc, strings.NewReader("x"), opts); err == nil {
				t.Fatal("accepted")
			}
			if dkc.Len() != 0 {
				t.Fatal("wrote output before validating the options")
			}
		})
	}
}

// Spec §38: an I_ACCESS is generated for one capsule only and never reused.
func TestPortableKeysAreNeverReused(t *testing.T) {
	var dkcs [2][]byte
	var keys [2]*accesskey.AccessKey
	for i := range 2 {
		opts := past(t, 1000)
		opts.Policy, opts.NewPortableKey = capsule.TimeAndKey, true
		var b bytes.Buffer
		res, err := encrypt(t, &b, "x", opts)
		if err != nil {
			t.Fatal(err)
		}
		dkcs[i], keys[i] = b.Bytes(), res.PortableKey
	}
	if bytes.Equal(keys[0].Material, keys[1].Material) || keys[0].CredentialID == keys[1].CredentialID || keys[0].CapsuleID == keys[1].CapsuleID {
		t.Fatal("two capsules share an I_ACCESS, credential_id or capsule_id")
	}
	// The .dkk of capsule A is refused for capsule B before any request, and
	// its identity cannot open B's access layer either.
	o := defaultOpen(1000)
	o.AccessKey = keys[0]
	src := testkit.NewSource(testkit.Release(1000))
	o.Source = src
	if _, err := open(t, dkcs[1], o); !errors.Is(err, datekeys.ErrAccessInvalid) || src.Calls != 0 {
		t.Fatalf("foreign .dkk: %v (requests: %d)", err, src.Calls)
	}
	id, _ := keys[0].Identity()
	o = defaultOpen(1000)
	o.Identities = []age.Identity{id}
	if _, err := open(t, dkcs[1], o); !errors.Is(err, datekeys.ErrAccessInvalid) {
		t.Fatalf("foreign identity: %v", err)
	}
}

func TestFutureCapsuleStaysLockedWithoutRequests(t *testing.T) {
	p := profile.Quicknet()
	now := time.Date(2026, 9, 25, 12, 0, 0, 0, time.UTC)
	opts := capsule.EncryptOptions{Profile: p, UnlockAt: now.Add(time.Hour), Now: testkit.Fixed(now), TestVectors: true}
	var dkc bytes.Buffer
	res, err := encrypt(t, &dkc, "secret", opts)
	if err != nil {
		t.Fatal(err)
	}
	if res.UnlockAt.Before(opts.UnlockAt) || res.UnlockAt.Sub(opts.UnlockAt) >= p.Period {
		t.Fatalf("unsafe rounding: %s for %s", res.UnlockAt, opts.UnlockAt)
	}
	src := testkit.NewSource()
	o := capsule.OpenOptions{Registry: testkit.Registry(), Source: src, Now: testkit.Fixed(now)}
	if _, err := open(t, dkc.Bytes(), o); !errors.Is(err, datekeys.ErrReleaseUnavailable) || src.Calls != 0 {
		t.Fatalf("locked capsule: %v (requests: %d)", err, src.Calls)
	}
	// Inspection works on a locked capsule and reports its condition.
	in, err := capsule.Inspect(bytes.NewReader(dkc.Bytes()), capsule.InspectOptions{Registry: testkit.Registry()})
	if err != nil || in.Header.DateKey != res.DateKey || !in.UnlockAt.Equal(res.UnlockAt) {
		t.Fatalf("inspect: %+v %v", in, err)
	}
}

func TestExtensionsRoundTrip(t *testing.T) {
	hExt, _ := extension.New("org.example.public", 1, []byte("public"))
	cExt, _ := extension.New("org.example.sealed", 3, []byte{0xa1, 0x00, 0x42, 0x01, 0x02})
	opts := past(t, 1000)
	opts.Noncritical = []extension.Extension{hExt}
	opts.ControlNoncritical = []extension.Extension{cExt}
	var dkc bytes.Buffer
	if _, err := encrypt(t, &dkc, "x", opts); err != nil {
		t.Fatal(err)
	}
	in, _ := capsule.Inspect(bytes.NewReader(dkc.Bytes()), capsule.InspectOptions{Registry: testkit.Registry()})
	if len(in.Header.Noncritical) != 1 || !bytes.Equal(in.Header.Noncritical[0].Data, hExt.Data) {
		t.Fatal("header extension lost")
	}
	var out bytes.Buffer
	opened, err := capsule.Open(context.Background(), &out, bytes.NewReader(dkc.Bytes()), defaultOpen(1000))
	if err != nil || len(opened.ControlNoncritical) != 1 || !bytes.Equal(opened.ControlNoncritical[0].Data, cExt.Data) {
		t.Fatalf("control extension lost: %v", err)
	}
}

func TestOpenRequiresOptions(t *testing.T) {
	f := loadFixture(t, "time_only")
	for name, o := range map[string]capsule.OpenOptions{
		"no source":   {Registry: testkit.Registry(), Now: time.Now},
		"no clock":    {Registry: testkit.Registry(), Source: testkit.NewSource()},
		"no registry": {Source: testkit.NewSource(), Now: time.Now},
	} {
		if _, err := capsule.Open(context.Background(), io.Discard, bytes.NewReader(f.dkc), o); err == nil {
			t.Errorf("%s: accepted", name)
		}
	}
}

func TestEncryptWriteError(t *testing.T) {
	opts := past(t, 1000)
	opts.Length = 1
	if _, err := capsule.Encrypt(failingWriter{}, strings.NewReader("x"), opts); err == nil {
		t.Fatal("write error ignored")
	}
	if _, err := capsule.Encrypt(io.Discard, failingReader{}, opts); err == nil {
		t.Fatal("read error ignored")
	}
}

// Spec §39, §62.1 rule 3: from 1 to 16 credentials; 16 of them leave no
// dummy, and each opens the capsule on its own.
func TestCredentialBounds(t *testing.T) {
	opts := past(t, 1000)
	opts.Policy, opts.NewPortableKey = capsule.TimeAndKey, true
	var ids []age.Identity
	for range 15 {
		id, _ := age.GenerateX25519Identity()
		opts.Recipients = append(opts.Recipients, id.Recipient())
		ids = append(ids, id)
	}
	var dkc bytes.Buffer
	res, err := encrypt(t, &dkc, "sixteen", opts)
	if err != nil {
		t.Fatal(err)
	}
	k, _ := res.PortableKey.Identity()
	for i, id := range append(ids, k) {
		o := defaultOpen(1000)
		o.Identities = []age.Identity{id}
		if got, err := open(t, dkc.Bytes(), o); err != nil || string(got) != "sixteen" {
			t.Fatalf("credential %d: %v", i, err)
		}
	}
}

// Spec §62.1 rules 6 and 9: L is known before the control is sealed, and a
// source that delivers another number of bytes is an error, not a capsule
// that would fail at step 17 after the date.
func TestEncryptSourceLength(t *testing.T) {
	for _, tc := range []struct {
		name    string
		content string
		length  int64
		ok      bool
	}{
		{"exact", "content", 7, true},
		{"empty", "", 0, true},
		{"one byte short", "content", 8, false},
		{"one byte more", "content", 6, false},
		{"length not set", "content", 0, false},
	} {
		opts := past(t, 1000)
		opts.Length = tc.length
		var dkc bytes.Buffer
		res, err := capsule.Encrypt(&dkc, strings.NewReader(tc.content), opts)
		if (err == nil) != tc.ok {
			t.Fatalf("%s: %v", tc.name, err)
		}
		if !tc.ok {
			continue
		}
		if res.Format != capsule.Format2 || res.Length != uint64(tc.length) || res.Padding != capsule.Reforzado || res.PaddedLength != 256 {
			t.Fatalf("%s: result %+v", tc.name, res)
		}
		if got, err := open(t, dkc.Bytes(), defaultOpen(1000)); err != nil || string(got) != tc.content {
			t.Fatalf("%s: open: %q, %v", tc.name, got, err)
		}
	}
}

type failingWriter struct{}

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

type failingReader struct{}

func (failingReader) Read([]byte) (int, error) { return 0, errors.New("read error") }