"""
Tests for Grover's Algorithm Implementation
"""

import pytest
import math
from chronos.core.simulate.grovers_algorithm import (
    GroversAlgorithm,
    SearchOracle,
    SearchSpaceReducer,
    QuantumResourceEstimator,
    AmplificationVisualizer,
    ClassicalBenchmark,
    create_grovers_algorithm,
    SymmetricKeyType,
    OracleType,
    SearchPhase,
)


class TestSearchOracle:
    """Test oracle implementation."""

    def test_oracle_creation(self):
        """Test creating oracle."""
        target = b'secret_key_here1'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        assert oracle.target_key == target
        assert oracle.oracle_calls == 0

    def test_oracle_query_exact_match(self):
        """Test oracle query with exact match."""
        target = b'secret_key_here1'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        
        result = oracle.query(target)
        assert result is True
        assert oracle.oracle_calls == 1

    def test_oracle_query_no_match(self):
        """Test oracle query with no match."""
        target = b'secret_key_here1'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        
        result = oracle.query(b'wrong_key_here12')
        assert result is False
        assert oracle.oracle_calls == 1

    def test_oracle_call_tracking(self):
        """Test oracle call tracking."""
        target = b'secret_key_here1'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        
        for i in range(100):
            oracle.query(f'key_{i:04d}'.encode()[:16].ljust(16, b'0'))
        
        assert oracle.oracle_calls == 100

    def test_weak_encryption_oracle(self):
        """Test weak encryption pattern detection."""
        target = b'\x00' * 16
        oracle = SearchOracle(OracleType.WEAK_CIPHERTEXT, target)
        
        # All zeros should be detected
        assert oracle.query(b'\x00' * 16) is True
        
        # All ones should be detected
        assert oracle.query(b'\xff' * 16) is True


class TestGroversAlgorithm:
    """Test Grover's algorithm implementation."""

    def test_algorithm_creation(self):
        """Test creating Grover's algorithm instance."""
        algo = create_grovers_algorithm()
        assert algo is not None
        assert algo.algorithm_id == "grover-search"

    def test_algorithm_custom_id(self):
        """Test creating algorithm with custom ID."""
        algo = create_grovers_algorithm("custom-id-123")
        assert algo.algorithm_id == "custom-id-123"

    def test_calculate_iterations_128bit(self):
        """Test iteration calculation for 128-bit key."""
        algo = GroversAlgorithm()
        total = 2 ** 128
        iterations = algo._calculate_iterations(total, 1)
        
        # Should be approximately π/4 * √(2^128)
        expected = int(math.pi / 4 * math.sqrt(total))
        assert iterations > 0
        assert abs(iterations - expected) < expected * 0.1  # 10% tolerance

    def test_calculate_iterations_multiple_marked(self):
        """Test iteration calculation with multiple marked items."""
        algo = GroversAlgorithm()
        total = 2 ** 128
        marked = 1000
        
        iterations = algo._calculate_iterations(total, marked)
        expected = int(math.pi / 4 * math.sqrt(total / marked))
        
        assert iterations > 0
        assert abs(iterations - expected) < expected * 0.1

    def test_search_key_aes128(self):
        """Test searching for AES-128 key."""
        algo = create_grovers_algorithm()
        target_key = b'0123456789abcdef'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target_key)
        
        result = algo.search_key(
            SymmetricKeyType.AES_128,
            oracle,
            marked_items=1
        )
        
        assert result.success is True
        assert result.key_found == target_key
        assert result.verification_passed is True
        assert result.speedup_factor > 1.0

    def test_search_key_aes256(self):
        """Test searching for AES-256 key."""
        algo = create_grovers_algorithm()
        target_key = b'0123456789abcdef' * 2
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target_key)
        
        result = algo.search_key(
            SymmetricKeyType.AES_256,
            oracle,
            marked_items=1
        )
        
        assert result is not None
        assert result.iterations_needed > 0

    def test_search_multiple_marked_items(self):
        """Test search with multiple marked items."""
        algo = create_grovers_algorithm()
        target_key = b'secret_key_here1'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target_key)
        
        result = algo.search_key(
            SymmetricKeyType.AES_128,
            oracle,
            marked_items=10
        )
        
        assert result.iterations_needed > 0
        assert result.search_space_reduced > 0

    def test_classical_time_estimation(self):
        """Test classical time estimation."""
        algo = GroversAlgorithm()
        
        # 128-bit key
        time_128 = algo._classical_time(2**128, 1)
        assert time_128 > 0
        
        # 256-bit key should take longer
        time_256 = algo._classical_time(2**256, 1)
        assert time_256 > time_128

    def test_verify_key(self):
        """Test key verification."""
        algo = GroversAlgorithm()
        target = b'secret_key_here1'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        
        # Correct key
        assert algo._verify_key(target, oracle) is True
        
        # Wrong key
        assert algo._verify_key(b'wrong_key_here12', oracle) is False
        
        # None key
        assert algo._verify_key(None, oracle) is False

    def test_execution_log(self):
        """Test execution log creation."""
        algo = create_grovers_algorithm("test-log")
        target = b'0123456789abcdef'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        
        result = algo.search_key(
            SymmetricKeyType.AES_128,
            oracle,
            marked_items=1
        )
        
        log = algo.get_execution_log()
        assert log is not None
        assert log.algorithm_id == "test-log"
        assert log.total_oracle_calls > 0
        assert log.success is True


class TestSearchSpaceReducer:
    """Test search space reduction calculations."""

    def test_reduction_single_marked(self):
        """Test reduction factor with single marked item."""
        info = SearchSpaceReducer.calculate_reduction(128, 1)
        
        assert info.total_search_space == 2**128
        assert info.marked_items == 1
        assert info.reduction_factor == math.sqrt(2**128)
        assert info.grover_iterations > 0

    def test_reduction_multiple_marked(self):
        """Test reduction factor with multiple marked items."""
        info = SearchSpaceReducer.calculate_reduction(256, 1000)
        
        assert info.total_search_space == 2**256
        assert info.marked_items == 1000
        assert info.grover_iterations > 0
        assert info.marked_item_ratio == 1000 / (2**256)

    def test_speedup_analysis_128(self):
        """Test speedup analysis for 128-bit key."""
        speedup = SearchSpaceReducer.speedup_analysis(128)
        
        assert speedup["classical_average_attempts"] == 2**127
        assert speedup["quantum_iterations"] > 0
        assert speedup["speedup_factor"] == 2**64
        assert speedup["speedup_ratio"] > 1.0

    def test_speedup_analysis_256(self):
        """Test speedup analysis for 256-bit key."""
        speedup = SearchSpaceReducer.speedup_analysis(256)
        
        assert speedup["speedup_factor"] == 2**128
        assert speedup["reduction_percent"] > 0
        assert speedup["reduction_percent"] < 100


class TestQuantumResourceEstimator:
    """Test quantum resource estimation."""

    def test_estimate_aes128(self):
        """Test resource estimate for AES-128."""
        estimate = QuantumResourceEstimator.estimate_resources(
            SymmetricKeyType.AES_128,
            marked_items=1
        )
        
        assert estimate.key_size == 128
        assert estimate.logical_qubits_needed == 2 * 128 + 1
        assert estimate.physical_qubits_needed > estimate.logical_qubits_needed
        assert estimate.circuit_depth > 0
        assert estimate.oracle_queries > 0
        assert estimate.success_probability > 0.9

    def test_estimate_aes256(self):
        """Test resource estimate for AES-256."""
        estimate = QuantumResourceEstimator.estimate_resources(
            SymmetricKeyType.AES_256,
            marked_items=1
        )
        
        assert estimate.key_size == 256
        assert estimate.logical_qubits_needed == 2 * 256 + 1
        assert estimate.physical_qubits_needed > 0
        assert estimate.success_probability > 0.9

    def test_estimate_error_correction_overhead(self):
        """Test error correction overhead increases with key size."""
        estimate_128 = QuantumResourceEstimator.estimate_resources(
            SymmetricKeyType.AES_128
        )
        estimate_256 = QuantumResourceEstimator.estimate_resources(
            SymmetricKeyType.AES_256
        )
        
        # Larger key should have higher overhead percentage
        assert estimate_256.error_correction_overhead > 0
        assert estimate_128.error_correction_overhead > 0

    def test_resource_speedup_factor(self):
        """Test speedup factor in resource estimate."""
        estimate = QuantumResourceEstimator.estimate_resources(
            SymmetricKeyType.AES_128
        )
        
        assert estimate.speedup_factor > 1.0


class TestAmplificationVisualizer:
    """Test visualization functions."""

    def test_draw_amplitude_growth(self):
        """Test amplitude growth visualization."""
        viz = AmplificationVisualizer.draw_amplitude_growth(
            iterations=10,
            total_keys=2**128,
            marked_items=1
        )
        
        assert viz is not None
        assert len(viz) > 0
        assert "Grover's Algorithm" in viz
        assert "Round" in viz

    def test_draw_search_space(self):
        """Test search space visualization."""
        viz = AmplificationVisualizer.draw_search_space(
            key_size=128,
            marked_items=1
        )
        
        assert viz is not None
        assert len(viz) > 0
        assert "Search Space" in viz
        assert "Speedup" in viz

    def test_draw_oracle_phase_kickback(self):
        """Test oracle phase kickback visualization."""
        viz = AmplificationVisualizer.draw_oracle_phase_kickback()
        
        assert viz is not None
        assert len(viz) > 0
        assert "Phase Kickback" in viz
        assert "Oracle" in viz


class TestClassicalBenchmark:
    """Test classical benchmark calculations."""

    def test_benchmark_aes128(self):
        """Test classical benchmark for AES-128."""
        bench = ClassicalBenchmark.estimate_classical_time(
            SymmetricKeyType.AES_128,
            marked_items=1
        )
        
        assert bench.key_size == 128
        assert bench.classical_attempts_average > 0
        assert bench.classical_time_seconds > 0
        assert bench.quantum_time_seconds > 0
        assert bench.speedup_ratio > 0

    def test_benchmark_aes256(self):
        """Test classical benchmark for AES-256."""
        bench = ClassicalBenchmark.estimate_classical_time(
            SymmetricKeyType.AES_256,
            marked_items=1
        )
        
        assert bench.key_size == 256
        assert bench.quantum_advantage is True  # Should have advantage

    def test_benchmark_quantum_advantage(self):
        """Test quantum advantage flag."""
        bench_128 = ClassicalBenchmark.estimate_classical_time(
            SymmetricKeyType.AES_128
        )
        
        # Quantum should be faster (speedup > 1.1)
        assert bench_128.quantum_advantage or bench_128.speedup_ratio <= 1.1

    def test_benchmark_different_hash_rates(self):
        """Test benchmark with different hash rates."""
        bench_slow = ClassicalBenchmark.estimate_classical_time(
            SymmetricKeyType.AES_128,
            hash_rate_per_sec=100_000
        )
        
        bench_fast = ClassicalBenchmark.estimate_classical_time(
            SymmetricKeyType.AES_128,
            hash_rate_per_sec=1_000_000
        )
        
        # Faster hash rate should result in less classical time
        assert bench_slow.classical_time_seconds > bench_fast.classical_time_seconds


class TestEnumerations:
    """Test enumeration definitions."""

    def test_symmetric_key_types(self):
        """Test symmetric key type enumeration."""
        assert SymmetricKeyType.AES_128.value == 128
        assert SymmetricKeyType.AES_256.value == 256
        assert SymmetricKeyType.CHACHA20.value == 256

    def test_oracle_types(self):
        """Test oracle type enumeration."""
        assert OracleType.KNOWN_PLAINTEXT.value == "known_plaintext"
        assert OracleType.PATTERN_MATCHING.value == "pattern_matching"

    def test_search_phases(self):
        """Test search phase enumeration."""
        assert SearchPhase.INITIALIZATION.value == "initialization"
        assert SearchPhase.COMPLETED.value == "completed"
        assert SearchPhase.FAILED.value == "failed"


class TestIntegration:
    """Integration tests for complete workflows."""

    def test_complete_key_search_workflow(self):
        """Test complete key search workflow."""
        # Create algorithm
        algo = create_grovers_algorithm("integration-test")
        
        # Create oracle
        target = b'0123456789abcdef'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        
        # Perform search
        result = algo.search_key(
            SymmetricKeyType.AES_128,
            oracle,
            marked_items=1,
            simulate_quantum_noise=True
        )
        
        # Verify results
        assert result.success is True
        assert result.key_found == target
        assert result.speedup_factor > 0
        assert len(result.amplification_rounds) > 0

    def test_resource_estimation_workflow(self):
        """Test resource estimation workflow."""
        # Estimate resources
        estimate = QuantumResourceEstimator.estimate_resources(
            SymmetricKeyType.AES_256,
            marked_items=1
        )
        
        # Check resource estimates
        assert estimate.logical_qubits_needed > 0
        assert estimate.physical_qubits_needed >= estimate.logical_qubits_needed
        assert estimate.oracle_queries > 0

        # Compare with classical
        bench = ClassicalBenchmark.estimate_classical_time(
            SymmetricKeyType.AES_256
        )
        
        assert bench.speedup_ratio > 0

    def test_amplitude_analysis_workflow(self):
        """Test amplitude analysis workflow."""
        # Get search space info
        info = SearchSpaceReducer.calculate_reduction(128, 1)
        
        # Draw visualization
        viz = AmplificationVisualizer.draw_amplitude_growth(
            iterations=info.grover_iterations,
            total_keys=info.total_search_space,
            marked_items=info.marked_items
        )
        
        # Verify visualization
        assert viz is not None
        assert "amplitude" in viz.lower() or "round" in viz.lower()


class TestEdgeCases:
    """Test edge cases and error handling."""

    def test_search_single_item_space(self):
        """Test searching in single item space."""
        algo = GroversAlgorithm()
        
        # N=1, only one possible key
        iterations = algo._calculate_iterations(1, 1)
        assert iterations == 1

    def test_all_items_marked(self):
        """Test when all items are marked."""
        algo = GroversAlgorithm()
        
        # All 2^128 items are marked
        iterations = algo._calculate_iterations(2**128, 2**128)
        assert iterations == 1

    def test_oracle_with_none_key(self):
        """Test oracle behavior with None key."""
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, None)
        
        result = oracle.query(b'test')
        assert result is False

    def test_thread_safety(self):
        """Test thread safety of algorithm."""
        import threading
        
        algo = create_grovers_algorithm("thread-test")
        target = b'0123456789abcdef'
        oracle = SearchOracle(OracleType.KNOWN_PLAINTEXT, target)
        
        results = []
        
        def search():
            result = algo.search_key(
                SymmetricKeyType.AES_128,
                oracle,
                marked_items=1
            )
            results.append(result)
        
        threads = [threading.Thread(target=search) for _ in range(3)]
        
        for t in threads:
            t.start()
        
        for t in threads:
            t.join()
        
        assert len(results) == 3
        assert all(r.success for r in results)


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
