package capsule_test

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"io"
	"strings"
	"testing"

	"filippo.io/age"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/capsule"
	"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"
	"g.activething.com/go/DateKeys/provider"
)

// mutation is one entry of the mutation corpus (spec §64): a function over a
// valid fixture that must fail with one exact normative error at one step.
type mutation struct {
	name string
	// spec is true for the twenty mutations listed in spec §64.
	spec bool
	make func(t *testing.T, env *env) (dkc []byte, opts capsule.OpenOptions)
	want *datekeys.Error
	step int
	// network reports whether the failure may happen after a release was
	// requested. Failures of steps 1 to 8 and of the access pre-checks must
	// not cause any request (spec §27, §63).
	network bool
}

type env struct {
	to, tk   *fixture // time_only and time_and_key_portable fixtures
	toParts  testkit.Parts
	tkParts  testkit.Parts
	sibling  []byte // another time_only capsule for the same round
	stranger *age.X25519Identity
}

func newEnv(t *testing.T) *env {
	e := &env{to: loadFixture(t, "time_only"), tk: loadFixture(t, "time_and_key_portable")}
	var err error
	if e.toParts, err = testkit.Split(e.to.dkc); err != nil {
		t.Fatal(err)
	}
	if e.tkParts, err = testkit.Split(e.tk.dkc); err != nil {
		t.Fatal(err)
	}
	var b bytes.Buffer
	p := profile.Quicknet()
	unlock, _ := datekey.RoundTime(p, 1000)
	if _, err := capsule.Encrypt(&b, strings.NewReader("sibling"), capsule.EncryptOptions{Profile: p, UnlockAt: unlock, Now: testkit.Fixed(testkit.Genesis())}); err != nil {
		t.Fatal(err)
	}
	e.sibling = b.Bytes()
	e.stranger, _ = age.GenerateX25519Identity()
	return e
}

func withSource(o capsule.OpenOptions, s provider.ReleaseSource) capsule.OpenOptions {
	o.Source = s
	return o
}

func set(b []byte, i int, v byte) []byte {
	c := bytes.Clone(b)
	c[i] = v
	return c
}

func xorLast(b []byte) []byte {
	c := bytes.Clone(b)
	c[len(c)-1] ^= 0x01
	return c
}

// build returns a capsule made by testkit.Build and the options to open it.
func build(t *testing.T, b testkit.Build) ([]byte, capsule.OpenOptions) {
	t.Helper()
	if b.Plaintext == nil {
		b.Plaintext = []byte("malicious creator")
	}
	out, err := b.Make()
	if err != nil {
		t.Fatal(err)
	}
	return out.DKC, capsule.OpenOptions{Registry: testkit.Registry(), Source: testkit.NewSource(testkit.Release(1000)), Now: testkit.Fixed(testkit.Genesis().AddDate(1, 0, 0))}
}

func headerWithDateKey(t *testing.T, e *env, dk string) []byte {
	t.Helper()
	h, err := capsule.DecodeHeader(e.toParts.Header)
	if err != nil {
		t.Fatal(err)
	}
	raw, err := testkit.RawHeader(h.CapsuleID, dk, 0)
	if err != nil {
		t.Fatal(err)
	}
	return testkit.Reframe(e.toParts.Prelude, raw, e.toParts.Sealed, e.toParts.Payload)
}

func policyByte(t *testing.T, header []byte) int {
	// The access_policy entry is the last one of a header without extensions: 0x04 <value>.
	i := len(header) - 2
	if header[i] != 0x04 {
		t.Fatalf("unexpected header layout %x", header[i:])
	}
	return i + 1
}

var mutations = []mutation{
	// ---- The twenty mutations of spec §64 -------------------------------
	{name: "PUBLIC_HEADER_A + SEALED_CONTROL_B", spec: true, want: datekeys.ErrHeaderBinding, step: 15, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			b, _ := testkit.Split(e.sibling)
			return testkit.Reframe(e.toParts.Prelude, e.toParts.Header, b.Sealed, b.Payload), e.to.openOptions(t)
		}},
	{name: "SEALED_CONTROL_A + PAYLOAD_AGE_B", spec: true, want: datekeys.ErrIntegrity, step: 17, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			b, _ := testkit.Split(e.sibling)
			return testkit.Join(e.toParts.Prelude, e.toParts.Header, e.toParts.Sealed, b.Payload), e.to.openOptions(t)
		}},
	{name: "DateKey A + release of round B", spec: true, want: datekeys.ErrRoundMismatch, step: 10, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			src := provider.ReleaseSourceFunc(func(context.Context, *profile.Profile, provider.Condition) (provider.Release, error) {
				return testkit.Release(1001), nil
			})
			return e.to.dkc, withSource(e.to.openOptions(t), src)
		}},
	{name: "chain hash changed", spec: true, want: datekeys.ErrProfileMismatch, step: 8,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			p := profile.Quicknet()
			other := strings.Repeat("ab", 32)
			return bytes.Replace(e.to.dkc, []byte(p.ChainHashHex()), []byte(other), 1), e.to.openOptions(t)
		}},
	{name: "version changed", spec: true, want: datekeys.ErrUnsupportedVersion, step: 2,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return set(e.to.dkc, 4, 2), e.to.openOptions(t)
		}},
	{name: "flags != 0", spec: true, want: datekeys.ErrInvalidFlags, step: 2,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return set(e.to.dkc, 5, 0x80), e.to.openOptions(t)
		}},
	{name: "reserved != 0", spec: true, want: datekeys.ErrInvalidFlags, step: 2,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return set(e.to.dkc, 7, 1), e.to.openOptions(t)
		}},
	{name: "payload truncated", spec: true, want: datekeys.ErrIntegrity, step: 17, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return e.to.dkc[:len(e.to.dkc)-1], e.to.openOptions(t)
		}},
	{name: "payload age modified", spec: true, want: datekeys.ErrIntegrity, step: 17, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return xorLast(e.to.dkc), e.to.openOptions(t)
		}},
	{name: "control modified", spec: true, want: datekeys.ErrIntegrity, step: 11, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return testkit.Join(e.toParts.Prelude, e.toParts.Header, xorLast(e.toParts.Sealed), e.toParts.Payload), e.to.openOptions(t)
		}},
	{name: "non-canonical dk1_ JSON", spec: true, want: datekeys.ErrDateKeyNonCanonical, step: 4,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			dk := datekey.Prefix + b64(`{"version":1, "network":"datekeys:quicknet:v1", "round":1000}`)
			return headerWithDateKey(t, e, dk), e.to.openOptions(t)
		}},
	{name: "unknown profile", spec: true, want: datekeys.ErrUnknownProfile, step: 4,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			dk := datekey.DateKey{ProfileID: "datekeys:evmnet:v1", Round: 1000}
			return headerWithDateKey(t, e, dk.Compact()), e.to.openOptions(t)
		}},
	{name: "release of another round", spec: true, want: datekeys.ErrReleaseInvalid, step: 10, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			forged := provider.Release{Round: 1000, Signature: testkit.Release(1001).Signature}
			return e.to.dkc, withSource(e.to.openOptions(t), testkit.NewSource(forged))
		}},
	{name: "access_policy=time_only with time_and_key structure", spec: true, want: datekeys.ErrPolicyStructureMismatch, step: 12, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			h := set(e.tkParts.Header, policyByte(t, e.tkParts.Header), 0)
			return testkit.Join(e.tkParts.Prelude, h, e.tkParts.Sealed, e.tkParts.Payload), e.tk.openOptions(t)
		}},
	{name: "access_policy=time_and_key with time_only structure", spec: true, want: datekeys.ErrPolicyStructureMismatch, step: 12, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			h := set(e.toParts.Header, policyByte(t, e.toParts.Header), 1)
			o := e.to.openOptions(t)
			o.Identities = []age.Identity{e.stranger}
			return testkit.Join(e.toParts.Prelude, h, e.toParts.Sealed, e.toParts.Payload), o
		}},
	{name: "extra stanza in OUTER_TIME_AGE", spec: true, want: datekeys.ErrPolicyStructureMismatch, step: 5,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return build(t, testkit.Build{EditOuter: func(fk []byte, s []*age.Stanza) []*age.Stanza {
				extra, _, _ := testkit.X25519Stanza(fk)
				return append(s, extra)
			}})
		}},
	{name: "extra stanza in PAYLOAD_AGE", spec: true, want: datekeys.ErrPolicyStructureMismatch, step: 6,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return build(t, testkit.Build{EditPayload: func(fk []byte, s []*age.Stanza) []*age.Stanza {
				extra, _, _ := testkit.X25519Stanza(fk)
				return append(s, extra)
			}})
		}},
	{name: "non-X25519 stanza in INNER_ACCESS_AGE", spec: true, want: datekeys.ErrPolicyStructureMismatch, step: 12, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			dkc, o := build(t, testkit.Build{Declared: capsule.TimeAndKey, Structure: capsule.TimeAndKey,
				AccessRecipients: []age.Recipient{e.stranger.Recipient()},
				EditInner: func(fk []byte, s []*age.Stanza) []*age.Stanza {
					return append(s, &age.Stanza{Type: "scrypt", Args: []string{"c2FsdHNhbHRzYWx0c2FsdA", "10"}, Body: make([]byte, 32)})
				}})
			o.Identities = []age.Identity{e.stranger}
			return dkc, o
		}},
	{name: "tlock stanza round differs from DateKey.round", spec: true, want: datekeys.ErrRoundMismatch, step: 8,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return build(t, testkit.Build{EditOuter: func(_ []byte, s []*age.Stanza) []*age.Stanza { s[0].Args[0] = "1001"; return s }})
		}},
	{name: "tlock stanza chain hash differs from the pinned profile", spec: true, want: datekeys.ErrProfileMismatch, step: 8,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return build(t, testkit.Build{EditOuter: func(_ []byte, s []*age.Stanza) []*age.Stanza {
				s[0].Args[1] = "dbd506d6ef76e5f386f41c651dcb808c5bcbd75471cc4eafa3f4df7ad4e4c493" // drand default chain
				return s
			}})
		}},

	// ---- Further cases ----------------------------------------------------
	{name: "magic", want: datekeys.ErrInvalidMagic, step: 1,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return set(e.to.dkc, 0, 'X'), e.to.openOptions(t)
		}},
	{name: "a .dkk offered as a .dkc", want: datekeys.ErrInvalidMagic, step: 1,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return append([]byte("DKK1"), e.to.dkc[4:]...), e.to.openOptions(t)
		}},
	{name: "empty file", want: datekeys.ErrInvalidMagic, step: 1,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) { return nil, e.to.openOptions(t) }},
	{name: "truncated prelude", want: datekeys.ErrIntegrity, step: 1,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) { return e.to.dkc[:10], e.to.openOptions(t) }},
	{name: "PUBLIC_HEADER_LEN above the limit", want: datekeys.ErrIntegrity, step: 2,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			c := bytes.Clone(e.to.dkc)
			binary.BigEndian.PutUint32(c[8:12], capsule.MaxPublicHeaderLen+1)
			return c, e.to.openOptions(t)
		}},
	{name: "SEALED_CONTROL_LEN above the limit", want: datekeys.ErrIntegrity, step: 2,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			c := bytes.Clone(e.to.dkc)
			binary.BigEndian.PutUint32(c[12:16], capsule.MaxSealedControlLen+1)
			return c, e.to.openOptions(t)
		}},
	{name: "truncated inside SEALED_CONTROL", want: datekeys.ErrIntegrity, step: 5,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return e.to.dkc[:len(e.toParts.Prelude)+len(e.toParts.Header)+10], e.to.openOptions(t)
		}},
	{name: "header schema version changed", want: datekeys.ErrUnsupportedVersion, step: 4,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			// a5 00 6a "datekeycap" 01 <version>
			return set(e.to.dkc, capsule.PreludeSize+14, 2), e.to.openOptions(t)
		}},
	{name: "unknown key in PUBLIC_HEADER", want: datekeys.ErrNonCanonicalCBOR, step: 4,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			h := append(bytes.Clone(e.toParts.Header), 0x07, 0x00)
			h[0]++ // one more map entry
			return testkit.Reframe(e.toParts.Prelude, h, e.toParts.Sealed, e.toParts.Payload), e.to.openOptions(t)
		}},
	{name: "undefined access_policy", want: datekeys.ErrNonCanonicalCBOR, step: 4,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return testkit.Join(e.toParts.Prelude, set(e.toParts.Header, policyByte(t, e.toParts.Header), 2), e.toParts.Sealed, e.toParts.Payload), e.to.openOptions(t)
		}},
	{name: "unknown critical PUBLIC_HEADER extension", want: datekeys.ErrExtensionCriticalUnknown, step: 4,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return build(t, testkit.Build{HeaderCritical: []extension.Extension{{ID: "org.example.must-understand", Version: 1}}})
		}},
	{name: "unknown critical CONTROL_CBOR extension", want: datekeys.ErrExtensionCriticalUnknown, step: 14, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return build(t, testkit.Build{ControlCritical: []extension.Extension{{ID: "org.example.must-understand", Version: 1}}})
		}},
	{name: "time_and_key without credentials", want: datekeys.ErrAccessRequired, step: 9,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			o := e.tk.openOptions(t)
			o.AccessKey = nil
			return e.tk.dkc, o
		}},
	{name: ".dkk of another capsule", want: datekeys.ErrAccessInvalid, step: 9,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			o := e.tk.openOptions(t)
			other := loadFixture(t, "time_and_key_recipients")
			o.AccessKey = other.dkk
			return e.tk.dkc, o
		}},
	{name: "capsule_digest of the .dkk does not match", want: datekeys.ErrAccessInvalid, step: 9,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return xorLast(e.tk.dkc), e.tk.openOptions(t)
		}},
	{name: "identity that is not a recipient", want: datekeys.ErrAccessInvalid, step: 13, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			o := e.tk.openOptions(t)
			o.AccessKey = nil
			o.Identities = []age.Identity{e.stranger}
			return e.tk.dkc, o
		}},
	{name: "round not reached yet", want: datekeys.ErrReleaseUnavailable, step: 9,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			o := e.to.openOptions(t)
			o.Now = testkit.Fixed(e.to.unlock(t).Add(-1))
			return e.to.dkc, o
		}},
	{name: "release source unavailable", want: datekeys.ErrReleaseUnavailable, step: 9, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return e.to.dkc, withSource(e.to.openOptions(t), testkit.NewSource())
		}},
	{name: "trailing data after PAYLOAD_AGE", want: datekeys.ErrIntegrity, step: 17, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return append(bytes.Clone(e.to.dkc), 0), e.to.openOptions(t)
		}},
	{name: "payload stanza body modified", want: datekeys.ErrIntegrity, step: 17, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			n, err := testkit.HeaderLen(e.toParts.Payload)
			if err != nil {
				t.Fatal(err)
			}
			// Flip a byte of the wrapped file key: the last body line before "---".
			i := bytes.LastIndex(e.toParts.Payload[:n], []byte("\n---")) - 10
			c := byte('A')
			if e.toParts.Payload[i] == 'A' {
				c = 'B'
			}
			p := set(e.toParts.Payload, i, c)
			return testkit.Join(e.toParts.Prelude, e.toParts.Header, e.toParts.Sealed, p), e.to.openOptions(t)
		}},
	{name: "tlock round edited by a third party", want: datekeys.ErrRoundMismatch, step: 8,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return bytes.Replace(e.to.dkc, []byte("-> tlock 1000 "), []byte("-> tlock 1001 "), 1), e.to.openOptions(t)
		}},
	{name: "empty registry", want: datekeys.ErrUnknownProfile, step: 4,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			o := e.to.openOptions(t)
			o.Registry, _ = profile.NewRegistry()
			return e.to.dkc, o
		}},
	{name: "time_only declared, time_and_key built by the creator", want: datekeys.ErrPolicyStructureMismatch, step: 12, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			return build(t, testkit.Build{Declared: capsule.TimeOnly, Structure: capsule.TimeAndKey, AccessRecipients: []age.Recipient{e.stranger.Recipient()}})
		}},
	{name: "time_and_key declared, time_only built by the creator", want: datekeys.ErrPolicyStructureMismatch, step: 12, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			dkc, o := build(t, testkit.Build{Declared: capsule.TimeAndKey, Structure: capsule.TimeOnly})
			o.Identities = []age.Identity{e.stranger}
			return dkc, o
		}},
	{name: "two INNER_ACCESS_AGE stanzas for one recipient", want: datekeys.ErrPolicyStructureMismatch, step: 13, network: true,
		make: func(t *testing.T, e *env) ([]byte, capsule.OpenOptions) {
			dkc, o := build(t, testkit.Build{Declared: capsule.TimeAndKey, Structure: capsule.TimeAndKey,
				AccessRecipients: []age.Recipient{e.stranger.Recipient()},
				EditInner: func(fk []byte, s []*age.Stanza) []*age.Stanza {
					again, _ := e.stranger.Recipient().Wrap(fk)
					return append(s, again[0])
				}})
			o.Identities = []age.Identity{e.stranger}
			return dkc, o
		}},
}

func TestMutationCorpus(t *testing.T) {
	e := newEnv(t)
	n := 0
	for _, m := range mutations {
		if m.spec {
			n++
		}
		t.Run(m.name, func(t *testing.T) {
			dkc, o := m.make(t, e)
			counter := &countingSource{inner: o.Source}
			o.Source = counter
			opened, err := capsule.Open(context.Background(), io.Discard, bytes.NewReader(dkc), o)
			if err == nil {
				t.Fatal("mutation accepted")
			}
			if !errors.Is(err, m.want) {
				t.Fatalf("got %v, want %v", err, m.want)
			}
			if !m.network && counter.calls != 0 {
				t.Fatalf("an invalid capsule caused %d release requests", counter.calls)
			}
			if m.step != 0 {
				checks := checksOf(opened, dkc, o)
				last := checks[len(checks)-1]
				if last.OK || last.Step != m.step || last.Error != m.want.Code() {
					t.Fatalf("failed at %+v, want step %d", last, m.step)
				}
			}
		})
	}
	if n != 20 {
		t.Fatalf("spec §64 lists 20 mutations, the corpus has %d", n)
	}
}

// checksOf returns the checks recorded for a failed Open; failures inside the
// inspection return no Opened, so the inspection is repeated for them.
func checksOf(opened *capsule.Opened, dkc []byte, o capsule.OpenOptions) []capsule.CheckResult {
	if opened != nil {
		return opened.Inspection.Checks
	}
	in, _ := capsule.Inspect(bytes.NewReader(dkc), capsule.InspectOptions{Registry: o.Registry, Extensions: o.Extensions})
	return in.Checks
}

type countingSource struct {
	inner provider.ReleaseSource
	calls int
}

func (c *countingSource) Fetch(ctx context.Context, p *profile.Profile, cond provider.Condition) (provider.Release, error) {
	c.calls++
	return c.inner.Fetch(ctx, p, cond)
}

// Known critical extensions are accepted when the application declares them.
func TestKnownCriticalExtensions(t *testing.T) {
	crit := []extension.Extension{{ID: "org.example.must-understand", Version: 1}}
	for _, b := range []testkit.Build{{HeaderCritical: crit}, {ControlCritical: crit}} {
		dkc, o := build(t, b)
		o.Extensions = extension.Set{"org.example.must-understand": {1}}
		var out bytes.Buffer
		if _, err := capsule.Open(context.Background(), &out, bytes.NewReader(dkc), o); err != nil || out.String() != "malicious creator" {
			t.Fatalf("known critical extension rejected: %v", err)
		}
	}
}

func b64(s string) string { return base64.RawURLEncoding.EncodeToString([]byte(s)) }