package axon_test

import (
	"crypto/ed25519"
	"encoding/hex"
	"encoding/json"
	"os"
	"path/filepath"
	"testing"
	"time"

	"code.northwest.io/axon"
)

// --------------------------------------------------------------------------
// Helpers
// --------------------------------------------------------------------------

func mustDecodeHex(t *testing.T, s string) []byte {
	t.Helper()
	b, err := hex.DecodeString(s)
	if err != nil {
		t.Fatalf("hex decode %q: %v", s, err)
	}
	return b
}

func hexOf(b []byte) string { return hex.EncodeToString(b) }

// knownSeed returns a fixed 32-byte seed so tests are fully deterministic.
func knownSeed(n byte) []byte {
	seed := make([]byte, 32)
	for i := range seed {
		seed[i] = n + byte(i)
	}
	return seed
}

// --------------------------------------------------------------------------
// Test: canonical_tags
// --------------------------------------------------------------------------

func TestCanonicalTags(t *testing.T) {
	vectors := loadVectors(t)
	tags := vectors.TagsInput

	enc, err := axon.CanonicalTags(tags)
	if err != nil {
		t.Fatalf("CanonicalTags: %v", err)
	}
	if len(enc) == 0 {
		t.Fatal("CanonicalTags returned empty bytes")
	}

	h, err := axon.CanonicalTagsHash(tags)
	if err != nil {
		t.Fatalf("CanonicalTagsHash: %v", err)
	}
	if h == [32]byte{} {
		t.Fatal("CanonicalTagsHash returned zero hash")
	}

	// Verify sort order: tags with the same name are sorted by first value.
	// "e"/"reply-id" < "e"/"root-id" lexicographically. "e" < "p".
	// So expected order: e/reply-id, e/root-id, p/alice
	expectHash := mustDecodeHex(t, vectors.CanonicalTagsHash)
	if hexOf(h[:]) != hexOf(expectHash) {
		t.Errorf("CanonicalTagsHash mismatch:\n  got  %s\n  want %s", hexOf(h[:]), hexOf(expectHash))
	}

	expectEnc := mustDecodeHex(t, vectors.CanonicalTagsHex)
	if hexOf(enc) != hexOf(expectEnc) {
		t.Errorf("canonical_tags bytes mismatch:\n  got  %s\n  want %s", hexOf(enc), hexOf(expectEnc))
	}
}

func TestCanonicalTagsDuplicateError(t *testing.T) {
	tags := []axon.Tag{
		{Name: "e", Values: []string{"same-id", "root"}},
		{Name: "e", Values: []string{"same-id", "reply"}},
	}
	_, err := axon.CanonicalTags(tags)
	if err == nil {
		t.Fatal("expected error for duplicate (name, first_value) pair")
	}
}

func TestCanonicalTagsEmpty(t *testing.T) {
	enc, err := axon.CanonicalTags(nil)
	if err != nil {
		t.Fatalf("CanonicalTags(nil): %v", err)
	}
	// Should encode as uint16(0) = 2 zero bytes.
	if len(enc) != 2 || enc[0] != 0 || enc[1] != 0 {
		t.Errorf("empty tags encoding: got %x, want 0000", enc)
	}
}

// --------------------------------------------------------------------------
// Test: canonical_payload
// --------------------------------------------------------------------------

func TestCanonicalPayload(t *testing.T) {
	vectors := loadVectors(t)
	kp := axon.NewKeyPairFromSeed(mustDecodeHex(t, vectors.SeedHex))
	content := mustDecodeHex(t, vectors.ContentHex)

	payload, err := axon.CanonicalPayload(kp.PubKey, vectors.CreatedAt, vectors.Kind, content, vectors.TagsInput)
	if err != nil {
		t.Fatalf("CanonicalPayload: %v", err)
	}

	// Validate structure: first 2 bytes must be 0x00 0x20 (= 32).
	if payload[0] != 0 || payload[1] != 32 {
		t.Errorf("payload[0:2] = %x, want 0020", payload[0:2])
	}
	// Bytes [2:34] = pubkey.
	if hexOf(payload[2:34]) != hexOf(kp.PubKey) {
		t.Errorf("pubkey in payload mismatch")
	}
	// Total length: 2 + 32 + 8 + 2 + 4 + len(content) + 32
	want := 2 + 32 + 8 + 2 + 4 + len(content) + 32
	if len(payload) != want {
		t.Errorf("payload length = %d, want %d", len(payload), want)
	}

	expectPayload := mustDecodeHex(t, vectors.CanonicalPayload)
	if hexOf(payload) != hexOf(expectPayload) {
		t.Errorf("canonical_payload mismatch:\n  got  %s\n  want %s", hexOf(payload), hexOf(expectPayload))
	}
}

// --------------------------------------------------------------------------
// Test: event id
// --------------------------------------------------------------------------

func TestEventID(t *testing.T) {
	vectors := loadVectors(t)
	kp := axon.NewKeyPairFromSeed(mustDecodeHex(t, vectors.SeedHex))
	content := mustDecodeHex(t, vectors.ContentHex)

	id, err := axon.EventID(kp.PubKey, vectors.CreatedAt, vectors.Kind, content, vectors.TagsInput)
	if err != nil {
		t.Fatalf("EventID: %v", err)
	}
	if len(id) != 32 {
		t.Fatalf("EventID returned %d bytes, want 32", len(id))
	}

	expectID := mustDecodeHex(t, vectors.EventID)
	if hexOf(id) != hexOf(expectID) {
		t.Errorf("event id mismatch:\n  got  %s\n  want %s", hexOf(id), hexOf(expectID))
	}
}

// --------------------------------------------------------------------------
// Test: sign and verify
// --------------------------------------------------------------------------

func TestSignVerify(t *testing.T) {
	kp := axon.NewKeyPairFromSeed(knownSeed(0x01))

	e := &axon.Event{
		CreatedAt: 1700000000,
		Kind:      axon.KindMessage,
		Content:   []byte("hello axon"),
		Tags:      []axon.Tag{{Name: "p", Values: []string{"alice"}}},
	}
	if err := axon.Sign(e, kp); err != nil {
		t.Fatalf("Sign: %v", err)
	}

	if len(e.ID) != 32 {
		t.Errorf("ID is %d bytes, want 32", len(e.ID))
	}
	if len(e.Sig) != ed25519.SignatureSize {
		t.Errorf("Sig is %d bytes, want 64", len(e.Sig))
	}

	if err := axon.Verify(e); err != nil {
		t.Fatalf("Verify: %v", err)
	}

	// Tampering with content must invalidate.
	e2 := *e
	e2.Content = []byte("tampered")
	if err := axon.Verify(&e2); err == nil {
		t.Fatal("Verify should fail after content tamper")
	}

	// Tampering with signature must invalidate.
	e3 := *e
	sig := make([]byte, len(e.Sig))
	copy(sig, e.Sig)
	sig[0] ^= 0xff
	e3.Sig = sig
	if err := axon.Verify(&e3); err == nil {
		t.Fatal("Verify should fail after sig tamper")
	}
}

func TestVerifyWrongKey(t *testing.T) {
	kp1 := axon.NewKeyPairFromSeed(knownSeed(0x01))
	kp2 := axon.NewKeyPairFromSeed(knownSeed(0x02))

	e := &axon.Event{
		CreatedAt: time.Now().Unix(),
		Kind:      axon.KindMessage,
		Content:   []byte("test"),
	}
	if err := axon.Sign(e, kp1); err != nil {
		t.Fatalf("Sign: %v", err)
	}

	// Replace pubkey with kp2's pubkey; should fail.
	e.PubKey = kp2.PubKey
	if err := axon.Verify(e); err == nil {
		t.Fatal("Verify should fail with wrong pubkey")
	}
}

// --------------------------------------------------------------------------
// Test: challenge sign/verify
// --------------------------------------------------------------------------

func TestChallengeSignVerify(t *testing.T) {
	kp := axon.NewKeyPairFromSeed(knownSeed(0x03))
	nonce := make([]byte, 32)
	for i := range nonce {
		nonce[i] = byte(i)
	}
	relayURL := "wss://relay.example.com"

	sig := axon.SignChallenge(kp, nonce, relayURL)
	if !axon.VerifyChallenge(kp.PubKey, nonce, relayURL, sig) {
		t.Fatal("VerifyChallenge failed for valid signature")
	}
	if axon.VerifyChallenge(kp.PubKey, nonce, "wss://other.relay.com", sig) {
		t.Fatal("VerifyChallenge should fail for different relay URL")
	}
}

// --------------------------------------------------------------------------
// Test: X25519 key conversion
// --------------------------------------------------------------------------

func TestX25519Conversion(t *testing.T) {
	kp := axon.NewKeyPairFromSeed(knownSeed(0x01))

	// Converting same Ed25519 keypair from both sides should yield same shared secret.
	kp2 := axon.NewKeyPairFromSeed(knownSeed(0x02))

	shared1, err := axon.DHSharedSecret(kp.PrivKey, kp2.PubKey)
	if err != nil {
		t.Fatalf("DHSharedSecret(kp->kp2): %v", err)
	}
	shared2, err := axon.DHSharedSecret(kp2.PrivKey, kp.PubKey)
	if err != nil {
		t.Fatalf("DHSharedSecret(kp2->kp): %v", err)
	}
	if shared1 != shared2 {
		t.Errorf("DH not commutative:\n  kp->kp2: %s\n  kp2->kp: %s",
			hexOf(shared1[:]), hexOf(shared2[:]))
	}
	if shared1 == [32]byte{} {
		t.Fatal("shared secret is all zeros")
	}
}

// --------------------------------------------------------------------------
// Test: DM encrypt/decrypt round-trip
// --------------------------------------------------------------------------

func TestDMEncryptDecrypt(t *testing.T) {
	alice := axon.NewKeyPairFromSeed(knownSeed(0x0a))
	bob := axon.NewKeyPairFromSeed(knownSeed(0x0b))

	plaintext := []byte("secret message from alice to bob")

	ct, err := axon.EncryptDM(alice.PrivKey, bob.PubKey, plaintext)
	if err != nil {
		t.Fatalf("EncryptDM: %v", err)
	}

	pt, err := axon.DecryptDM(bob.PrivKey, alice.PubKey, ct)
	if err != nil {
		t.Fatalf("DecryptDM: %v", err)
	}
	if string(pt) != string(plaintext) {
		t.Errorf("decrypt mismatch: got %q, want %q", pt, plaintext)
	}

	// Tamper with ciphertext — authentication should fail.
	ct2 := make([]byte, len(ct))
	copy(ct2, ct)
	ct2[len(ct2)-1] ^= 0xff
	if _, err := axon.DecryptDM(bob.PrivKey, alice.PubKey, ct2); err == nil {
		t.Fatal("DecryptDM should fail on tampered ciphertext")
	}

	// Decrypt with wrong key should fail.
	carol := axon.NewKeyPairFromSeed(knownSeed(0x0c))
	if _, err := axon.DecryptDM(carol.PrivKey, alice.PubKey, ct); err == nil {
		t.Fatal("DecryptDM should fail with wrong recipient key")
	}
}

// --------------------------------------------------------------------------
// Test: MessagePack encode/decode
// --------------------------------------------------------------------------

func TestMsgPackRoundTrip(t *testing.T) {
	kp := axon.NewKeyPairFromSeed(knownSeed(0x05))
	e := &axon.Event{
		CreatedAt: 1700000000,
		Kind:      axon.KindMessage,
		Content:   []byte("msgpack test"),
		Tags: []axon.Tag{
			{Name: "p", Values: []string{"recipient"}},
			{Name: "e", Values: []string{"ref-id", "root"}},
		},
	}
	if err := axon.Sign(e, kp); err != nil {
		t.Fatalf("Sign: %v", err)
	}

	data, err := axon.MarshalEvent(e)
	if err != nil {
		t.Fatalf("MarshalEvent: %v", err)
	}
	if len(data) == 0 {
		t.Fatal("MarshalEvent returned empty bytes")
	}

	e2, err := axon.UnmarshalEvent(data)
	if err != nil {
		t.Fatalf("UnmarshalEvent: %v", err)
	}

	// Verify round-tripped event is still valid.
	if err := axon.Verify(e2); err != nil {
		t.Fatalf("Verify after msgpack round-trip: %v", err)
	}

	if hexOf(e.ID) != hexOf(e2.ID) {
		t.Errorf("ID mismatch after round-trip")
	}
	if hexOf(e.Sig) != hexOf(e2.Sig) {
		t.Errorf("Sig mismatch after round-trip")
	}
	if string(e.Content) != string(e2.Content) {
		t.Errorf("Content mismatch after round-trip: got %q, want %q", e2.Content, e.Content)
	}
}

// --------------------------------------------------------------------------
// Test vector generation and validation
// --------------------------------------------------------------------------

// Vectors is the schema for testdata/vectors.json.
type Vectors struct {
	// canonical_tags vector
	TagsInput         []axon.Tag `json:"tags_input"`
	CanonicalTagsHex  string     `json:"canonical_tags_hex"`
	CanonicalTagsHash string     `json:"canonical_tags_hash"`

	// canonical_payload vector
	SeedHex          string `json:"seed_hex"`
	PubKeyHex        string `json:"pubkey_hex"`
	CreatedAt        int64  `json:"created_at"`
	Kind             uint16 `json:"kind"`
	ContentHex       string `json:"content_hex"`
	CanonicalPayload string `json:"canonical_payload"`

	// event id
	EventID string `json:"event_id"`

	// signature (sign with known seed; Ed25519 is deterministic so sig is stable)
	SigHex string `json:"sig_hex"`
}

func vectorsPath() string {
	return filepath.Join("testdata", "vectors.json")
}

func loadVectors(t *testing.T) Vectors {
	t.Helper()
	data, err := os.ReadFile(vectorsPath())
	if err != nil {
		t.Fatalf("load vectors: %v (run TestGenerateVectors first)", err)
	}
	var v Vectors
	if err := json.Unmarshal(data, &v); err != nil {
		t.Fatalf("parse vectors: %v", err)
	}
	return v
}

// TestGenerateVectors generates testdata/vectors.json from known inputs.
// Run once with -run TestGenerateVectors to produce the file, then the other
// tests validate against it. The file is committed to the repo.
func TestGenerateVectors(t *testing.T) {
	seed := knownSeed(0x01)
	kp := axon.NewKeyPairFromSeed(seed)

	tags := []axon.Tag{
		{Name: "p", Values: []string{"alice"}},
		{Name: "e", Values: []string{"root-id", "root"}},
		{Name: "e", Values: []string{"reply-id", "reply"}},
	}
	content := []byte("hello axon")
	createdAt := int64(1700000000)
	kind := axon.KindMessage

	// canonical_tags
	tagsEnc, err := axon.CanonicalTags(tags)
	if err != nil {
		t.Fatalf("CanonicalTags: %v", err)
	}
	tagsHash, err := axon.CanonicalTagsHash(tags)
	if err != nil {
		t.Fatalf("CanonicalTagsHash: %v", err)
	}

	// canonical_payload
	payload, err := axon.CanonicalPayload(kp.PubKey, createdAt, kind, content, tags)
	if err != nil {
		t.Fatalf("CanonicalPayload: %v", err)
	}

	// event id
	id, err := axon.EventID(kp.PubKey, createdAt, kind, content, tags)
	if err != nil {
		t.Fatalf("EventID: %v", err)
	}

	// sign
	e := &axon.Event{
		CreatedAt: createdAt,
		Kind:      kind,
		Content:   content,
		Tags:      tags,
	}
	if err := axon.Sign(e, kp); err != nil {
		t.Fatalf("Sign: %v", err)
	}
	// Verify the event so we know the sig is valid.
	if err := axon.Verify(e); err != nil {
		t.Fatalf("Verify generated event: %v", err)
	}

	v := Vectors{
		TagsInput:         tags,
		CanonicalTagsHex:  hexOf(tagsEnc),
		CanonicalTagsHash: hexOf(tagsHash[:]),

		SeedHex:          hexOf(seed),
		PubKeyHex:        hexOf(kp.PubKey),
		CreatedAt:        createdAt,
		Kind:             kind,
		ContentHex:       hexOf(content),
		CanonicalPayload: hexOf(payload),

		EventID: hexOf(id),
		SigHex:  hexOf(e.Sig),
	}

	data, err := json.MarshalIndent(v, "", "  ")
	if err != nil {
		t.Fatalf("marshal vectors: %v", err)
	}

	if err := os.MkdirAll("testdata", 0o755); err != nil {
		t.Fatalf("mkdir testdata: %v", err)
	}
	if err := os.WriteFile(vectorsPath(), data, 0o644); err != nil {
		t.Fatalf("write vectors: %v", err)
	}

	t.Logf("vectors written to %s", vectorsPath())
	t.Logf("canonical_tags_hash: %s", hexOf(tagsHash[:]))
	t.Logf("event_id:            %s", hexOf(id))
}

// TestVectors validates all vectors in testdata/vectors.json. This is the
// primary cross-language compatibility check.
func TestVectors(t *testing.T) {
	v := loadVectors(t)

	seed := mustDecodeHex(t, v.SeedHex)
	kp := axon.NewKeyPairFromSeed(seed)

	if hexOf(kp.PubKey) != v.PubKeyHex {
		t.Errorf("pubkey mismatch:\n  got  %s\n  want %s", hexOf(kp.PubKey), v.PubKeyHex)
	}

	// canonical_tags
	tagsEnc, err := axon.CanonicalTags(v.TagsInput)
	if err != nil {
		t.Fatalf("CanonicalTags: %v", err)
	}
	if hexOf(tagsEnc) != v.CanonicalTagsHex {
		t.Errorf("canonical_tags bytes mismatch:\n  got  %s\n  want %s",
			hexOf(tagsEnc), v.CanonicalTagsHex)
	}

	tagsHash, err := axon.CanonicalTagsHash(v.TagsInput)
	if err != nil {
		t.Fatalf("CanonicalTagsHash: %v", err)
	}
	if hexOf(tagsHash[:]) != v.CanonicalTagsHash {
		t.Errorf("canonical_tags_hash mismatch:\n  got  %s\n  want %s",
			hexOf(tagsHash[:]), v.CanonicalTagsHash)
	}

	// canonical_payload
	content := mustDecodeHex(t, v.ContentHex)
	payload, err := axon.CanonicalPayload(kp.PubKey, v.CreatedAt, v.Kind, content, v.TagsInput)
	if err != nil {
		t.Fatalf("CanonicalPayload: %v", err)
	}
	if hexOf(payload) != v.CanonicalPayload {
		t.Errorf("canonical_payload mismatch:\n  got  %s\n  want %s",
			hexOf(payload), v.CanonicalPayload)
	}

	// event id
	id, err := axon.EventID(kp.PubKey, v.CreatedAt, v.Kind, content, v.TagsInput)
	if err != nil {
		t.Fatalf("EventID: %v", err)
	}
	if hexOf(id) != v.EventID {
		t.Errorf("event_id mismatch:\n  got  %s\n  want %s", hexOf(id), v.EventID)
	}

	// verify the stored signature
	sig := mustDecodeHex(t, v.SigHex)
	e := &axon.Event{
		ID:        id,
		PubKey:    kp.PubKey,
		CreatedAt: v.CreatedAt,
		Kind:      v.Kind,
		Content:   content,
		Sig:       sig,
		Tags:      v.TagsInput,
	}
	if err := axon.Verify(e); err != nil {
		t.Errorf("Verify stored vector event: %v", err)
	}
}