package crypto

import (
	"crypto/ed25519"
	"crypto/rand"
	"encoding/base64"
	"strings"
	"testing"
)

func generateTestKeys(t *testing.T) (ed25519.PublicKey, ed25519.PrivateKey) {
	t.Helper()
	pub, pri, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		t.Fatalf("generate ed25519 key pair: %v", err)
	}
	return pub, pri
}

func TestSetKeys(t *testing.T) {
	pub, pri := generateTestKeys(t)

	err := SetKeys(pub, pri)
	if err != nil {
		t.Fatalf("SetKeys: %v", err)
	}
	if !KeysSet() {
		t.Fatal("expected KeysSet() = true")
	}
}

func TestSetKeys_InvalidSize(t *testing.T) {
	err := SetKeys([]byte("short"), []byte("key"))
	if err != ErrInvalidKeySize {
		t.Fatalf("expected ErrInvalidKeySize, got %v", err)
	}
}

func TestEncryptDecrypt_Basic(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	plaintext := "sk-abc123my-secret-api-key"

	ciphertext, err := Encrypt(plaintext)
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}

	if ciphertext == plaintext {
		t.Fatal("ciphertext should differ from plaintext")
	}

	decrypted, err := Decrypt(ciphertext)
	if err != nil {
		t.Fatalf("Decrypt: %v", err)
	}

	if decrypted != plaintext {
		t.Fatalf("expected %q, got %q", plaintext, decrypted)
	}
}

func TestEncryptDecrypt_EmptyString(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	ciphertext, err := Encrypt("")
	if err != nil {
		t.Fatalf("Encrypt empty: %v", err)
	}

	decrypted, err := Decrypt(ciphertext)
	if err != nil {
		t.Fatalf("Decrypt empty: %v", err)
	}

	if decrypted != "" {
		t.Fatalf("expected empty string, got %q", decrypted)
	}
}

func TestEncryptDecrypt_Unicode(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	plaintext := "密钥-🔑-тест-αβγ"

	ciphertext, err := Encrypt(plaintext)
	if err != nil {
		t.Fatalf("Encrypt: %v", err)
	}

	decrypted, err := Decrypt(ciphertext)
	if err != nil {
		t.Fatalf("Decrypt: %v", err)
	}

	if decrypted != plaintext {
		t.Fatalf("expected %q, got %q", plaintext, decrypted)
	}
}

func TestEncryptDecrypt_LongKey(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	plaintext := strings.Repeat("a", 10000)

	ciphertext, err := Encrypt(plaintext)
	if err != nil {
		t.Fatalf("Encrypt long: %v", err)
	}

	decrypted, err := Decrypt(ciphertext)
	if err != nil {
		t.Fatalf("Decrypt long: %v", err)
	}

	if decrypted != plaintext {
		t.Fatalf("long string mismatch: got len %d, want len %d", len(decrypted), len(plaintext))
	}
}

func TestEncrypt_NoKeysSet(t *testing.T) {
	ResetKeys()

	_, err := Encrypt("test")
	if err != ErrNoKeysSet {
		t.Fatalf("expected ErrNoKeysSet, got %v", err)
	}
}

func TestDecrypt_NoKeysSet(t *testing.T) {
	ResetKeys()

	_, err := Decrypt("dGVzdA==")
	if err != ErrNoKeysSet {
		t.Fatalf("expected ErrNoKeysSet, got %v", err)
	}
}

func TestDecrypt_InvalidBase64(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	_, err := Decrypt("not-valid-base64!!!")
	if err == nil {
		t.Fatal("expected error for invalid base64")
	}
}

func TestDecrypt_WrongKey(t *testing.T) {
	pub1, pri1 := generateTestKeys(t)
	pub2, pri2 := generateTestKeys(t)

	if err := SetKeys(pub1, pri1); err != nil {
		t.Fatal(err)
	}

	ciphertext, err := Encrypt("secret-data")
	if err != nil {
		t.Fatal(err)
	}

	if err := SetKeys(pub2, pri2); err != nil {
		t.Fatal(err)
	}

	_, err = Decrypt(ciphertext)
	if err == nil {
		t.Fatal("expected decryption to fail with wrong key")
	}
}

func TestEncryptDecrypt_DifferentKeys(t *testing.T) {
	keys := make([]struct {
		pub ed25519.PublicKey
		pri ed25519.PrivateKey
	}, 5)

	for i := range keys {
		keys[i].pub, keys[i].pri, _ = ed25519.GenerateKey(rand.Reader)
	}

	for i, k := range keys {
		if err := SetKeys(k.pub, k.pri); err != nil {
			t.Fatal(err)
		}

		plaintext := "key-test-" + string(rune('A'+i))

		ciphertext, err := Encrypt(plaintext)
		if err != nil {
			t.Fatalf("Encrypt with key %d: %v", i, err)
		}

		decrypted, err := Decrypt(ciphertext)
		if err != nil {
			t.Fatalf("Decrypt with key %d: %v", i, err)
		}

		if decrypted != plaintext {
			t.Fatalf("key %d: expected %q, got %q", i, plaintext, decrypted)
		}
	}
}

func TestEncrypt_ProducesDifferentCiphertext(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	plaintext := "same-input"

	ct1, err := Encrypt(plaintext)
	if err != nil {
		t.Fatal(err)
	}

	ct2, err := Encrypt(plaintext)
	if err != nil {
		t.Fatal(err)
	}

	if ct1 == ct2 {
		t.Fatal("encrypting the same plaintext twice should produce different ciphertexts (random nonce)")
	}
}

func TestEncryptDecrypt_DeterministicRoundtrip(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	plaintext := "roundtrip-test-key"

	ct, err := Encrypt(plaintext)
	if err != nil {
		t.Fatal(err)
	}

	for i := 0; i < 10; i++ {
		dec, err := Decrypt(ct)
		if err != nil {
			t.Fatalf("attempt %d: Decrypt: %v", i, err)
		}
		if dec != plaintext {
			t.Fatalf("attempt %d: expected %q, got %q", i, plaintext, dec)
		}
	}
}

func TestIsEncrypted(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	ct, err := Encrypt("test")
	if err != nil {
		t.Fatal(err)
	}

	if !IsEncrypted(ct) {
		t.Fatal("expected ciphertext to be detected as encrypted")
	}

	if IsEncrypted("plain-text") {
		t.Fatal("expected plain text to not be detected as encrypted")
	}

	if IsEncrypted("") {
		t.Fatal("expected empty string to not be detected as encrypted")
	}

	if IsEncrypted("not!!base64") {
		t.Fatal("expected invalid base64 to not be detected as encrypted")
	}
}

func TestEncryptDecrypt_ThreeRounds(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	plaintext := "verify-three-rounds-encryption"

	ct, err := Encrypt(plaintext)
	if err != nil {
		t.Fatal(err)
	}

	data, err := base64.StdEncoding.DecodeString(ct)
	if err != nil {
		t.Fatal(err)
	}

	for round := encryptionRounds - 1; round >= 1; round-- {
		decrypted, err := open(data)
		if err != nil {
			t.Fatalf("manual round %d: %v", round, err)
		}
		data = decrypted
	}

	intermediate, err := open(data)
	if err != nil {
		t.Fatalf("final manual round failed: %v", err)
	}

	if string(intermediate) != plaintext {
		t.Fatalf("after peeling all 3 rounds, expected %q, got %q", plaintext, string(intermediate))
	}
}

func TestKeysSet_AfterReset(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}
	if !KeysSet() {
		t.Fatal("expected keys to be set")
	}

	ResetKeys()
	if KeysSet() {
		t.Fatal("expected keys to not be set after reset")
	}
}

func TestDecrypt_TooShort(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	shortCiphertext := "AQID"
	_, err := Decrypt(shortCiphertext)
	if err == nil {
		t.Fatal("expected error for too-short ciphertext")
	}
}

func TestEncryptDecrypt_SpecialChars(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	testCases := []string{
		"sk-proj-abc123_DEF-456/ghi=789+jkl",
		"key\twith\nspecial\rchars",
		"null\x00byte",
		`{"api_key": "sk-test-123", "model": "gpt-4"}`,
	}

	for _, plaintext := range testCases {
		ct, err := Encrypt(plaintext)
		if err != nil {
			t.Fatalf("Encrypt(%q): %v", plaintext, err)
		}

		dec, err := Decrypt(ct)
		if err != nil {
			t.Fatalf("Decrypt(%q): %v", plaintext, err)
		}

		if dec != plaintext {
			t.Fatalf("roundtrip mismatch for %q: got %q", plaintext, dec)
		}
	}
}

func TestDecrypt_TamperedCiphertext(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	ct, err := Encrypt("secret")
	if err != nil {
		t.Fatal(err)
	}

	decoded, _ := base64.StdEncoding.DecodeString(ct)
	if len(decoded) > 60 {
		decoded[50] ^= 0xFF
		tampered := base64.StdEncoding.EncodeToString(decoded)
		_, err = Decrypt(tampered)
		if err == nil {
			t.Fatal("expected decryption to fail with tampered ciphertext")
		}
	}
}

func TestResetKeys_ClearsState(t *testing.T) {
	pub, pri := generateTestKeys(t)
	if err := SetKeys(pub, pri); err != nil {
		t.Fatal(err)
	}

	ResetKeys()

	if KeysSet() {
		t.Fatal("KeysSet should return false after ResetKeys")
	}

	_, err := Encrypt("test")
	if err != ErrNoKeysSet {
		t.Fatalf("expected ErrNoKeysSet after reset, got %v", err)
	}

	_, err = Decrypt("dGVzdA==")
	if err != ErrNoKeysSet {
		t.Fatalf("expected ErrNoKeysSet after reset, got %v", err)
	}
}