"""Tests for the job payload encryption module."""

from __future__ import annotations

import base64
import json
import os
import time
from unittest.mock import patch

import pytest

from oclawma.encryption import (
    DecryptionError,
    EncryptionManager,
    EncryptionMetadata,
    KeyNotFoundError,
    generate_key,
)


class TestKeyGeneration:
    """Test key generation utilities."""

    def test_generate_key_returns_bytes(self):
        """Test that generate_key returns 32 bytes."""
        key = generate_key()
        assert isinstance(key, bytes)
        assert len(key) == 32  # AES-256 requires 32 bytes

    def test_generate_key_is_random(self):
        """Test that generate_key produces different keys."""
        key1 = generate_key()
        key2 = generate_key()
        assert key1 != key2


class TestEncryptionMetadata:
    """Test EncryptionMetadata model."""

    def test_default_creation(self):
        """Test creating metadata with defaults."""
        meta = EncryptionMetadata()
        assert meta.encrypted is True
        assert meta.key_version == "1"
        assert meta.algorithm == "AES-256-GCM"

    def test_custom_creation(self):
        """Test creating metadata with custom values."""
        meta = EncryptionMetadata(
            encrypted=True,
            key_version="v2",
            algorithm="AES-256-GCM",
        )
        assert meta.key_version == "v2"

    def test_to_from_dict(self):
        """Test serialization round-trip."""
        original = EncryptionMetadata(
            encrypted=True,
            key_version="v3",
            algorithm="AES-256-GCM",
        )
        data = original.to_dict()
        restored = EncryptionMetadata.from_dict(data)
        assert restored.encrypted == original.encrypted
        assert restored.key_version == original.key_version
        assert restored.algorithm == original.algorithm


class TestEncryptionManagerBasic:
    """Test basic EncryptionManager functionality."""

    def test_init_with_generated_key(self):
        """Test initialization with auto-generated key."""
        manager = EncryptionManager()
        assert manager.current_key_version == "1"
        assert "1" in manager.keys
        assert len(manager.keys["1"]) == 32

    def test_init_with_provided_key(self):
        """Test initialization with provided key."""
        key = generate_key()
        manager = EncryptionManager(key=key)
        assert manager.keys["1"] == key

    def test_init_with_key_from_env(self):
        """Test initialization with key from environment variable."""
        key = generate_key()
        key_b64 = base64.b64encode(key).decode()
        with patch.dict(os.environ, {"OCLAWMA_ENCRYPTION_KEY": key_b64}):
            manager = EncryptionManager()
            assert manager.keys["1"] == key

    def test_init_with_key_from_file(self, tmp_path):
        """Test initialization with key from file."""
        key = generate_key()
        key_b64 = base64.b64encode(key).decode()
        key_file = tmp_path / "encryption.key"
        key_file.write_text(key_b64)
        manager = EncryptionManager(key_file=str(key_file))
        assert manager.keys["1"] == key

    def test_init_invalid_key_size(self):
        """Test that invalid key size raises error."""
        with pytest.raises(ValueError, match="32 bytes"):
            EncryptionManager(key=b"too_short")

    def test_encrypt_decrypt_roundtrip(self):
        """Test basic encrypt/decrypt roundtrip."""
        manager = EncryptionManager()
        plaintext = b"Hello, World!"
        ciphertext = manager.encrypt(plaintext)
        assert ciphertext != plaintext
        assert isinstance(ciphertext, bytes)
        decrypted = manager.decrypt(ciphertext)
        assert decrypted == plaintext

    def test_encrypt_different_results(self):
        """Test that encrypting same data twice gives different ciphertext."""
        manager = EncryptionManager()
        plaintext = b"Hello, World!"
        ct1 = manager.encrypt(plaintext)
        ct2 = manager.encrypt(plaintext)
        assert ct1 != ct2  # Due to random nonce

    def test_encrypt_empty_data(self):
        """Test encrypting empty data."""
        manager = EncryptionManager()
        ciphertext = manager.encrypt(b"")
        decrypted = manager.decrypt(ciphertext)
        assert decrypted == b""

    def test_encrypt_binary_data(self):
        """Test encrypting binary data."""
        manager = EncryptionManager()
        plaintext = bytes(range(256))
        ciphertext = manager.encrypt(plaintext)
        decrypted = manager.decrypt(ciphertext)
        assert decrypted == plaintext


class TestEncryptionManagerKeyRotation:
    """Test key rotation functionality."""

    def test_add_new_key_version(self):
        """Test adding a new key version."""
        manager = EncryptionManager()
        key = generate_key()
        manager.add_key("v2", key)
        assert "v2" in manager.keys
        assert manager.keys["v2"] == key

    def test_add_duplicate_key_version_raises(self):
        """Test that adding duplicate key version raises error."""
        manager = EncryptionManager()
        key = generate_key()
        with pytest.raises(ValueError, match="already exists"):
            manager.add_key("1", key)

    def test_encrypt_with_specific_key_version(self):
        """Test encrypting with a specific key version."""
        manager = EncryptionManager()
        key_v2 = generate_key()
        manager.add_key("v2", key_v2)

        plaintext = b"secret data"
        ciphertext = manager.encrypt(plaintext, key_version="v2")

        # Parse ciphertext to verify key version
        data = json.loads(ciphertext.decode())
        assert data["key_version"] == "v2"

    def test_decrypt_with_old_key_version(self):
        """Test decrypting data encrypted with old key version."""
        manager = EncryptionManager()

        # Encrypt with v1
        plaintext = b"secret data"
        ciphertext_v1 = manager.encrypt(plaintext)

        # Add v2 key
        key_v2 = generate_key()
        manager.add_key("v2", key_v2)
        manager.current_key_version = "v2"

        # Should still decrypt v1 data
        decrypted = manager.decrypt(ciphertext_v1)
        assert decrypted == plaintext

    def test_current_key_version_affects_new_encryption(self):
        """Test that changing current_key_version affects new encryptions."""
        manager = EncryptionManager()
        key_v2 = generate_key()
        manager.add_key("v2", key_v2)
        manager.current_key_version = "v2"

        ciphertext = manager.encrypt(b"data")
        data = json.loads(ciphertext.decode())
        assert data["key_version"] == "v2"


class TestEncryptionManagerErrors:
    """Test error handling."""

    def test_decrypt_with_missing_key_version(self):
        """Test decrypting when key version doesn't exist."""
        manager = EncryptionManager()

        # Create ciphertext with non-existent key version
        fake_data = {
            "key_version": "nonexistent",
            "nonce": base64.b64encode(os.urandom(12)).decode(),
            "ciphertext": base64.b64encode(b"data").decode(),
            "tag": base64.b64encode(os.urandom(16)).decode(),
        }
        ciphertext = json.dumps(fake_data).encode()

        with pytest.raises(KeyNotFoundError, match="nonexistent"):
            manager.decrypt(ciphertext)

    def test_decrypt_corrupted_data(self):
        """Test decrypting corrupted data."""
        manager = EncryptionManager()
        plaintext = b"Hello, World!"
        ciphertext = manager.encrypt(plaintext)

        # Corrupt the ciphertext
        corrupted = ciphertext[:-5] + b"XXXXX"

        with pytest.raises(DecryptionError):
            manager.decrypt(corrupted)

    def test_decrypt_invalid_json(self):
        """Test decrypting invalid JSON."""
        manager = EncryptionManager()
        with pytest.raises(DecryptionError):
            manager.decrypt(b"not valid json")

    def test_decrypt_tampered_ciphertext(self):
        """Test detecting tampered ciphertext."""
        manager = EncryptionManager()
        plaintext = b"Hello, World!"
        ciphertext = manager.encrypt(plaintext)

        # Parse and modify ciphertext
        data = json.loads(ciphertext.decode())
        ct_bytes = base64.b64decode(data["ciphertext"])
        ct_bytes = bytes([b ^ 0xFF for b in ct_bytes[:5]]) + ct_bytes[5:]
        data["ciphertext"] = base64.b64encode(ct_bytes).decode()
        tampered = json.dumps(data).encode()

        with pytest.raises(DecryptionError):
            manager.decrypt(tampered)

    def test_decrypt_tampered_tag(self):
        """Test detecting tampered authentication tag."""
        manager = EncryptionManager()
        plaintext = b"Hello, World!"
        ciphertext = manager.encrypt(plaintext)

        # Parse and modify tag
        data = json.loads(ciphertext.decode())
        tag_bytes = base64.b64decode(data["tag"])
        tag_bytes = bytes([b ^ 0xFF for b in tag_bytes[:5]]) + tag_bytes[5:]
        data["tag"] = base64.b64encode(tag_bytes).decode()
        tampered = json.dumps(data).encode()

        with pytest.raises(DecryptionError):
            manager.decrypt(tampered)


class TestEncryptionManagerStringAPI:
    """Test string-based encrypt/decrypt API."""

    def test_encrypt_string(self):
        """Test encrypting string."""
        manager = EncryptionManager()
        plaintext = "Hello, Unicode! 🎉"
        ciphertext = manager.encrypt_string(plaintext)
        assert isinstance(ciphertext, str)

        decrypted = manager.decrypt_string(ciphertext)
        assert decrypted == plaintext

    def test_encrypt_dict(self):
        """Test encrypting dictionary."""
        manager = EncryptionManager()
        plaintext = {"key": "value", "number": 42, "nested": {"a": 1}}
        ciphertext = manager.encrypt_dict(plaintext)
        assert isinstance(ciphertext, str)

        decrypted = manager.decrypt_dict(ciphertext)
        assert decrypted == plaintext

    def test_decrypt_dict_invalid_json(self):
        """Test decrypting dict when plaintext isn't valid JSON."""
        manager = EncryptionManager()
        ciphertext = manager.encrypt_string("not valid json")
        with pytest.raises(json.JSONDecodeError):
            manager.decrypt_dict(ciphertext)


class TestEncryptionManagerPerformance:
    """Test performance characteristics."""

    def test_encrypt_performance_small_data(self):
        """Test encryption performance for small data."""
        manager = EncryptionManager()
        plaintext = b"x" * 1024  # 1KB

        start = time.time()
        for _ in range(100):
            ciphertext = manager.encrypt(plaintext)
            manager.decrypt(ciphertext)
        elapsed = time.time() - start

        # Should complete 100 encrypt+decrypt cycles in under 1 second
        assert elapsed < 1.0

    def test_encrypt_performance_large_data(self):
        """Test encryption performance for larger data."""
        manager = EncryptionManager()
        plaintext = b"x" * 1024 * 1024  # 1MB

        start = time.time()
        ciphertext = manager.encrypt(plaintext)
        elapsed_encrypt = time.time() - start

        start = time.time()
        manager.decrypt(ciphertext)
        elapsed_decrypt = time.time() - start

        # Should complete 1MB encrypt in under 0.5 seconds
        assert elapsed_encrypt < 0.5
        assert elapsed_decrypt < 0.5

    def test_ciphertext_overhead(self):
        """Test that ciphertext overhead is reasonable."""
        manager = EncryptionManager()
        plaintext = b"x" * 1000
        ciphertext = manager.encrypt(plaintext)

        # Ciphertext includes:
        # - Base64 encoding of plaintext (~33% overhead)
        # - 12 byte nonce (base64: 16 bytes)
        # - 16 byte tag (base64: 24 bytes)
        # - JSON structure (~100 bytes)
        # Total overhead: ~33% of plaintext + ~140 bytes fixed
        overhead = len(ciphertext) - len(plaintext)
        # For 1000 bytes: ~330 + 140 = 470 bytes overhead
        assert overhead < 600


class TestEncryptionManagerConfiguration:
    """Test configuration options."""

    def test_encrypted_types_default(self):
        """Test default encrypted_types is empty set."""
        manager = EncryptionManager()
        assert manager.encrypted_types == set()
        assert manager.should_encrypt("any_type") is False

    def test_encrypted_types_custom(self):
        """Test custom encrypted_types."""
        manager = EncryptionManager(encrypted_types={"sensitive", "private"})
        assert manager.should_encrypt("sensitive") is True
        assert manager.should_encrypt("private") is True
        assert manager.should_encrypt("normal") is False

    def test_add_encrypted_type(self):
        """Test adding encrypted type."""
        manager = EncryptionManager()
        manager.add_encrypted_type("new_type")
        assert manager.should_encrypt("new_type") is True

    def test_remove_encrypted_type(self):
        """Test removing encrypted type."""
        manager = EncryptionManager(encrypted_types={"type1", "type2"})
        manager.remove_encrypted_type("type1")
        assert manager.should_encrypt("type1") is False
        assert manager.should_encrypt("type2") is True


class TestEncryptionManagerKeyPersistence:
    """Test key persistence."""

    def test_save_and_load_keys(self, tmp_path):
        """Test saving and loading keys."""
        manager = EncryptionManager()
        key_v2 = generate_key()
        manager.add_key("v2", key_v2)

        key_file = tmp_path / "keys.json"
        manager.save_keys(str(key_file))

        # Load in new manager
        new_manager = EncryptionManager(key_file=str(key_file))
        assert "1" in new_manager.keys
        assert "v2" in new_manager.keys
        assert new_manager.keys["v2"] == key_v2

    def test_generate_key_file_if_not_exists(self, tmp_path):
        """Test that key file is generated if it doesn't exist."""
        key_file = tmp_path / "new_keys.json"
        manager = EncryptionManager(key_file=str(key_file))
        assert key_file.exists()

        # Verify we can load from it
        new_manager = EncryptionManager(key_file=str(key_file))
        assert new_manager.keys["1"] == manager.keys["1"]

    def test_base64_key_encoding(self, tmp_path):
        """Test that keys are properly base64 encoded in file."""
        key = generate_key()
        key_file = tmp_path / "key.b64"
        key_file.write_text(base64.b64encode(key).decode())

        manager = EncryptionManager(key_file=str(key_file))
        assert manager.keys["1"] == key


class TestEncryptionManagerContextManager:
    """Test context manager support."""

    def test_context_manager(self):
        """Test using EncryptionManager as context manager."""
        with EncryptionManager() as manager:
            assert "1" in manager.keys
            plaintext = b"test"
            ciphertext = manager.encrypt(plaintext)
            assert manager.decrypt(ciphertext) == plaintext
