package capsule_test

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"io"
	"runtime"
	"testing"

	"filippo.io/age"

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

// The tests of this file cover the reader of format 3 (spec v0.10): step 17,
// its substeps and their precedence, and step 18 (spec §29.2 to §29.7, §63).

// capsule3 describes a format 3 capsule of round 1000 to build, with edits
// at each level: the head, HEAD_CBOR, BODY, the plaintext with its padding,
// and PAYLOAD_AGE.
type capsule3 struct {
	paths           []string
	contents        [][]byte
	comment, author string
	security        []byte // nil: the empty security that writers write
	head            func(h *capsule.Head)
	headCBOR        func(b []byte) []byte
	body            func(b []byte) []byte
	plain           func(p []byte) []byte
	payload         func(t *testing.T, p []byte) []byte
	structure       capsule.Policy
	recipients      []age.Recipient
}

// build returns the .dkc and its BODY, before any edit of the plaintext.
func (c capsule3) build(t *testing.T) ([]byte, []byte) {
	t.Helper()
	h := testkit.Head3(c.comment, c.author, c.paths, c.contents)
	if c.head != nil {
		c.head(h)
	}
	hb, err := capsule.EncodeHead(h)
	if err != nil {
		t.Fatal(err)
	}
	if c.headCBOR != nil {
		hb = c.headCBOR(hb)
	}
	sec := c.security
	if sec == nil {
		sec = capsule.EncodeSecurity()
	}
	body := testkit.Body3(capsule.AreaLen, sec, hb, c.contents...)
	if c.body != nil {
		body = c.body(body)
	}
	b, err := testkit.Build{Format: capsule.Format3, Declared: c.structure, Structure: c.structure,
		AccessRecipients: c.recipients, Plaintext: body, EditPlaintext: c.plain}.Make()
	if err != nil {
		t.Fatal(err)
	}
	if c.payload == nil {
		return b.DKC, body
	}
	return testkit.Join(b.Prelude[:], b.PublicHeader, b.Sealed, c.payload(t, bytes.Clone(b.Payload))), body
}

type result3 struct {
	opened *capsule.Opened
	sink   *testkit.MemorySink
	err    error
	unread int // bytes of the .dkc that Open did not read
}

// open3 opens a capsule of round 1000 into sink.
func open3(t *testing.T, dkc []byte, sink capsule.Sink, ids ...age.Identity) result3 {
	t.Helper()
	o := defaultOpen(1000)
	o.Identities, o.Sink = ids, sink
	r := bytes.NewReader(dkc)
	opened, err := capsule.Open(context.Background(), nil, r, o)
	res := result3{opened: opened, err: err, unread: r.Len()}
	if m, ok := sink.(*testkit.MemorySink); ok {
		res.sink = m
	}
	return res
}

// chunk returns the offset in the age file p of its STREAM chunk i: after
// the header come a 16-byte nonce and chunks of 64 KiB, each with a 16-byte
// tag (C2SP age.md).
func chunk(t *testing.T, p []byte, i int) int {
	t.Helper()
	n, err := testkit.HeaderLen(p)
	if err != nil {
		t.Fatal(err)
	}
	return n + 16 + i*(64<<10+16)
}

// Spec §29.2 to §29.7, §63 steps 17 and 18: a format 3 capsule opens, its
// files reach the Sink in the order of the head, and Commit comes last.
func TestOpen3(t *testing.T) {
	photo := bytes.Repeat([]byte("playa"), 30000) // three STREAM chunks
	c := capsule3{
		paths:    []string{"fotos/playa.jpg", "nota.txt", "vacío.txt"},
		contents: [][]byte{photo, []byte("Hola.\n"), {}},
		comment:  "Para ti ❤\ufe0f", author: "Ana López",
	}
	dkc, body := c.build(t)

	// Steps 1 to 8 need no Sink.
	in, err := capsule.Inspect(bytes.NewReader(dkc), capsule.InspectOptions{Registry: testkit.Registry()})
	if err != nil || in.Prelude.Format != capsule.Format3 {
		t.Fatalf("Inspect: format %d, %v", in.Prelude.Format, err)
	}

	res := open3(t, dkc, &testkit.MemorySink{})
	if res.err != nil {
		t.Fatal(res.err)
	}
	o, s := res.opened, res.sink
	if !s.Committed || s.Aborted {
		t.Errorf("committed %v, aborted %v", s.Committed, s.Aborted)
	}
	for i, want := range c.contents {
		if !bytes.Equal(s.Files[i], want) {
			t.Errorf("file %d: %d bytes, want %d", i+1, len(s.Files[i]), len(want))
		}
	}
	p, _ := capsule.PaddedLength(uint64(len(body)), capsule.Reforzado)
	switch {
	case o.Format != capsule.Format3 || o.AreaLen != capsule.AreaLen:
		t.Errorf("format %d, area %d", o.Format, o.AreaLen)
	case o.PayloadLength != uint64(len(body)) || o.PaddedLength != p || o.Padding != capsule.Reforzado:
		t.Errorf("L = %d, P = %d, padding %s; want %d, %d", o.PayloadLength, o.PaddedLength, o.Padding, len(body), p)
	case o.Head == nil || o.Head != s.Head || o.Head.Comment != c.comment || o.Head.Author != c.author || len(o.Head.Files) != 3:
		t.Errorf("head %+v", o.Head)
	case o.Verdicts != capsule.Verdicts{Signature: capsule.VerdictNoSignature, Seal: capsule.VerdictNoSeal}:
		t.Errorf("verdicts %+v", o.Verdicts)
	}
	checks := o.Inspection.Checks
	if n := len(checks); n < 2 || checks[n-2].Step != 17 || !checks[n-2].OK || checks[n-1].Step != 18 || !checks[n-1].OK {
		t.Errorf("last checks %+v", checks[len(checks)-2:])
	}

	// The examples of spec §29.2: L and P of three small capsules.
	for _, tc := range []struct {
		name string
		c    capsule3
		l, p uint64
	}{
		{"no files and no comment", capsule3{}, 577, 768},
		{"a comment of one byte", capsule3{comment: "a"}, 580, 768},
		{"nota.txt of 1000 bytes, with mtime", capsule3{paths: []string{"nota.txt"}, contents: [][]byte{make([]byte, 1000)},
			head: func(h *capsule.Head) { h.Files[0].MTime, h.Files[0].HasMTime = 1790000000, true }}, 1641, 1792},
	} {
		dkc, _ := tc.c.build(t)
		res := open3(t, dkc, &testkit.MemorySink{})
		if res.err != nil || res.opened.PayloadLength != tc.l || res.opened.PaddedLength != tc.p || !res.sink.Committed {
			t.Errorf("%s: L = %d, P = %d, %v; want %d, %d", tc.name, res.opened.PayloadLength, res.opened.PaddedLength, res.err, tc.l, tc.p)
		}
	}

	// time_and_key: format 3 has the 16 stanzas of format 2 (spec §29.1, §39).
	id, _ := age.GenerateX25519Identity()
	dkc, _ = capsule3{paths: []string{"a"}, contents: [][]byte{[]byte("x")},
		structure: capsule.TimeAndKey, recipients: []age.Recipient{id.Recipient()}}.build(t)
	if res := open3(t, dkc, &testkit.MemorySink{}, id); res.err != nil || string(res.sink.Files[0]) != "x" {
		t.Errorf("time_and_key: %v", res.err)
	}
}

// Spec §63 step 17: a failure of age, or a plaintext whose length is not P,
// prevails; otherwise the first substep that fails decides. A code other
// than ERR_INTEGRITY is reported only after reading PAYLOAD_AGE to its end.
// Nothing reaches Commit, and a Sink that got Begin gets Abort.
func TestOpen3Substeps(t *testing.T) {
	two := capsule3{paths: []string{"a.txt", "b.txt"}, contents: [][]byte{[]byte("uno"), []byte("dos")}}
	// With a big file the head is in STREAM chunk 0, and two chunks follow.
	big := capsule3{paths: []string{"a.bin"}, contents: [][]byte{bytes.Repeat([]byte{0x5a}, 100000)}}
	with := func(c capsule3, edit func(c *capsule3)) capsule3 { edit(&c); return c }
	u32 := func(at int, v uint32) func(b []byte) []byte {
		return func(b []byte) []byte { binary.BigEndian.PutUint32(b[at:], v); return b }
	}
	dotdot := func(h *capsule.Head) { h.Files[0].Path = ".." }
	lastPadding := func(p []byte) []byte { p[len(p)-1] = 1; return p }
	short := func(p []byte) []byte { return p[:len(p)-1] }
	long := func(p []byte) []byte { return append(p, make([]byte, 256)...) }
	areaByte := func(b []byte) []byte { b[capsule.BodyFrameSize+capsule.AreaLen-1] = 1; return b }
	// HEAD_CBOR starts with the map, key 0, the text header and the 13 bytes
	// of "datekeys-head", then key 1 and the version.
	const headTag, headVersion = 3, 17

	for _, tc := range []struct {
		name  string
		c     capsule3
		want  error // nil: opens
		begun bool  // the Sink got Begin, and so Abort
		eof   bool  // Open read the .dkc to its end
	}{
		// 17.2: the frame and the area.
		{"AREA_LEN 511", with(two, func(c *capsule3) { c.body = u32(0, 511) }), datekeys.ErrIntegrity, false, false},
		{"AREA_LEN 513", with(two, func(c *capsule3) { c.body = u32(0, 513) }), datekeys.ErrIntegrity, false, false},
		{"AREA_LEN 66048", with(two, func(c *capsule3) { c.body = u32(0, 66048) }), datekeys.ErrIntegrity, false, false},
		{"SECURITY_LEN 0", with(two, func(c *capsule3) { c.body = u32(4, 0) }), datekeys.ErrIntegrity, false, false},
		{"SECURITY_LEN 513", with(two, func(c *capsule3) { c.body = u32(4, 513) }), datekeys.ErrIntegrity, false, false},
		{"HEAD_LEN 0", with(two, func(c *capsule3) { c.body = u32(8, 0) }), datekeys.ErrIntegrity, false, false},
		{"HEAD_LEN 2^24 + 1", with(two, func(c *capsule3) { c.body = u32(8, 1<<24+1) }), datekeys.ErrIntegrity, false, false},
		{"12 + AREA_LEN + HEAD_LEN = L + 1", with(capsule3{}, func(c *capsule3) { c.body = short }), datekeys.ErrIntegrity, false, false},
		{"L < 12", with(two, func(c *capsule3) { c.body = func(b []byte) []byte { return b[:11] } }), datekeys.ErrIntegrity, false, false},
		{"a byte of the area not zero", with(two, func(c *capsule3) { c.body = areaByte }), datekeys.ErrIntegrity, false, false},
		{"a byte of the area not zero, and path ..", with(two, func(c *capsule3) { c.head, c.body = dotdot, areaByte }), datekeys.ErrIntegrity, false, false},

		// 17.3 and 17.6: security never fails.
		{"security unreadable", with(two, func(c *capsule3) { c.security = []byte{0xa0} }), nil, false, true},

		// 17.4: the head, whose codes wait for the end of PAYLOAD_AGE.
		{"head of version 2", with(two, func(c *capsule3) {
			c.headCBOR = func(b []byte) []byte { b[headVersion] = 2; return b }
		}), datekeys.ErrUnsupportedVersion, false, true},
		{"head of another type tag", with(two, func(c *capsule3) {
			c.headCBOR = func(b []byte) []byte { b[headTag] = 'D'; return b }
		}), datekeys.ErrNonCanonicalCBOR, false, true},
		{"head with one byte more within HEAD_LEN", with(two, func(c *capsule3) {
			c.headCBOR = func(b []byte) []byte { return append(b, 0) }
		}), datekeys.ErrNonCanonicalCBOR, false, true},
		{"paths b and a, in that order", with(two, func(c *capsule3) { c.paths = []string{"b.txt", "a.txt"} }), datekeys.ErrNonCanonicalCBOR, false, true},
		{"path ..", with(two, func(c *capsule3) { c.head = dotdot }), datekeys.ErrHeadInvalid, false, true},
		{"paths A.txt and a.txt", with(two, func(c *capsule3) { c.paths = []string{"A.txt", "a.txt"} }), datekeys.ErrHeadInvalid, false, true},
		{"comment with U+202E", with(two, func(c *capsule3) { c.comment = "a\u202eb" }), datekeys.ErrHeadInvalid, false, true},
		{"start of an entry not the end of the one before", with(two, func(c *capsule3) {
			c.head = func(h *capsule.Head) { h.Files[1].Start, h.Files[1].End = 2, 5 }
		}), datekeys.ErrHeadInvalid, false, true},
		{"an unknown critical extension", with(two, func(c *capsule3) {
			c.head = func(h *capsule.Head) { h.Critical = []extension.Extension{{ID: "x.head", Version: 1}} }
		}), datekeys.ErrExtensionCriticalUnknown, false, true},

		// 17.5: the files fill CONTENT.
		{"end of the last file not C", with(two, func(c *capsule3) {
			c.body = func(b []byte) []byte { return append(b, 'x') }
		}), datekeys.ErrIntegrity, false, false},

		// 17.7: the SHA-256 of each file, and 17.8: the padding, up to P.
		{"a byte of a file changed", with(two, func(c *capsule3) {
			c.body = func(b []byte) []byte { b[len(b)-1] ^= 1; return b }
		}), datekeys.ErrIntegrity, true, false},
		{"a padding byte not zero", with(two, func(c *capsule3) { c.plain = lastPadding }), datekeys.ErrIntegrity, true, false},
		{"plaintext of P - 1 bytes", with(two, func(c *capsule3) { c.plain = short }), datekeys.ErrIntegrity, true, true},
		{"plaintext of P + 256 bytes", with(two, func(c *capsule3) { c.plain = long }), datekeys.ErrIntegrity, true, false},

		// Precedence: the head fails first, and what follows decides only when
		// it is a failure of age or of the length.
		{"path .. and a padding byte not zero", with(two, func(c *capsule3) { c.head, c.plain = dotdot, lastPadding }), datekeys.ErrHeadInvalid, false, true},
		{"path .. and a plaintext of P - 1 bytes", with(two, func(c *capsule3) { c.head, c.plain = dotdot, short }), datekeys.ErrIntegrity, false, true},
		{"path .. and a plaintext of P + 256 bytes", with(two, func(c *capsule3) { c.head, c.plain = dotdot, long }), datekeys.ErrIntegrity, false, true},
		{"path .. and the next STREAM chunk corrupt", with(big, func(c *capsule3) {
			c.head = dotdot
			c.payload = func(t *testing.T, p []byte) []byte { p[chunk(t, p, 1)+100] ^= 1; return p }
		}), datekeys.ErrIntegrity, false, false},
		{"path .. and a cut right after its chunk", with(big, func(c *capsule3) {
			c.head = dotdot
			c.payload = func(t *testing.T, p []byte) []byte { return p[:chunk(t, p, 1)] }
		}), datekeys.ErrIntegrity, false, true},
		{"path .. and a big file", with(big, func(c *capsule3) { c.head = dotdot }), datekeys.ErrHeadInvalid, false, true},
	} {
		dkc, _ := tc.c.build(t)
		res := open3(t, dkc, &testkit.MemorySink{})
		o, s := res.opened, res.sink
		if tc.want == nil {
			if res.err != nil || !s.Committed || o.Verdicts.Signature != capsule.VerdictUnreadable || len(o.Verdicts.Lines()) != 1 {
				t.Errorf("%s: %v, verdicts %+v", tc.name, res.err, o.Verdicts)
			}
			continue
		}
		step := failedStep(t, o.Inspection.Checks, res.err)
		expectStep(t, tc.name, step, res.err, tc.want, 17)
		switch {
		case s.Committed || o.Head != nil:
			t.Errorf("%s: committed %v, head %v", tc.name, s.Committed, o.Head)
		case (s.Head != nil) != tc.begun || s.Aborted != tc.begun:
			t.Errorf("%s: begun %v, aborted %v, want %v", tc.name, s.Head != nil, s.Aborted, tc.begun)
		case tc.eof && res.unread != 0:
			t.Errorf("%s: %d bytes of the .dkc not read", tc.name, res.unread)
		}
	}
}

// failSink fails at one point: "begin", "create", "write", "close" or
// "commit", at file 2 for "create", "write" and "close".
type failSink struct {
	testkit.MemorySink
	at     string
	aborts int
}

var errDiskFull = errors.New("disk full")

func (s *failSink) Begin(h *capsule.Head) error {
	if s.at == "begin" {
		return errDiskFull
	}
	return s.MemorySink.Begin(h)
}

func (s *failSink) Create(i int) (io.WriteCloser, error) {
	if s.at == "create" && i == 1 {
		return nil, errDiskFull
	}
	w, err := s.MemorySink.Create(i)
	return &failWriter{WriteCloser: w, fail: i == 1, at: s.at}, err
}

func (s *failSink) Commit() error {
	if s.at == "commit" {
		return errDiskFull
	}
	return s.MemorySink.Commit()
}

func (s *failSink) Abort() { s.aborts++; s.MemorySink.Abort() }

type failWriter struct {
	io.WriteCloser
	fail bool
	at   string
}

func (w *failWriter) Write(b []byte) (int, error) {
	if w.fail && w.at == "write" {
		return 0, errDiskFull
	}
	return w.WriteCloser.Write(b)
}

func (w *failWriter) Close() error {
	if w.fail && w.at == "close" {
		return errDiskFull
	}
	return w.WriteCloser.Close()
}

// A failure of the Sink is the caller's own error with ERR_INTEGRITY, as a
// failure of dst is in formats 1 and 2, and after Begin it gets Abort, once.
func TestOpen3SinkFailures(t *testing.T) {
	dkc, _ := capsule3{paths: []string{"a.txt", "b.txt"}, contents: [][]byte{[]byte("uno"), []byte("dos")}}.build(t)
	for _, at := range []string{"begin", "create", "write", "close", "commit"} {
		s := &failSink{at: at}
		res := open3(t, dkc, s)
		wantAborts := 1
		if at == "begin" {
			wantAborts = 0
		}
		switch {
		case !errors.Is(res.err, errDiskFull) || datekeys.Code(res.err) != datekeys.Code(datekeys.ErrIntegrity):
			t.Errorf("%s: %v", at, res.err)
		case s.aborts != wantAborts || s.Committed || res.opened.Head != nil:
			t.Errorf("%s: %d aborts, committed %v", at, s.aborts, s.Committed)
		}
	}
	if res := open3(t, dkc, &failSink{}); res.err != nil {
		t.Fatal(res.err)
	}
}

// Spec §57: a reader reserves no memory by what BODY declares before it
// receives the bytes. Each capsule declares 16 MiB, in HEAD_LEN or in the
// size of a file, and holds a plaintext of 2000 bytes; Open fails, and
// allocates about what a small capsule makes it allocate.
func TestOpen3DeclaredLengths(t *testing.T) {
	const declared = 16 << 20
	cut := func(p []byte) []byte { return p[:2000] }
	for _, tc := range []struct {
		name string
		c    capsule3
	}{
		{"HEAD_LEN", capsule3{plain: cut, body: func(b []byte) []byte {
			binary.BigEndian.PutUint32(b[8:], declared)
			return append(b, make([]byte, declared)...)
		}}},
		{"size", capsule3{plain: cut, paths: []string{"a.bin"}, contents: [][]byte{make([]byte, declared)}}},
	} {
		dkc, _ := tc.c.build(t)
		var before, after runtime.MemStats
		runtime.GC()
		runtime.ReadMemStats(&before)
		res := open3(t, dkc, &testkit.MemorySink{})
		runtime.ReadMemStats(&after)
		if n := after.TotalAlloc - before.TotalAlloc; n > 4<<20 || !errors.Is(res.err, datekeys.ErrIntegrity) {
			t.Errorf("%s: %d bytes allocated, %v", tc.name, n, res.err)
		}
	}
}