package capsule

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

	"filippo.io/age"

	datekeys "g.activething.com/go/DateKeys"
	"g.activething.com/go/DateKeys/agewrap"
)

// Spec §39, §62 step 6, §76 change 3: the 16 recipients of INNER_ACCESS_AGE,
// the credentials and the dummies, are in a uniformly random order. The slot
// of the first element, and of the last, follows the uniform distribution
// over the 16: the chi-square statistic of 32 000 permutations, with 15
// degrees of freedom, stays below 60, which a uniform distribution passes
// with probability above 1 - 10^-6 and a slot fixed or much favoured does not.
// fillSlots adds 15 distinct dummies to one credential, and none to 16.
func TestStanzaOrderIsUniform(t *testing.T) {
	const draws = 32000
	var first, last [agewrap.AccessSlots]int
	for range draws {
		s := make([]int, agewrap.AccessSlots)
		for i := range s {
			s[i] = i
		}
		if err := permute(s); err != nil {
			t.Fatal(err)
		}
		for i, v := range s {
			switch v {
			case 0:
				first[i]++
			case agewrap.AccessSlots - 1:
				last[i]++
			}
		}
	}
	for _, count := range [][agewrap.AccessSlots]int{first, last} {
		expected := float64(draws) / agewrap.AccessSlots
		chi2 := 0.0
		for _, c := range count {
			d := float64(c) - expected
			chi2 += d * d / expected
		}
		if chi2 > 60 {
			t.Fatalf("slot counts %v, chi-square %.1f", count, chi2)
		}
	}

	credential, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatal(err)
	}
	slots, err := fillSlots([]age.Recipient{credential.Recipient()})
	if err != nil {
		t.Fatal(err)
	}
	distinct := map[string]bool{}
	for _, r := range slots {
		distinct[r.(*age.X25519Recipient).String()] = true
	}
	if len(slots) != agewrap.AccessSlots || len(distinct) != agewrap.AccessSlots || !distinct[credential.Recipient().String()] {
		t.Fatalf("%d slots, %d distinct", len(slots), len(distinct))
	}

	var sixteen []age.Recipient
	seen := map[string]bool{}
	for range agewrap.AccessSlots {
		id, _ := age.GenerateX25519Identity()
		sixteen = append(sixteen, id.Recipient())
		seen[id.Recipient().String()] = true
	}
	slots, err = fillSlots(sixteen)
	if err != nil {
		t.Fatal(err)
	}
	for _, r := range slots {
		if !seen[r.(*age.X25519Recipient).String()] {
			t.Fatal("a dummy among 16 credentials")
		}
		delete(seen, r.(*age.X25519Recipient).String())
	}
}

// Spec §62.1 rule 11, §76 change 13: the self-checks of the writer reject
// what the reader would reject after the date.
func TestEncryptSelfCheck(t *testing.T) {
	portable, _ := age.GenerateX25519Identity()
	other, _ := age.GenerateX25519Identity()
	control := []byte("control")
	slots, err := fillSlots([]age.Recipient{portable.Recipient()})
	if err != nil {
		t.Fatal(err)
	}
	good, err := encryptAll(control, slots...)
	if err != nil {
		t.Fatal(err)
	}
	if err := selfCheckInner(good, control, portable); err != nil {
		t.Fatalf("a valid INNER_ACCESS_AGE: %v", err)
	}
	// The portable key and 14 others: with one more, 15 stanzas; with the
	// portable key again, 16 stanzas of which it opens two.
	var others []age.Recipient
	for range 14 {
		id, _ := age.GenerateX25519Identity()
		others = append(others, id.Recipient())
	}
	fifteen, _ := encryptAll(control, append([]age.Recipient{portable.Recipient()}, others...)...)
	twice, _ := encryptAll(control, append([]age.Recipient{portable.Recipient(), portable.Recipient()}, others...)...)
	for name, tc := range map[string]struct {
		inner   []byte
		control []byte
		id      *age.X25519Identity
	}{
		"15 stanzas":                          {fifteen, control, portable},
		"two stanzas for the portable key":    {twice, control, portable},
		"a portable key that is no recipient": {good, control, other},
		"another control":                     {good, []byte("other"), portable},
	} {
		if err := selfCheckInner(tc.inner, tc.control, tc.id); err == nil {
			t.Errorf("INNER_ACCESS_AGE with %s passed the self-check", name)
		}
	}

	// PAYLOAD_AGE: the length P gives, and a header I_PAYLOAD opens.
	payloadID, _ := age.GenerateX25519Identity()
	raw, _ := agewrap.RawX25519Identity(payloadID)
	write := func(plaintext []byte, to age.Recipient) *payloadWriter {
		p := &payloadWriter{w: io.Discard}
		aw, err := age.Encrypt(p, to)
		if err != nil {
			t.Fatal(err)
		}
		if _, err := aw.Write(plaintext); err != nil {
			t.Fatal(err)
		}
		if err := aw.Close(); err != nil {
			t.Fatal(err)
		}
		return p
	}
	if err := selfCheckPayload(write(make([]byte, 256), payloadID.Recipient()), raw, 256); err != nil {
		t.Fatalf("a valid PAYLOAD_AGE: %v", err)
	}
	if err := selfCheckPayload(write(make([]byte, 255), payloadID.Recipient()), raw, 256); err == nil {
		t.Error("a PAYLOAD_AGE one byte short passed the self-check")
	}
	if err := selfCheckPayload(write(make([]byte, 256), other.Recipient()), raw, 256); err == nil {
		t.Error("a PAYLOAD_AGE for another recipient passed the self-check")
	}
}

// checkPadding reads the padding of a format 2 plaintext: exactly p - l zero
// bytes, found in any read pattern.
func TestCheckPadding(t *testing.T) {
	for _, tc := range []struct {
		name string
		rest string
		ok   bool
	}{
		{"exact zeros", strings.Repeat("\x00", 200), true},
		{"one byte short", strings.Repeat("\x00", 199), false},
		{"one byte more", strings.Repeat("\x00", 201), false},
		{"first byte not zero", "\x01" + strings.Repeat("\x00", 199), false},
		{"last byte not zero", strings.Repeat("\x00", 199) + "\x80", false},
	} {
		for _, r := range []io.Reader{strings.NewReader(tc.rest), &oneByteReader{strings.NewReader(tc.rest)}} {
			err := checkPadding(r, 56, 256)
			if (err == nil) != tc.ok || (err != nil && !errors.Is(err, datekeys.ErrIntegrity)) {
				t.Errorf("%s: %v", tc.name, err)
			}
		}
	}
	if err := checkPadding(&failingAfter{n: 10}, 56, 256); err == nil || datekeys.Code(err) != "" {
		t.Errorf("a read error is not a padding error: %v", err)
	}
}

type oneByteReader struct{ r io.Reader }

func (o *oneByteReader) Read(p []byte) (int, error) { return o.r.Read(p[:min(1, len(p))]) }

type failingAfter struct{ n int }

func (f *failingAfter) Read(p []byte) (int, error) {
	if f.n == 0 {
		return 0, errors.New("stream failure")
	}
	k := min(f.n, len(p))
	clear(p[:k])
	f.n -= k
	return k, nil
}

// decryptAll reads the plaintext into one buffer of the ciphertext's size, so
// that no outgrown buffer keeps a copy of I_PAYLOAD, and still reports a
// truncated STREAM, which the age reader signals with io.ErrUnexpectedEOF.
func TestDecryptAll(t *testing.T) {
	id, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatal(err)
	}
	const chunk = 64 << 10
	for _, size := range []int{0, 1, 511, 512, 513, chunk, chunk + 1, 2*chunk + 100} {
		plaintext := bytes.Repeat([]byte{0xab}, size)
		ct, err := encryptAll(plaintext, id.Recipient())
		if err != nil {
			t.Fatal(err)
		}
		out, err := decryptAll(ct, id)
		if err != nil || !bytes.Equal(out, plaintext) {
			t.Fatalf("%d bytes: %v", size, err)
		}
		if cap(out) != len(ct) {
			t.Errorf("%d bytes: a buffer of %d bytes for a ciphertext of %d", size, cap(out), len(ct))
		}
		if size == 0 {
			continue
		}
		// The last chunk removed: a truncation at a chunk boundary.
		last := size % chunk
		if last == 0 {
			last = chunk
		}
		for _, cut := range []int{len(ct) - 1, len(ct) - (last + 16)} {
			if _, err := decryptAll(ct[:cut], id); !errors.Is(err, datekeys.ErrIntegrity) {
				t.Errorf("%d bytes cut to %d of %d: %v", size, cut, len(ct), err)
			}
		}
	}
}

// Review: a panic of a parser of security fails only its own part. recovered
// tells a function that returns from one that panics.
func TestRecovered(t *testing.T) {
	if !recovered(func() {}) || recovered(func() { panic("a parser") }) {
		t.Error("recovered")
	}
}