package tce

import (
	"bytes"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"testing"
)

// Unit tests for the TCE primitives, mapping to PROTOCOL.md sections 2, 4,
// 5, 6, 7 and 12. The frozen vectors are exercised separately in
// vectors_test.go; these tests cover the individual rules with hand-built
// inputs.

func mustEncode(t *testing.T, build func(e *Encoder)) []byte {
	t.Helper()
	e := NewEncoder()
	build(e)
	b, err := e.Bytes()
	if err != nil {
		t.Fatalf("encode: %v", err)
	}
	return b
}

// claimBytes builds a minimal valid claim TCE for splicing into malformed
// inputs. It is the same shape as the frozen claim/boolean vector.
func claimBytes(t *testing.T) []byte {
	t.Helper()
	issuer := make([]byte, PubKeySize)
	subject := make([]byte, PubKeySize)
	copy(issuer, []byte{0x8a})
	copy(subject, []byte{0x81})
	return mustEncode(t, func(e *Encoder) {
		e.Header(TagClaim)
		e.Identity("issuer", issuer)
		e.Identity("subject", subject)
		e.Map("claims", map[string]Value{"example.flag": Bool(true)}, 1)
		e.Timestamp("created_at", 1_700_000_000, false)
		e.Timestamp("expires_at", 1_700_086_400, true)
		e.Uvarint(1)
		e.FixedBytes("nonce", []byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15}, NonceSize)
	})
}

func TestHeader(t *testing.T) {
	good := claimBytes(t)

	// The first 22 bytes are magic plus the object tag, then the version.
	if !bytes.HasPrefix(good, []byte(Magic)) {
		t.Fatal("object does not start with the magic")
	}

	t.Run("round trip", func(t *testing.T) {
		d := NewDecoder(good)
		tag, err := d.Header()
		if err != nil {
			t.Fatalf("header: %v", err)
		}
		if tag != TagClaim {
			t.Fatalf("tag = %v, want claim", tag)
		}
	})

	t.Run("empty input", func(t *testing.T) {
		d := NewDecoder(nil)
		if _, err := d.Header(); !errors.Is(err, ErrTruncated) {
			t.Fatalf("empty input: err = %v, want ErrTruncated", err)
		}
	})

	t.Run("truncated magic", func(t *testing.T) {
		d := NewDecoder(good[:MagicLen-1])
		if _, err := d.Header(); !errors.Is(err, ErrTruncated) {
			t.Fatalf("truncated magic: err = %v", err)
		}
	})

	t.Run("wrong magic", func(t *testing.T) {
		// Bump the framing version digit inside the magic.
		bad := bytes.Replace(good, []byte("tce/1\x00"), []byte("tce/2\x00"), 1)
		d := NewDecoder(bad)
		if _, err := d.Header(); !errors.Is(err, ErrMagic) {
			t.Fatalf("wrong magic: err = %v", err)
		}
	})

	for name, tag := range map[string]byte{
		"unknown tag 0x7f":  0x7f,
		"reserved tag 0x00": 0x00,
	} {
		t.Run(name, func(t *testing.T) {
			mut := make([]byte, len(good))
			copy(mut, good)
			mut[MagicLen] = tag
			d := NewDecoder(mut)
			if _, err := d.Header(); !errors.Is(err, ErrObjectTag) {
				t.Fatalf("err = %v, want ErrObjectTag", err)
			}
		})
	}

	t.Run("unknown version", func(t *testing.T) {
		d := NewDecoder(good)
		if _, err := d.Header(); err != nil {
			t.Fatal(err)
		}
		// Re-run a decoder whose version field says 2.
		mut := splatVersion(t, good, 2)
		if _, err := NewDecoder(mut).Header(); !errors.Is(err, ErrVersion) {
			t.Fatalf("unknown version: err = %v", err)
		}
	})

	t.Run("non-minimal version uvarint", func(t *testing.T) {
		// Replace the single-byte version (0x01) with its two-byte spelling.
		off := MagicLen + 1
		mut := make([]byte, 0, len(good)+1)
		mut = append(mut, good[:off]...)
		mut = append(mut, 0x81, 0x00)
		mut = append(mut, good[off+1:]...)
		if _, err := NewDecoder(mut).Header(); !errors.Is(err, ErrNonMinimal) {
			t.Fatalf("non-minimal uvarint: err = %v", err)
		}
	})
}

// splatVersion rewrites the version uvarint of an object to v.
func splatVersion(t *testing.T, b []byte, v uint64) []byte {
	t.Helper()
	off := MagicLen + 1
	out := make([]byte, 0, len(b))
	out = append(out, b[:off]...)
	out = AppendUvarint(out, v)
	out = append(out, b[off+1:]...)
	return out
}

func TestUvarintRoundTrip(t *testing.T) {
	values := []uint64{0, 1, 2, 127, 128, 16383, 16384, 1 << 16, 1 << 32, (1 << 63) - 1, 1<<64 - 1}
	for _, n := range values {
		enc := AppendUvarint(nil, n)
		if len(enc) != UvarintLen(n) {
			t.Errorf("UvarintLen(%d) = %d, want %d", n, UvarintLen(n), len(enc))
		}
		got, err := NewDecoder(enc).Uvarint()
		if err != nil {
			t.Fatalf("decode of %d: %v", n, err)
		}
		if got != n {
			t.Fatalf("round trip of %d gave %d", n, got)
		}
	}
}

func TestUvarintRejections(t *testing.T) {
	cases := []struct {
		name string
		in   []byte
		want error
	}{
		{"empty", nil, ErrTruncated},
		{"truncated mid-continuation", []byte{0x80}, ErrTruncated},
		{"non-minimal 0x81 0x00", []byte{0x81, 0x00}, ErrNonMinimal},
		{"non-minimal two redundant groups", []byte{0xff, 0x81, 0x00}, ErrNonMinimal},
		{"overlong ten continuations overflow first", []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01}, ErrOverflow},
		{"overflow bit 63 set as 2", []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x02}, ErrOverflow},
		// 2^64-1 fits in 10 bytes and is the largest accepted value.
		{"max value accepted", []byte{0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01}, nil},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			_, err := NewDecoder(c.in).Uvarint()
			if c.want == nil {
				if err != nil {
					t.Fatalf("want success, got %v", err)
				}
				return
			}
			if !errors.Is(err, c.want) {
				t.Fatalf("err = %v, want %v", err, c.want)
			}
		})
	}
}

func TestStringRules(t *testing.T) {
	valid := []string{
		"", "a", "hello world", "NikoCraft",
		// U+FEFF is an ordinary character, not a BOM to strip.
		"\ufeffbom",
		strings.Repeat("a", MaxStringValue),
		// Multibyte characters are fine as long as they are not controls.
		"café ☕ 中文",
	}
	for _, s := range valid {
		if err := ValidateString(s, MaxStringValue); err != nil {
			t.Errorf("ValidateString(%q) = %v, want nil", s, err)
		}
	}

	tooLong := strings.Repeat("a", MaxStringValue+1)
	if err := ValidateString(tooLong, MaxStringValue); !errors.Is(err, ErrTooLong) {
		t.Errorf("overlong string: err = %v", err)
	}
	for _, s := range []string{"a\xc3\x28", "\xed\xa0\x80", "a\xff\xfe"} {
		if err := ValidateString(s, 100); !errors.Is(err, ErrUTF8) {
			t.Errorf("invalid utf8 %q: err = %v", s, err)
		}
	}
	for _, r := range []rune{0x00, 0x1f, 0x7f, 0x80, 0x9f} {
		if err := ValidateString(string(r), 100); !errors.Is(err, ErrControlChar) {
			t.Errorf("control U+%04X: err = %v", r, err)
		}
	}
}

func TestKeyGrammar(t *testing.T) {
	valid := []string{
		"a", "x0", "example.flag", "a.first", "a-b_c.d", "z", strings.Repeat("k", MaxKeyLen),
	}
	for _, k := range valid {
		if err := ValidateKey(k); err != nil {
			t.Errorf("ValidateKey(%q) = %v, want nil", k, err)
		}
	}
	invalid := []string{
		"A", "1a", ".a", "a.", "a..b", "a--b", "a__b", "a ", "a b", "_a", "-a",
	}
	for _, k := range invalid {
		if err := ValidateKey(k); !errors.Is(err, ErrKeyGrammar) {
			t.Errorf("ValidateKey(%q) err = %v, want ErrKeyGrammar", k, err)
		}
	}
	// Empty and over-long keys are also invalid, though they surface as the
	// length limit first, which is equally a rejection.
	for _, k := range []string{"", strings.Repeat("k", MaxKeyLen+1)} {
		if err := ValidateKey(k); err == nil {
			t.Errorf("ValidateKey(%q bytes=%d) accepted", k, len(k))
		}
	}
}

func TestTimestampRange(t *testing.T) {
	for _, ts := range []uint64{MinTimestamp, MaxTimestamp, 1_700_000_000} {
		if err := ValidateTimestamp(ts, false); err != nil {
			t.Errorf("timestamp %d rejected: %v", ts, err)
		}
	}
	for _, ts := range []uint64{0, MinTimestamp - 1, MaxTimestamp + 1} {
		if err := ValidateTimestamp(ts, false); !errors.Is(err, ErrTimestamp) {
			t.Errorf("timestamp %d accepted", ts)
		}
	}
	if err := ValidateTimestamp(0, true); err != nil {
		t.Errorf("zero with allowZero rejected: %v", err)
	}
}

func TestEncodeRejectsMalformedField(t *testing.T) {
	enc := NewEncoder()
	enc.FixedBytes("nonce", make([]byte, 15), NonceSize)
	if !errors.Is(enc.Err(), ErrFieldSize) {
		t.Fatalf("encoder accepted a short fixed field: %v", enc.Err())
	}

	enc = NewEncoder()
	enc.Identity("id", make([]byte, PubKeySize-1))
	if !errors.Is(enc.Err(), ErrFieldSize) {
		t.Fatalf("encoder accepted a short identity: %v", enc.Err())
	}

	enc = NewEncoder()
	enc.Timestamp("ts", 0, false)
	if !errors.Is(enc.Err(), ErrTimestamp) {
		t.Fatalf("encoder accepted zero timestamp: %v", enc.Err())
	}

	enc = NewEncoder()
	enc.Value("v", String("x"))
	if err := enc.Err(); err != nil {
		t.Fatal(err)
	}
}

func TestValueEncoding(t *testing.T) {
	cases := []struct {
		name string
		v    Value
		want []byte
	}{
		{"null", Null(), []byte{0x00}},
		{"false", Bool(false), []byte{0x01}},
		{"true", Bool(true), []byte{0x02}},
		{"string", String("hi"), []byte{byte(ValString), 0x02, 'h', 'i'}},
		{"number canonicalised", Number("1.0"), []byte{byte(ValNumber), 0x01, '1'}},
		{"number preserved as text", Number("42"), []byte{byte(ValNumber), 0x02, '4', '2'}},
	}
	for _, c := range cases {
		t.Run(c.name, func(t *testing.T) {
			b := mustEncode(t, func(e *Encoder) { e.Value("v", c.v) })
			if !bytes.Equal(b, c.want) {
				t.Fatalf("got %x want %x", b, c.want)
			}
		})
	}

	// A number that cannot be canonicalized must fail the encoder.
	enc := NewEncoder()
	enc.Value("v", Number("1e99999"))
	if !errors.Is(enc.Err(), ErrNumberRange) {
		t.Fatalf("encoder accepted a non-canonicalisable number: %v", enc.Err())
	}

	if !Null().Equal(Null()) || !Number("1.0").Equal(Number("1")) || String("a").Equal(String("b")) {
		t.Fatal("Value.Equal disagrees")
	}
	if Bool(false).Equal(Bool(true)) {
		t.Fatal("false equals true")
	}
}

func TestMapEncodingSortsAndRejects(t *testing.T) {
	t.Run("sorted canonical form", func(t *testing.T) {
		m := map[string]Value{
			"z.last":  Null(),
			"a.first": Bool(false),
			"m.mid":   Number("2"),
		}
		enc := NewEncoder()
		enc.Map("claims", m, 1)
		b, err := enc.Bytes()
		if err != nil {
			t.Fatal(err)
		}
		// 3 entries, keys in alphabetical order.
		if !bytes.HasPrefix(b, []byte{0x03, 0x07}) {
			t.Fatalf("expected count 3 then len 7 for a.first, got %x", b)
		}
		dec := NewDecoder(b)
		got, err := dec.Map(1)
		if err != nil {
			t.Fatal(err)
		}
		// Re-encoding the decoded map reproduces the bytes exactly.
		enc2 := NewEncoder()
		enc2.Map("claims", got, 1)
		b2, err := enc2.Bytes()
		if err != nil {
			t.Fatal(err)
		}
		if !bytes.Equal(b, b2) {
			t.Fatalf("decode(encode) not stable:\n%x\n%x", b, b2)
		}
	})

	t.Run("duplicate key rejected on decode", func(t *testing.T) {
		// count=2, then two entries with the same key length and content.
		raw := []byte{0x02, 0x01, 'a', byte(ValTrue), 0x01, 'a', byte(ValFalse)}
		if _, err := NewDecoder(raw).Map(1); !errors.Is(err, ErrDuplicateKey) {
			t.Fatalf("err = %v, want ErrDuplicateKey", err)
		}
	})

	t.Run("unsorted keys rejected on decode", func(t *testing.T) {
		raw := []byte{0x02, 0x01, 'b', byte(ValTrue), 0x01, 'a', byte(ValFalse)}
		if _, err := NewDecoder(raw).Map(1); !errors.Is(err, ErrKeyOrder) {
			t.Fatalf("err = %v, want ErrKeyOrder", err)
		}
	})

	t.Run("too many entries", func(t *testing.T) {
		// Encode 33 entries of a trivial key.
		m := make(map[string]Value, MaxMapEntries+1)
		for i := 0; i < MaxMapEntries+1; i++ {
			m[fmtKey(i)] = Bool(true)
		}
		enc := NewEncoder()
		enc.Map("claims", m, 1)
		if !errors.Is(enc.Err(), ErrTooLong) {
			t.Fatalf("encoder accepted %d map entries", MaxMapEntries+1)
		}
	})

	t.Run("empty map with minimum required", func(t *testing.T) {
		enc := NewEncoder()
		enc.Map("claims", nil, 1)
		if !errors.Is(enc.Err(), ErrEmptyMap) {
			t.Fatalf("encoder accepted an empty map with min 1: %v", enc.Err())
		}
	})

	t.Run("bare key grammar failure on decode", func(t *testing.T) {
		raw := []byte{0x01, 0x01, 'A', byte(ValTrue)}
		if _, err := NewDecoder(raw).Map(1); !errors.Is(err, ErrKeyGrammar) {
			t.Fatalf("err = %v, want ErrKeyGrammar", err)
		}
	})

	t.Run("length prefix exceeds remaining input", func(t *testing.T) {
		raw := []byte{0x01, 0x40, 'a'} // claims length 64 but only 1 byte remains
		if _, err := NewDecoder(raw).Map(1); !errors.Is(err, ErrTruncated) {
			t.Fatalf("err = %v, want ErrTruncated", err)
		}
	})
}

func fmtKey(i int) string {
	return fmt.Sprintf("k%d", i)
}

func TestValueDecoderRejectsReservedTag(t *testing.T) {
	for _, tag := range []byte{byte(ValResBytes), byte(ValResArray), byte(ValResMap), 0x08, 0xff} {
		d := NewDecoder([]byte{tag})
		if _, err := d.Value(); !errors.Is(err, ErrValueTag) {
			t.Fatalf("tag 0x%02x accepted", tag)
		}
	}

	// A number in a non-canonical spelling must be rejected, not re-canonicalized.
	raw := []byte{byte(ValNumber), 0x03, '1', '.', '0'}
	if _, err := NewDecoder(raw).Value(); !errors.Is(err, ErrNumberFormat) {
		t.Fatalf("non-canonical number accepted: %v", err)
	}
}

func TestFixedLengthDecode(t *testing.T) {
	// A correctly sized 16-byte nonce.
	raw := []byte{0x10, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
	d := NewDecoder(raw)
	b, err := d.FixedBytes(NonceSize)
	if err != nil {
		t.Fatal(err)
	}
	if len(b) != NonceSize {
		t.Fatal("wrong nonce length")
	}

	// Length prefix claims 8 bytes for a 16-byte field: the prefix is well
	// formed and within the remaining input, but the width is wrong.
	short := []byte{0x08, 1, 2, 3, 4, 5, 6, 7, 8}
	if _, err := NewDecoder(short).FixedBytes(NonceSize); !errors.Is(err, ErrFieldSize) {
		t.Fatalf("short fixed field: err = %v, want ErrFieldSize", err)
	}

	// Length prefix claims more than the remaining input.
	over := []byte{0x40, 1}
	if _, err := NewDecoder(over).FixedBytes(NonceSize); !errors.Is(err, ErrTruncated) {
		t.Fatalf("over-long fixed field accepted: %v", err)
	}
}

func TestTrailingBytesAreRejected(t *testing.T) {
	good := claimBytes(t)
	for _, n := range []int{1, 2} {
		trailing := append(append([]byte{}, good...), make([]byte, n)...)
		d := NewDecoder(trailing)
		if _, err := d.Header(); err != nil {
			t.Fatal(err)
		}
		if _, err := d.Identity(); err != nil {
			t.Fatal(err)
		}
		if _, err := d.Identity(); err != nil {
			t.Fatal(err)
		}
		if _, err := d.Map(1); err != nil {
			t.Fatal(err)
		}
		if _, err := d.Timestamp(false); err != nil {
			t.Fatal(err)
		}
		if _, err := d.Timestamp(true); err != nil {
			t.Fatal(err)
		}
		if _, err := d.Uvarint(); err != nil {
			t.Fatal(err)
		}
		if _, err := d.FixedBytes(NonceSize); err != nil {
			t.Fatal(err)
		}
		if err := d.End(); !errors.Is(err, ErrTrailing) {
			t.Fatalf("%d extra bytes: err = %v, want ErrTrailing", n, err)
		}
	}
}

func TestIdentityFieldRejectsBadAddressVersion(t *testing.T) {
	// address version 1 is not defined in v1.
	raw := []byte{0x01, 0x20}
	raw = append(raw, make([]byte, PubKeySize)...)
	if _, err := NewDecoder(raw).Identity(); !errors.Is(err, ErrAddressVersion) {
		t.Fatalf("err = %v, want ErrAddressVersion", err)
	}
}

func TestObjectID(t *testing.T) {
	x := claimBytes(t)
	id1 := ComputeID(x)
	id2 := ComputeID(append(append([]byte{}, x...), 0))
	if id1.Equal(id2) {
		t.Fatal("two different byte strings have the same ID")
	}
	if id1.String() != hex.EncodeToString(id1[:]) {
		t.Fatal("ID.String disagrees with hex")
	}
	parsed, err := ParseID(id1.String())
	if err != nil {
		t.Fatal(err)
	}
	if !parsed.Equal(id1) {
		t.Fatal("parse round trip failed")
	}
	for _, bad := range []string{"", "abcd", "ABCDABCDABCDABCDABCDABCDABCDABCDABCDABCDABCDABCDABCDABCDABCDABCD", "zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz"} {
		if _, err := ParseID(bad); err == nil {
			t.Fatalf("ParseID accepted %q", bad)
		}
	}
	if _, err := IDFromBytes(make([]byte, 31)); err == nil {
		t.Fatal("IDFromBytes accepted 31 bytes")
	}
}

func TestInjectiveEncodeDecode(t *testing.T) {
	b := claimBytes(t)

	// decode then encode must reproduce the bytes exactly, which is the
	// non-malleability guarantee of section 12.4.
	d := NewDecoder(b)
	if _, err := d.Header(); err != nil {
		t.Fatal(err)
	}
	issuer, err := d.Identity()
	if err != nil {
		t.Fatal(err)
	}
	subject, err := d.Identity()
	if err != nil {
		t.Fatal(err)
	}
	claims, err := d.Map(1)
	if err != nil {
		t.Fatal(err)
	}
	createdAt, err := d.Timestamp(false)
	if err != nil {
		t.Fatal(err)
	}
	expiresAt, err := d.Timestamp(true)
	if err != nil {
		t.Fatal(err)
	}
	serial, err := d.Uvarint()
	if err != nil {
		t.Fatal(err)
	}
	nonce, err := d.FixedBytes(NonceSize)
	if err != nil {
		t.Fatal(err)
	}
	if err := d.End(); err != nil {
		t.Fatal(err)
	}

	re := mustEncode(t, func(e *Encoder) {
		e.Header(TagClaim)
		e.Identity("issuer", issuer)
		e.Identity("subject", subject)
		e.Map("claims", claims, 1)
		e.Timestamp("created_at", createdAt, false)
		e.Timestamp("expires_at", expiresAt, true)
		e.Uvarint(serial)
		e.FixedBytes("nonce", nonce, NonceSize)
	})
	if !bytes.Equal(b, re) {
		t.Fatalf("encode(decode(b)) != b:\n%x\n%x", b, re)
	}
}