package capsule_test

import (
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"io"
	"os"
	"path/filepath"
	"reflect"
	"slices"
	"testing"
	"time"

	"filippo.io/age"

	"g.activething.com/go/DateKeys/accesskey"
	"g.activething.com/go/DateKeys/agewrap"
	"g.activething.com/go/DateKeys/capsule"
	"g.activething.com/go/DateKeys/extension"
	"g.activething.com/go/DateKeys/internal/testkit"
	"g.activething.com/go/DateKeys/profile"
	"g.activething.com/go/DateKeys/provider"
)

const fixtureDir = "../testdata/fixtures"

// fixtureNames are the official fixtures: the format 1 ones of spec v0.8.2,
// kept for compatibility, and the format 2 ones (spec §67).
var fixtureNames = []string{
	"time_only", "time_only_extensions", "time_and_key_portable", "time_and_key_recipients", "empty_payload",
	"format2_time_only", "format2_time_only_bloque256", "format2_empty_payload", "format2_time_only_extensions",
	"format2_time_and_key_portable", "format2_time_and_key_recipients", "format2_time_and_key_sixteen",
	"format3_single", "format3_tree", "format3_comment_only", "format3_bloque256", "format3_time_and_key_portable",
	"format3_area_1024", "format3_security_v2", "format3_signature_unsupported", "format3_seal_unsupported",
}

type fixture struct {
	testkit.DKCFixture
	dkc       []byte
	plaintext []byte
	release   provider.Release
	dkk       *accesskey.AccessKey
	ids       []age.Identity
}

func loadFixture(t testing.TB, name string) *fixture {
	t.Helper()
	f := &fixture{}
	if err := testkit.ReadJSON(filepath.Join(fixtureDir, name+".json"), &f.DKCFixture); err != nil {
		t.Fatal(err)
	}
	var err error
	if f.dkc, err = os.ReadFile(filepath.Join(fixtureDir, f.File)); err != nil {
		t.Fatal(err)
	}
	if f.plaintext, err = os.ReadFile(filepath.Join(fixtureDir, f.PlaintextFile)); err != nil {
		t.Fatal(err)
	}
	sig, err := hex.DecodeString(f.Release.Signature)
	if err != nil {
		t.Fatal(err)
	}
	f.release = provider.Release{Round: f.Release.Round, Signature: sig}
	if f.AccessKeyFile != "" {
		b, err := os.ReadFile(filepath.Join(fixtureDir, f.AccessKeyFile))
		if err != nil {
			t.Fatal(err)
		}
		if f.dkk, err = accesskey.Decode(bytes.NewReader(b)); err != nil {
			t.Fatal(err)
		}
	}
	for _, s := range f.Identities {
		id, err := age.ParseX25519Identity(s)
		if err != nil {
			t.Fatal(err)
		}
		f.ids = append(f.ids, id)
	}
	return f
}

func (f *fixture) format() capsule.Format { return capsule.Format(f.Format) }

// slots is the number of INNER_ACCESS_AGE stanzas of the format of f, 0
// meaning one or more.
func (f *fixture) slots() int {
	if f.format() != capsule.Format1 {
		return agewrap.AccessSlots
	}
	return 0
}

func (f *fixture) unlock(t testing.TB) time.Time {
	u, err := time.Parse(time.RFC3339, f.UnlockAt)
	if err != nil {
		t.Fatal(err)
	}
	return u
}

// openOptions returns the options that open f: its release, its unlock time
// and, for time_and_key, its .dkk or else its identities.
func (f *fixture) openOptions(t testing.TB) capsule.OpenOptions {
	o := capsule.OpenOptions{Registry: testkit.Registry(), Source: testkit.NewSource(f.release), Now: testkit.Fixed(f.unlock(t)), Sink: testkit.DiscardSink{}}
	if f.AccessPolicy == "time_and_key" {
		if f.dkk != nil {
			o.AccessKey = f.dkk
		} else {
			o.Identities = f.ids
		}
	}
	return o
}

// credentials returns the identities that open INNER_ACCESS_AGE: the .dkk
// first, when there is one, then the identities of the record.
func (f *fixture) credentials(t testing.TB) []age.Identity {
	var ids []age.Identity
	if f.dkk != nil {
		id, err := f.dkk.Identity()
		if err != nil {
			t.Fatal(err)
		}
		ids = append(ids, id)
	}
	return append(ids, f.ids...)
}

func decryptAge(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
}

func stanzaList(in []capsule.StanzaInfo) []testkit.FixtureStanza {
	out := make([]testkit.FixtureStanza, len(in))
	for i, s := range in {
		out[i] = testkit.FixtureStanza{Type: s.Type, Args: s.Args}
	}
	return out
}

// Spec §67: conformance is shown by decrypting and verifying each official
// fixture and comparing every intermediate value and the plaintext.
func TestConformanceFixtures(t *testing.T) {
	p := profile.Quicknet()
	for _, name := range fixtureNames {
		t.Run(name, func(t *testing.T) {
			f := loadFixture(t, name)
			if sum := sha256.Sum256(f.dkc); hex.EncodeToString(sum[:]) != f.SHA256 {
				t.Fatal("fixture bytes changed")
			}
			if sum := sha256.Sum256(f.plaintext); hex.EncodeToString(sum[:]) != f.PlaintextSHA256 {
				t.Fatal("plaintext file changed")
			}
			if err := provider.Verify(p, provider.Condition{Round: f.release.Round}, f.release); err != nil {
				t.Fatalf("embedded release: %v", err)
			}

			// PRELUDE, PUBLIC_HEADER and the pre-unlock view (steps 1 to 8).
			parts, err := testkit.Split(f.dkc)
			if err != nil {
				t.Fatal(err)
			}
			if hex.EncodeToString(parts.Prelude) != f.Prelude || hex.EncodeToString(parts.Header) != f.PublicHeader {
				t.Fatal("PRELUDE or PUBLIC_HEADER differ")
			}
			in, err := capsule.Inspect(bytes.NewReader(f.dkc), capsule.InspectOptions{Registry: testkit.Registry()})
			if err != nil {
				t.Fatal(err)
			}
			if in.Prelude.Format != f.format() || in.Header.DateKey.Compact() != f.DateKey || in.Header.CapsuleIDHex() != f.CapsuleID ||
				in.Header.Policy.String() != f.AccessPolicy || in.UnlockAt.Format(time.RFC3339) != f.UnlockAt {
				t.Fatalf("inspection differs: format %d, %+v", in.Prelude.Format, in.Header)
			}
			if !reflect.DeepEqual(stanzaList(in.OuterStanzas), f.OuterStanzas) || !reflect.DeepEqual(stanzaList(in.PayloadStanzas), f.PayloadStanzas) {
				t.Fatal("stanzas differ")
			}
			for _, c := range in.Checks {
				if !c.OK || c.Step > 8 {
					t.Fatalf("unexpected inspection check %+v", c)
				}
			}
			var pre [capsule.PreludeSize]byte
			copy(pre[:], parts.Prelude)
			if b := capsule.HeaderBinding(pre, parts.Header); hex.EncodeToString(b[:]) != f.HeaderBinding {
				t.Fatal("header_binding differs")
			}
			h, err := capsule.DecodeHeader(parts.Header)
			if err != nil {
				t.Fatal(err)
			}
			if re, _ := capsule.EncodeHeader(h); !bytes.Equal(re, parts.Header) {
				t.Fatal("EncodeHeader(DecodeHeader(x)) != x")
			}
			if got := fixtureExts(h.Critical, h.Noncritical); !reflect.DeepEqual(got, f.HeaderExtensions) {
				t.Fatalf("header extensions differ:\n got %+v\nwant %+v", got, f.HeaderExtensions)
			}

			// Opening layer by layer: OUTER_TIME_AGE, INNER_ACCESS_AGE, CONTROL_CBOR.
			timeID, _ := agewrap.NewTimeIdentity(p, f.release.Round, f.release)
			inner := decryptAge(t, parts.Sealed, timeID)
			control := inner
			if f.Structure == "time_and_key" {
				st, err := agewrap.Stanzas(bytes.NewReader(inner))
				if err != nil {
					t.Fatal(err)
				}
				got := make([]testkit.FixtureStanza, len(st))
				for i, s := range st {
					got[i] = testkit.FixtureStanza{Type: s.Type, Args: s.Args}
				}
				if !reflect.DeepEqual(got, f.InnerStanzas) {
					t.Fatal("INNER_ACCESS_AGE stanzas differ")
				}
				if err := agewrap.CheckAccessStanzas(st, f.slots()); err != nil {
					t.Fatal(err)
				}
				if f.format() != capsule.Format1 {
					checkCredentialStanzas(t, f, st)
				}
				acc, _ := agewrap.NewAccessIdentity(f.slots(), f.credentials(t)[0])
				control = decryptAge(t, inner, acc)
			}
			if hex.EncodeToString(control) != f.ControlCBOR {
				t.Fatal("CONTROL_CBOR differs")
			}
			ctrl, err := capsule.DecodeControl(control, f.format())
			if err != nil {
				t.Fatal(err)
			}
			if hex.EncodeToString(ctrl.PayloadIdentity[:]) != f.PayloadIdentity || hex.EncodeToString(ctrl.HeaderBinding[:]) != f.HeaderBinding {
				t.Fatal("I_PAYLOAD or header_binding in CONTROL_CBOR differ")
			}
			if re, _ := capsule.EncodeControl(ctrl, f.format()); !bytes.Equal(re, control) {
				t.Fatal("EncodeControl(DecodeControl(x)) != x")
			}
			if got := fixtureExts(ctrl.Critical, ctrl.Noncritical); !reflect.DeepEqual(got, f.ControlExt) {
				t.Fatalf("control extensions differ:\n got %+v\nwant %+v", got, f.ControlExt)
			}

			// PAYLOAD_AGE: the content, BODY in format 3, and in formats 2 and 3
			// its padding (spec §29.1, §29.2).
			payloadID, _ := agewrap.NewPayloadIdentity(ctrl.PayloadIdentity[:])
			plaintext := decryptAge(t, parts.Payload, payloadID)
			if f.PayloadLength != uint64(len(f.plaintext)) {
				t.Fatalf("payload_length %d, the content is %d bytes", f.PayloadLength, len(f.plaintext))
			}
			switch f.format() {
			case capsule.Format1:
				if !bytes.Equal(plaintext, f.plaintext) || ctrl.PayloadLength != 0 || ctrl.Padding != 0 || f.Padding != 0 || f.PaddedLength != 0 {
					t.Fatal("PAYLOAD_AGE plaintext differs")
				}
			case capsule.Format2, capsule.Format3:
				padded, err := capsule.PaddedLength(f.PayloadLength, capsule.Padding(f.Padding))
				if err != nil {
					t.Fatal(err)
				}
				if ctrl.PayloadLength != f.PayloadLength || ctrl.Padding != capsule.Padding(f.Padding) || padded != f.PaddedLength ||
					uint64(len(plaintext)) != padded || !bytes.Equal(plaintext[:f.PayloadLength], f.plaintext) ||
					len(bytes.Trim(plaintext[f.PayloadLength:], "\x00")) != 0 {
					t.Fatal("PAYLOAD_AGE plaintext is not the content and its padding")
				}
				if n := uint64(len(parts.Payload)); n != capsule.PayloadAgeLength(padded) {
					t.Fatalf("PAYLOAD_AGE of %d bytes, P = %d gives %d", n, padded, capsule.PayloadAgeLength(padded))
				}
			}
			if f.format() == capsule.Format3 {
				checkBody3(t, f)
			}

			// The complete flow through the public API, stage by stage: the
			// content, or in format 3 the files, the head and the verdicts.
			var out bytes.Buffer
			o := f.openOptions(t)
			sink := &testkit.MemorySink{}
			o.Sink = sink
			opened, err := capsule.Open(context.Background(), &out, bytes.NewReader(f.dkc), o)
			if err != nil {
				t.Fatal(err)
			}
			if f.format() == capsule.Format3 {
				if out.Len() != 0 || !sink.Committed || len(sink.Files) != len(f.Files) {
					t.Fatal("Open did not deliver the files to the Sink")
				}
				for i, r := range f.Files {
					if !bytes.Equal(sink.Files[i], f.plaintext[f.ContentOffset+r.Start:f.ContentOffset+r.End]) {
						t.Fatalf("file %q differs", r.Path)
					}
				}
				v := opened.Verdicts
				if string(v.Signature) != f.Verdicts.Signature || string(v.Seal) != f.Verdicts.Seal || opened.AreaLen != f.AreaLen ||
					opened.Head.Comment != f.Comment || opened.Head.Author != f.Author || hex.EncodeToString(opened.Head.Salt[:]) != f.Salt {
					t.Fatalf("Open reports another head or verdicts: %+v", v)
				}
			} else if !bytes.Equal(out.Bytes(), f.plaintext) {
				t.Fatal("plaintext differs")
			}
			if opened.Format != f.format() || opened.PayloadLength != f.PayloadLength ||
				opened.Padding != capsule.Padding(f.Padding) || opened.PaddedLength != f.PaddedLength {
				t.Fatalf("Open reports format %d, L = %d, padding %d, P = %d", opened.Format, opened.PayloadLength, opened.Padding, opened.PaddedLength)
			}
			var stages []testkit.FixtureStage
			for _, c := range opened.Inspection.Checks {
				stages = append(stages, testkit.FixtureStage{Step: c.Step, Name: c.Name, OK: c.OK, Error: c.Error})
			}
			if !reflect.DeepEqual(stages, f.Stages) {
				t.Fatalf("stages differ:\n got %+v\nwant %+v", stages, f.Stages)
			}
			if got := fixtureExts(opened.ControlCritical, opened.ControlNoncritical); !reflect.DeepEqual(got, f.ControlExt) {
				t.Fatal("Open does not report the control extensions")
			}
		})
	}
}

// checkBody3 checks BODY, the plaintext file of a format 3 fixture, against
// its record: the frame, the area, security and its verdicts, the head and
// each file with its SHA-256 (spec §29.2 to §29.7).
func checkBody3(t *testing.T, f *fixture) {
	t.Helper()
	body := f.plaintext
	l := uint64(len(body))
	frame, err := capsule.ParseBodyFrame(body[:capsule.BodyFrameSize], l)
	if err != nil {
		t.Fatal(err)
	}
	area := body[capsule.BodyFrameSize : capsule.BodyFrameSize+uint64(frame.AreaLen)]
	if err := capsule.CheckArea(area, frame.SecurityLen); err != nil {
		t.Fatal(err)
	}
	security := area[:frame.SecurityLen]
	offset := capsule.BodyFrameSize + uint64(frame.AreaLen) + uint64(frame.HeadLen)
	hb := body[capsule.BodyFrameSize+uint64(frame.AreaLen) : offset]
	if frame.AreaLen != f.AreaLen || hex.EncodeToString(security) != f.Security || hex.EncodeToString(hb) != f.Head || offset != f.ContentOffset {
		t.Fatal("the area, security, the head or the offset of CONTENT differ")
	}
	v := capsule.EvaluateSecurity(security)
	if f.Verdicts == nil || string(v.Signature) != f.Verdicts.Signature || string(v.Seal) != f.Verdicts.Seal || !reflect.DeepEqual(v.Lines(), f.Verdicts.Lines) {
		t.Fatalf("verdicts %+v, recorded %+v", v, f.Verdicts)
	}
	h, err := capsule.DecodeHead(hb, nil)
	if err != nil {
		t.Fatal(err)
	}
	if re, _ := capsule.EncodeHead(h); !bytes.Equal(re, hb) {
		t.Fatal("EncodeHead(DecodeHead(x)) != x")
	}
	if err := capsule.CheckHeadEnd(h, frame.ContentLength(l)); err != nil {
		t.Fatal(err)
	}
	if hex.EncodeToString(h.Salt[:]) != f.Salt || h.Comment != f.Comment || h.Author != f.Author || len(h.Files) != len(f.Files) {
		t.Fatal("the head differs from the record")
	}
	if got := fixtureExts(h.Critical, h.Noncritical); !reflect.DeepEqual(got, f.HeadExtensions) {
		t.Fatalf("head extensions differ:\n got %+v\nwant %+v", got, f.HeadExtensions)
	}
	for i, e := range h.Files {
		r := f.Files[i]
		recorded := r.MTime != nil
		if e.Path != r.Path || e.Size != r.Size || e.Start != r.Start || e.End != r.End || hex.EncodeToString(e.SHA256[:]) != r.SHA256 ||
			e.HasMTime != recorded || recorded && e.MTime != *r.MTime {
			t.Fatalf("file %d: %+v, recorded %+v", i+1, e, r)
		}
		if sha256.Sum256(body[offset+e.Start:offset+e.End]) != e.SHA256 {
			t.Fatalf("file %q: its bytes do not have its SHA-256", e.Path)
		}
	}
}

// checkCredentialStanzas checks the record of a time_and_key fixture of
// format 2 or 3: each credential opens exactly the stanza the record gives, and no
// credential opens any other one, which are the dummies (spec §39, §67).
func checkCredentialStanzas(t *testing.T, f *fixture, st []*age.Stanza) {
	t.Helper()
	want := slices.Clone(f.IdentityStanzas)
	if f.dkk != nil {
		if f.AccessKeyStanza == nil {
			t.Fatal("no access_key_stanza for the .dkk")
		}
		want = append([]int{*f.AccessKeyStanza}, want...)
	}
	creds := f.credentials(t)
	if len(creds) != len(want) {
		t.Fatalf("%d credentials, %d recorded stanzas", len(creds), len(want))
	}
	opened := map[int]bool{}
	for c, id := range creds {
		for i, s := range st {
			fk, err := id.Unwrap([]*age.Stanza{s})
			if errors.Is(err, age.ErrIncorrectIdentity) {
				continue
			}
			if err != nil || i != want[c] || opened[i] {
				t.Fatalf("credential %d opens stanza %d, recorded %d: %v", c, i, want[c], err)
			}
			opened[i] = true
			clear(fk)
		}
	}
	if len(opened) != len(creds) {
		t.Fatalf("%d credentials open %d stanzas", len(creds), len(opened))
	}
}

// fixtureExts lists extensions as the fixtures record them, with the exact
// data bytes (spec §67, §68).
func fixtureExts(critical, noncritical []extension.Extension) []testkit.FixtureExt {
	var out []testkit.FixtureExt
	for i, e := range append(slices.Clone(critical), noncritical...) {
		out = append(out, testkit.FixtureExt{Critical: i < len(critical), ID: e.ID, Version: e.Version, Data: hex.EncodeToString(e.Data)})
	}
	return out
}

// The extension data of time_only_extensions and format2_time_only_extensions,
// restated independently of their JSON: raw UTF-8 bytes in PUBLIC_HEADER, and
// CBOR in the profile of spec §58 in CONTROL_CBOR. The base protocol decodes
// neither.
func TestExtensionFixtureData(t *testing.T) {
	for _, name := range []string{"time_only_extensions", "format2_time_only_extensions"} {
		f := loadFixture(t, name)
		var out bytes.Buffer
		opened, err := capsule.Open(context.Background(), &out, bytes.NewReader(f.dkc), f.openOptions(t))
		if err != nil {
			t.Fatal(err)
		}
		h := opened.Inspection.Header
		if len(h.Critical) != 0 || len(h.Noncritical) != 1 || h.Noncritical[0].ID != "org.example.label" || h.Noncritical[0].Version != 1 ||
			string(h.Noncritical[0].Data) != "public label" {
			t.Fatalf("%s: header extension %+v", name, h.Noncritical)
		}
		c := opened.ControlNoncritical
		if len(opened.ControlCritical) != 0 || len(c) != 1 || c[0].ID != "org.example.note" || c[0].Version != 2 ||
			hex.EncodeToString(c[0].Data) != "a2000701667365616c6564" {
			t.Fatalf("%s: control extension %+v", name, c)
		}
	}
}

func loadAccessKey(t *testing.T, name string) *accesskey.AccessKey {
	t.Helper()
	b, err := os.ReadFile(filepath.Join(fixtureDir, name+".dkk"))
	if err != nil {
		t.Fatal(err)
	}
	k, err := accesskey.Decode(bytes.NewReader(b))
	if err != nil {
		t.Fatal(err)
	}
	return k
}

// Every known recipient of a multi-recipient fixture opens it on its own.
func TestFixtureRecipients(t *testing.T) {
	for _, c := range []struct {
		name string
		n    int
	}{{"time_and_key_recipients", 2}, {"format2_time_and_key_recipients", 3}, {"format2_time_and_key_sixteen", 16}} {
		f := loadFixture(t, c.name)
		if len(f.ids) != c.n {
			t.Fatalf("%s should have %d known recipients", c.name, c.n)
		}
		for i, id := range f.ids {
			o := f.openOptions(t)
			o.AccessKey = nil
			o.Identities = []age.Identity{id}
			var out bytes.Buffer
			if _, err := capsule.Open(context.Background(), &out, bytes.NewReader(f.dkc), o); err != nil || !bytes.Equal(out.Bytes(), f.plaintext) {
				t.Fatalf("%s, recipient %d: %v", c.name, i, err)
			}
		}
	}
}

// A non-seekable reader works too: only the capsule_digest shortcut is skipped.
func TestOpenFromPlainReader(t *testing.T) {
	for _, name := range []string{"time_only", "time_and_key_portable", "format2_time_only", "format2_time_and_key_portable"} {
		f := loadFixture(t, name)
		var out bytes.Buffer
		r := struct{ io.Reader }{bytes.NewReader(f.dkc)}
		if _, err := capsule.Open(context.Background(), &out, r, f.openOptions(t)); err != nil || !bytes.Equal(out.Bytes(), f.plaintext) {
			t.Fatalf("%s: %v", name, err)
		}
	}
}