import os
import unittest

from src.encryptor import Encryptor


class TestEncryptor(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        """Set up shared resources for all tests."""
        cls.test_key_path = "test_encryption.key"
        cls.test_invalid_key_path = "test_invalid_encryption.key"
        cls.test_file_path = "test_file.txt"
        cls.encrypted_file_path = "test_file_encrypted.txt"
        cls.sample_text = "This is a test message.\n" * 100

        # Create Encryptor instance with test key path
        cls.encryptor = Encryptor(cls.test_key_path)

    def setUp(self):
        """Prepare test environment before each test."""
        # Ensure test file exists
        with open(self.test_file_path, "w") as f:
            f.write(self.sample_text)

    def tearDown(self):
        """Clean up test files after each test."""
        for file in [self.test_file_path, self.test_key_path, self.test_invalid_key_path]:
            if os.path.exists(file):
                os.remove(file)

    def test_key_persistence(self):
        """Test if encryption key is stored and reused correctly."""
        encryptor1 = Encryptor(self.test_key_path)
        key1 = encryptor1.encryption_key

        encryptor2 = Encryptor(self.test_key_path)
        key2 = encryptor2.encryption_key

        self.assertEqual(key1, key2, "Keys should persist across instances.")

    def test_file_encryption_and_decryption(self):
        """Test file encryption and decryption process."""
        self.encryptor.encrypt_file(self.test_file_path)

        # Ensure file content is changed after encryption
        with open(self.test_file_path, "rb") as f:
            encrypted_data = f.read()
        self.assertNotEqual(encrypted_data, self.sample_text.encode(), "File should be encrypted.")

        self.encryptor.decrypt_file(self.test_file_path)

        # Ensure file content is restored correctly
        with open(self.test_file_path, "r") as f:
            decrypted_data = f.read()
        self.assertEqual(decrypted_data, self.sample_text, "Decrypted content should match original text.")

    def test_text_encryption_and_decryption(self):
        """Test direct text encryption and decryption methods."""
        encrypted_text = self.encryptor.encrypt_text(self.sample_text)
        decrypted_text = self.encryptor.decrypt_text(encrypted_text)

        self.assertNotEqual(encrypted_text, self.sample_text, "Encrypted text should not match plain text.")
        self.assertEqual(decrypted_text, self.sample_text, "Decrypted text should match the original text.")

    def test_invalid_decryption(self):
        """Test decryption with an invalid key to trigger an exception."""
        invalid_encryptor = Encryptor(self.test_invalid_key_path)
        encrypted_text = invalid_encryptor.encrypt_text(self.sample_text)

        with self.assertRaises(Exception):
            self.encryptor.decrypt_text(encrypted_text)  # Should fail due to key mismatch
