package ageio

import (
	"bytes"
	"os"
	"path/filepath"
	"testing"

	"filippo.io/age"
)

func TestEncryptDecryptRoundTrip(t *testing.T) {
	// Generate test key pair
	tmpDir := t.TempDir()
	identityPath, recipientPath, err := GenerateTestKeyPair(tmpDir)
	if err != nil {
		t.Fatalf("generate test key pair: %v", err)
	}

	// Create encryptor and decryptor
	enc, err := NewEncryptor(recipientPath)
	if err != nil {
		t.Fatalf("create encryptor: %v", err)
	}

	dec, err := NewDecryptor(identityPath)
	if err != nil {
		t.Fatalf("create decryptor: %v", err)
	}

	// Test data
	testCases := []struct {
		name      string
		plaintext []byte
	}{
		{
			name:      "simple text",
			plaintext: []byte("hello world"),
		},
		{
			name:      "json data",
			plaintext: []byte(`{"tool":"sequentialthinking","payload":{"thought":"test"}}`),
		},
		{
			name:      "large data",
			plaintext: bytes.Repeat([]byte("ABCDEFGHIJ"), 1000), // 10KB
		},
		{
			name:      "unicode",
			plaintext: []byte("Hello 世界 🌍"),
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			// Encrypt
			ciphertext, err := enc.Encrypt(tc.plaintext)
			if err != nil {
				t.Fatalf("encrypt: %v", err)
			}

			// Verify ciphertext is not empty and different from plaintext
			if len(ciphertext) == 0 {
				t.Fatal("ciphertext is empty")
			}

			if bytes.Equal(ciphertext, tc.plaintext) {
				t.Fatal("ciphertext equals plaintext (not encrypted)")
			}

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

			// Verify decrypted matches original
			if !bytes.Equal(decrypted, tc.plaintext) {
				t.Fatalf("decrypted data doesn't match original\ngot:  %q\nwant: %q", decrypted, tc.plaintext)
			}
		})
	}
}

func TestEncryptEmptyData(t *testing.T) {
	tmpDir := t.TempDir()
	_, recipientPath, err := GenerateTestKeyPair(tmpDir)
	if err != nil {
		t.Fatalf("generate test key pair: %v", err)
	}

	enc, err := NewEncryptor(recipientPath)
	if err != nil {
		t.Fatalf("create encryptor: %v", err)
	}

	_, err = enc.Encrypt([]byte{})
	if err == nil {
		t.Fatal("expected error encrypting empty data")
	}
}

func TestDecryptEmptyData(t *testing.T) {
	tmpDir := t.TempDir()
	identityPath, _, err := GenerateTestKeyPair(tmpDir)
	if err != nil {
		t.Fatalf("generate test key pair: %v", err)
	}

	dec, err := NewDecryptor(identityPath)
	if err != nil {
		t.Fatalf("create decryptor: %v", err)
	}

	_, err = dec.Decrypt([]byte{})
	if err == nil {
		t.Fatal("expected error decrypting empty data")
	}
}

func TestDecryptInvalidCiphertext(t *testing.T) {
	tmpDir := t.TempDir()
	identityPath, _, err := GenerateTestKeyPair(tmpDir)
	if err != nil {
		t.Fatalf("generate test key pair: %v", err)
	}

	dec, err := NewDecryptor(identityPath)
	if err != nil {
		t.Fatalf("create decryptor: %v", err)
	}

	// Try to decrypt garbage data
	_, err = dec.Decrypt([]byte("not a valid age ciphertext"))
	if err == nil {
		t.Fatal("expected error decrypting invalid ciphertext")
	}
}

func TestDecryptWrongKey(t *testing.T) {
	tmpDir := t.TempDir()

	// Generate two separate key pairs
	identity1Path := filepath.Join(tmpDir, "key1.age")
	recipient1Path := filepath.Join(tmpDir, "key1.pub")
	identity2Path := filepath.Join(tmpDir, "key2.age")

	// Create first key pair
	id1, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatalf("generate key 1: %v", err)
	}
	os.WriteFile(identity1Path, []byte(id1.String()+"\n"), 0600)
	os.WriteFile(recipient1Path, []byte(id1.Recipient().String()+"\n"), 0644)

	// Create second key pair
	id2, err := age.GenerateX25519Identity()
	if err != nil {
		t.Fatalf("generate key 2: %v", err)
	}
	os.WriteFile(identity2Path, []byte(id2.String()+"\n"), 0600)

	// Encrypt with key 1
	enc, err := NewEncryptor(recipient1Path)
	if err != nil {
		t.Fatalf("create encryptor: %v", err)
	}

	ciphertext, err := enc.Encrypt([]byte("secret message"))
	if err != nil {
		t.Fatalf("encrypt: %v", err)
	}

	// Try to decrypt with key 2 (should fail)
	dec, err := NewDecryptor(identity2Path)
	if err != nil {
		t.Fatalf("create decryptor: %v", err)
	}

	_, err = dec.Decrypt(ciphertext)
	if err == nil {
		t.Fatal("expected error decrypting with wrong key")
	}
}

func TestNewEncryptorInvalidPath(t *testing.T) {
	_, err := NewEncryptor("/nonexistent/path/to/recipients")
	if err == nil {
		t.Fatal("expected error with nonexistent recipients file")
	}
}

func TestNewDecryptorInvalidPath(t *testing.T) {
	_, err := NewDecryptor("/nonexistent/path/to/identity")
	if err == nil {
		t.Fatal("expected error with nonexistent identity file")
	}
}

func TestNewEncryptorEmptyPath(t *testing.T) {
	_, err := NewEncryptor("")
	if err == nil {
		t.Fatal("expected error with empty recipients path")
	}
}

func TestNewDecryptorEmptyPath(t *testing.T) {
	_, err := NewDecryptor("")
	if err == nil {
		t.Fatal("expected error with empty identity path")
	}
}

func TestStreamingEncryptDecrypt(t *testing.T) {
	// Generate test key pair
	tmpDir := t.TempDir()
	identityPath, recipientPath, err := GenerateTestKeyPair(tmpDir)
	if err != nil {
		t.Fatalf("generate test key pair: %v", err)
	}

	// Create encryptor and decryptor
	enc, err := NewEncryptor(recipientPath)
	if err != nil {
		t.Fatalf("create encryptor: %v", err)
	}

	dec, err := NewDecryptor(identityPath)
	if err != nil {
		t.Fatalf("create decryptor: %v", err)
	}

	// Test streaming encryption
	plaintext := []byte("streaming test data")
	var ciphertextBuf bytes.Buffer

	encWriter, err := enc.EncryptStream(&ciphertextBuf)
	if err != nil {
		t.Fatalf("create encrypt stream: %v", err)
	}

	if _, err := encWriter.Write(plaintext); err != nil {
		t.Fatalf("write to encrypt stream: %v", err)
	}

	if err := encWriter.Close(); err != nil {
		t.Fatalf("close encrypt stream: %v", err)
	}

	// Test streaming decryption
	decReader, err := dec.DecryptStream(&ciphertextBuf)
	if err != nil {
		t.Fatalf("create decrypt stream: %v", err)
	}

	decrypted := make([]byte, len(plaintext))
	n, err := decReader.Read(decrypted)
	if err != nil {
		t.Fatalf("read from decrypt stream: %v", err)
	}

	if !bytes.Equal(decrypted[:n], plaintext) {
		t.Fatalf("decrypted data doesn't match original\ngot:  %q\nwant: %q", decrypted[:n], plaintext)
	}
}

func TestConvenienceFunctions(t *testing.T) {
	// Generate test keys in memory
	identity, recipient, err := GenerateTestKeys()
	if err != nil {
		t.Fatalf("generate test keys: %v", err)
	}

	plaintext := []byte("test message")

	// Encrypt with convenience function
	ciphertext, err := EncryptBytes(plaintext, recipient)
	if err != nil {
		t.Fatalf("encrypt bytes: %v", err)
	}

	// Decrypt with convenience function
	decrypted, err := DecryptBytes(ciphertext, identity)
	if err != nil {
		t.Fatalf("decrypt bytes: %v", err)
	}

	if !bytes.Equal(decrypted, plaintext) {
		t.Fatalf("decrypted data doesn't match original\ngot:  %q\nwant: %q", decrypted, plaintext)
	}
}