package capsule_test

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

	"filippo.io/age"

	datekeys "github.com/datekeys/datekeys-go"
	"github.com/datekeys/datekeys-go/accesskey"
	"github.com/datekeys/datekeys-go/capsule"
	"github.com/datekeys/datekeys-go/datekey"
	"github.com/datekeys/datekeys-go/extension"
	"github.com/datekeys/datekeys-go/internal/testkit"
	"github.com/datekeys/datekeys-go/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())}
}

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

func TestEncryptRejectsInvalidOptions(t *testing.T) {
	x, _ := age.GenerateX25519Identity()
	scrypt, _ := age.NewScryptRecipient("pw")
	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}}
		}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			opts := past(t, 1000)
			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 := capsule.Encrypt(&b, strings.NewReader("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)}
	var dkc bytes.Buffer
	res, err := capsule.Encrypt(&dkc, strings.NewReader("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, []any{"a", uint64(1)})
	cExt, _ := extension.New("org.example.sealed", 3, map[string]any{"k": []byte{1, 2}})
	opts := past(t, 1000)
	opts.Noncritical = []extension.Extension{hExt}
	opts.ControlNoncritical = []extension.Extension{cExt}
	var dkc bytes.Buffer
	if _, err := capsule.Encrypt(&dkc, strings.NewReader("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)
	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")
	}
}

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") }