"""
Tests for KeyManager and crypto-shredding functionality.
"""

import pytest
from unittest.mock import patch

from cryptography.fernet import Fernet
from atomsql import Database

from atomingest.security.kms import KeyManager, KeyStoreModel


class TestKeyStoreModel:
    """Tests for KeyStoreModel."""

    @pytest.fixture
    def db(self):
        """Create an in-memory database for testing."""
        database = Database("sqlite:///:memory:")
        database.register(KeyStoreModel)
        yield database
        database.close()

    def test_keystore_model_has_required_fields(self, db):
        """Test that KeyStoreModel has all required fields."""
        assert hasattr(KeyStoreModel, "key_identifier")
        assert hasattr(KeyStoreModel, "wrapped_key")
        assert hasattr(KeyStoreModel, "created_at")

    def test_keystore_model_can_be_created(self, db):
        """Test that KeyStoreModel instances can be created."""
        keystore = KeyStoreModel(
            key_identifier="test_key_1",
            wrapped_key="encrypted_key_data",
        )
        keystore.save()

        assert keystore.id is not None
        assert keystore.key_identifier == "test_key_1"

    def test_keystore_model_can_be_queried(self, db):
        """Test that KeyStoreModel records can be queried."""
        keystore = KeyStoreModel(
            key_identifier="test_key_2",
            wrapped_key="encrypted_key_data_2",
        )
        keystore.save()

        # Query by key_identifier
        results = list(KeyStoreModel.objects().filter(key_identifier="test_key_2"))
        assert len(results) == 1
        assert results[0].wrapped_key == "encrypted_key_data_2"


class TestKeyManagerMasterCipher:
    """Tests for KeyManager master cipher functionality."""

    @pytest.fixture
    def master_key(self):
        """Generate a valid Fernet master key."""
        return Fernet.generate_key().decode("utf-8")

    @pytest.fixture
    def set_master_key(self, master_key, monkeypatch):
        """Set the master key in environment variables."""
        monkeypatch.setenv("ATOMINGEST_MASTER_KEY", master_key)
        # Reset the cached cipher
        KeyManager._master_cipher = None
        yield master_key
        KeyManager._master_cipher = None

    def test_missing_master_key_raises_error(self, monkeypatch):
        """Test that missing master key raises ValueError."""
        monkeypatch.delenv("ATOMINGEST_MASTER_KEY", raising=False)
        monkeypatch.delenv("ATOMINGEST_ENCRYPTION_KEY", raising=False)
        KeyManager._master_cipher = None

        with pytest.raises(ValueError) as exc_info:
            KeyManager.get_master_cipher()

        assert "ATOMINGEST_MASTER_KEY environment variable is not set" in str(
            exc_info.value
        )

    def test_invalid_master_key_raises_error(self, monkeypatch):
        """Test that invalid master key raises ValueError."""
        monkeypatch.setenv("ATOMINGEST_MASTER_KEY", "invalid_key_123")
        KeyManager._master_cipher = None

        with pytest.raises(ValueError) as exc_info:
            KeyManager.get_master_cipher()

        assert "Invalid Master Key" in str(exc_info.value)

    def test_valid_master_key_returns_cipher(self, set_master_key):
        """Test that valid master key returns a Fernet cipher."""
        cipher = KeyManager.get_master_cipher()

        assert isinstance(cipher, Fernet)

    def test_master_cipher_is_cached(self, set_master_key):
        """Test that master cipher is cached after first retrieval."""
        cipher1 = KeyManager.get_master_cipher()
        cipher2 = KeyManager.get_master_cipher()

        # Should return the same cached instance
        assert cipher1 is cipher2

    def test_fallback_to_legacy_encryption_key(self, monkeypatch):
        """Test fallback to ATOMINGEST_ENCRYPTION_KEY for backward compatibility."""
        legacy_key = Fernet.generate_key().decode("utf-8")
        monkeypatch.delenv("ATOMINGEST_MASTER_KEY", raising=False)
        monkeypatch.setenv("ATOMINGEST_ENCRYPTION_KEY", legacy_key)
        KeyManager._master_cipher = None

        cipher = KeyManager.get_master_cipher()

        assert isinstance(cipher, Fernet)


class TestKeyManagerKeyOperations:
    """Tests for KeyManager key creation, retrieval, and shredding."""

    @pytest.fixture
    def master_key(self):
        """Generate a valid Fernet master key."""
        return Fernet.generate_key().decode("utf-8")

    @pytest.fixture
    def db(self, master_key, monkeypatch):
        """Create an in-memory database with master key set."""
        monkeypatch.setenv("ATOMINGEST_MASTER_KEY", master_key)
        KeyManager._master_cipher = None

        database = Database("sqlite:///:memory:")
        database.register(KeyStoreModel)
        yield database
        database.close()
        KeyManager._master_cipher = None

    def test_create_key_generates_new_key(self, db):
        """Test that create_key generates and stores a new key."""
        key_id = "user_123"

        data_key = KeyManager.create_key(key_id)

        assert data_key is not None
        assert isinstance(data_key, bytes)
        assert len(data_key) > 0

        # Verify it's stored in database
        keystore_entry = list(KeyStoreModel.objects().filter(key_identifier=key_id))
        assert len(keystore_entry) == 1

    def test_get_key_retrieves_existing_key(self, db):
        """Test that get_key retrieves an existing key."""
        key_id = "user_456"

        # Create a key first
        original_key = KeyManager.create_key(key_id)

        # Retrieve the same key
        retrieved_key = KeyManager.get_key(key_id)

        assert retrieved_key == original_key

    def test_get_key_creates_if_not_exists(self, db):
        """Test that get_key creates a new key if it doesn't exist."""
        key_id = "user_789"

        # Key doesn't exist yet
        data_key = KeyManager.get_key(key_id)

        assert data_key is not None
        assert isinstance(data_key, bytes)

        # Verify it was created in database
        keystore_entries = list(KeyStoreModel.objects().filter(key_identifier=key_id))
        assert len(keystore_entries) == 1

    def test_wrapped_key_is_encrypted(self, db):
        """Test that wrapped keys are encrypted in the database."""
        key_id = "user_encrypted"

        data_key = KeyManager.create_key(key_id)

        # Get the wrapped key from database
        keystore_entry = list(KeyStoreModel.objects().filter(key_identifier=key_id))[0]

        # Wrapped key should not match the data key
        assert keystore_entry.wrapped_key != data_key.decode("utf-8")

        # Wrapped key should be decryptable with master cipher
        master = KeyManager.get_master_cipher()
        decrypted = master.decrypt(keystore_entry.wrapped_key.encode())
        assert decrypted == data_key

    def test_shred_key_deletes_key(self, db):
        """Test that shred_key permanently deletes a key."""
        key_id = "user_to_delete"

        # Create a key
        KeyManager.create_key(key_id)

        # Verify it exists
        before_count = len(list(KeyStoreModel.objects().filter(key_identifier=key_id)))
        assert before_count == 1

        # Shred the key
        result = KeyManager.shred_key(key_id)

        assert result is True

        # Verify it's deleted
        after_count = len(list(KeyStoreModel.objects().filter(key_identifier=key_id)))
        assert after_count == 0

    def test_shred_nonexistent_key_returns_false(self, db):
        """Test that shredding a non-existent key returns False."""
        key_id = "nonexistent_key"

        result = KeyManager.shred_key(key_id)

        assert result is False

    def test_multiple_keys_can_coexist(self, db):
        """Test that multiple keys can be stored independently."""
        key_ids = ["user_1", "user_2", "user_3"]

        keys = {}
        for key_id in key_ids:
            keys[key_id] = KeyManager.create_key(key_id)

        # Retrieve all keys
        for key_id in key_ids:
            retrieved_key = KeyManager.get_key(key_id)
            assert retrieved_key == keys[key_id]

    def test_shredding_one_key_doesnt_affect_others(self, db):
        """Test that shredding one key doesn't affect other keys."""
        key_id_1 = "user_keep"
        key_id_2 = "user_delete"

        key_1 = KeyManager.create_key(key_id_1)
        key_2 = KeyManager.create_key(key_id_2)

        # Shred key_2
        KeyManager.shred_key(key_id_2)

        # key_1 should still be retrievable
        retrieved_key_1 = KeyManager.get_key(key_id_1)
        assert retrieved_key_1 == key_1

        # key_2 should create a new key (since it was shredded)
        new_key_2 = KeyManager.get_key(key_id_2)
        assert new_key_2 != key_2


class TestKeyManagerIntegration:
    """Integration tests for KeyManager with encryption."""

    @pytest.fixture
    def master_key(self):
        """Generate a valid Fernet master key."""
        return Fernet.generate_key().decode("utf-8")

    @pytest.fixture
    def db(self, master_key, monkeypatch):
        """Create an in-memory database with master key set."""
        monkeypatch.setenv("ATOMINGEST_MASTER_KEY", master_key)
        KeyManager._master_cipher = None

        database = Database("sqlite:///:memory:")
        database.register(KeyStoreModel)
        yield database
        database.close()
        KeyManager._master_cipher = None

    def test_key_roundtrip_encryption_decryption(self, db):
        """Test complete roundtrip of key creation and encryption/decryption."""
        key_id = "roundtrip_test"

        # Create a key
        data_key = KeyManager.create_key(key_id)

        # Create a cipher from this key
        cipher = Fernet(data_key)

        # Encrypt some data
        plaintext = "sensitive information"
        encrypted = cipher.encrypt(plaintext.encode("utf-8"))

        # Retrieve the key
        retrieved_key = KeyManager.get_key(key_id)
        retrieved_cipher = Fernet(retrieved_key)

        # Decrypt the data
        decrypted = retrieved_cipher.decrypt(encrypted).decode("utf-8")

        assert decrypted == plaintext

    def test_crypto_shredding_makes_data_unreadable(self, db):
        """Test that crypto-shredding makes encrypted data unreadable."""
        key_id = "shred_test"

        # Create a key and encrypt data
        data_key = KeyManager.create_key(key_id)
        cipher = Fernet(data_key)
        plaintext = "top secret data"
        encrypted = cipher.encrypt(plaintext.encode("utf-8"))

        # Shred the key
        KeyManager.shred_key(key_id)

        # Try to decrypt with a newly created key (different key)
        new_key = KeyManager.get_key(key_id)  # Creates a new key
        new_cipher = Fernet(new_key)

        # Decryption should fail because it's a different key
        with pytest.raises(Exception):
            new_cipher.decrypt(encrypted)

    def test_different_key_ids_produce_different_keys(self, db):
        """Test that different key identifiers produce different data keys."""
        key_1 = KeyManager.get_key("tenant_1")
        key_2 = KeyManager.get_key("tenant_2")

        assert key_1 != key_2

    def test_same_key_id_returns_same_key(self, db):
        """Test that the same key identifier returns the same key."""
        key_id = "consistent_key"

        key_1 = KeyManager.get_key(key_id)
        key_2 = KeyManager.get_key(key_id)

        assert key_1 == key_2

    def test_key_recreation_after_shredding(self, db):
        """Test that a key can be recreated after shredding."""
        key_id = "recreate_test"

        # Create and shred
        original_key = KeyManager.create_key(key_id)
        KeyManager.shred_key(key_id)

        # Create again
        new_key = KeyManager.create_key(key_id)

        # Should be a different key
        assert new_key != original_key

    def test_concurrent_key_access(self, db):
        """Test that the same key can be accessed multiple times concurrently."""
        key_id = "concurrent_test"

        # Create initial key
        KeyManager.create_key(key_id)

        # Simulate multiple concurrent accesses
        keys = [KeyManager.get_key(key_id) for _ in range(10)]

        # All should be the same key
        assert all(k == keys[0] for k in keys)


class TestKeyManagerErrorHandling:
    """Tests for KeyManager error handling."""

    @pytest.fixture
    def master_key(self):
        """Generate a valid Fernet master key."""
        return Fernet.generate_key().decode("utf-8")

    @pytest.fixture
    def db(self, master_key, monkeypatch):
        """Create an in-memory database with master key set."""
        monkeypatch.setenv("ATOMINGEST_MASTER_KEY", master_key)
        KeyManager._master_cipher = None

        database = Database("sqlite:///:memory:")
        database.register(KeyStoreModel)
        yield database
        database.close()
        KeyManager._master_cipher = None

    def test_create_key_with_database_error(self, db):
        """Test error handling when database fails during key creation."""
        key_id = "db_error_test"

        # Mock KeyStoreModel.create to raise an exception
        with patch.object(
            KeyStoreModel, "create", side_effect=Exception("Database error")
        ):
            with pytest.raises(Exception) as exc_info:
                KeyManager.create_key(key_id)

            assert "Database error" in str(exc_info.value)

    def test_shred_key_with_database_error(self, db):
        """Test error handling when database fails during key shredding."""
        key_id = "shred_error_test"

        # Create a key first
        KeyManager.create_key(key_id)

        # Mock objects().filter().first() to raise an exception
        with patch.object(
            KeyStoreModel, "objects", side_effect=Exception("Database error")
        ):
            result = KeyManager.shred_key(key_id)

            # Should return False on error
            assert result is False

    def test_get_key_handles_corrupted_wrapped_key(self, db):
        """Test handling of corrupted wrapped keys in database."""
        key_id = "corrupted_test"

        # Manually create a keystore entry with corrupted data
        keystore = KeyStoreModel(
            key_identifier=key_id, wrapped_key="corrupted_invalid_data"
        )
        keystore.save()

        # Attempting to get this key should fail during decryption
        with pytest.raises(Exception):
            KeyManager.get_key(key_id)

    def test_master_cipher_caching_across_calls(self, db):
        """Test that master cipher is properly cached across multiple calls."""
        # First call
        cipher1 = KeyManager.get_master_cipher()

        # Second call
        cipher2 = KeyManager.get_master_cipher()

        # Should be the same instance
        assert cipher1 is cipher2

        # Create a key to verify cipher works
        key = KeyManager.create_key("cache_test")
        assert key is not None
