"""
Tests for Post-Quantum Cryptography Migration Engine

Comprehensive test suite for PQC migration system including:
- Kyber key encapsulation tests
- Dilithium signature tests
- Hybrid mode tests
- Migration planning and execution
- Performance benchmarking
- Rollback mechanisms
"""

import pytest
import json
from pathlib import Path
from datetime import datetime
from chronos.core.defend.pqc_migrator import (
    PQCAlgorithm,
    ClassicalAlgorithm,
    MigrationMode,
    MigrationStatus,
    PerformanceMetric,
    KeyPair,
    PerformanceBenchmark,
    MigrationSnapshot,
    MigrationPlan,
    KyberMigrator,
    DilithiumMigrator,
    HybridMigrator,
    MigrationRollback,
    PQCMigrator,
    analyze_migration_feasibility,
)


class TestPQCEnums:
    """Test PQC algorithm and status enumerations."""
    
    def test_pqc_algorithms_defined(self):
        """Test that PQC algorithms are properly defined."""
        assert PQCAlgorithm.KYBER_512.value == "Kyber512"
        assert PQCAlgorithm.KYBER_1024.value == "Kyber1024"
        assert PQCAlgorithm.DILITHIUM3.value == "Dilithium3"
        assert len(list(PQCAlgorithm)) >= 7
    
    def test_classical_algorithms_defined(self):
        """Test classical algorithm definitions."""
        assert ClassicalAlgorithm.RSA_2048.value == "RSA-2048"
        assert ClassicalAlgorithm.ECDSA_P256.value == "ECDSA-P256"
        assert len(list(ClassicalAlgorithm)) >= 6
    
    def test_migration_modes(self):
        """Test migration mode options."""
        assert MigrationMode.FULL_REPLACEMENT.value == "full_replacement"
        assert MigrationMode.HYBRID.value == "hybrid"
        assert MigrationMode.DUAL_ALGORITHM.value == "dual_algorithm"
    
    def test_migration_status_values(self):
        """Test migration status values."""
        assert MigrationStatus.PLANNED.value == "planned"
        assert MigrationStatus.COMPLETED.value == "completed"
        assert MigrationStatus.ROLLED_BACK.value == "rolled_back"


class TestMigrationPlan:
    """Test migration plan creation and management."""
    
    def test_migration_plan_creation(self):
        """Test creating a migration plan."""
        plan = MigrationPlan(
            source_algorithm=ClassicalAlgorithm.RSA_2048,
            target_algorithm=PQCAlgorithm.KYBER_1024,
        )
        
        assert plan.id is not None
        assert plan.source_algorithm == ClassicalAlgorithm.RSA_2048
        assert plan.target_algorithm == PQCAlgorithm.KYBER_1024
        assert plan.status == MigrationStatus.PLANNED
        assert plan.created_at is not None
    
    def test_migration_plan_to_dict(self):
        """Test converting migration plan to dictionary."""
        plan = MigrationPlan(
            source_algorithm=ClassicalAlgorithm.RSA_2048,
            target_algorithm=PQCAlgorithm.KYBER_1024,
        )
        
        plan_dict = plan.to_dict()
        
        assert 'id' in plan_dict
        assert plan_dict['source_algorithm'] == 'RSA-2048'
        assert plan_dict['target_algorithm'] == 'Kyber1024'
        assert plan_dict['status'] == 'planned'
    
    def test_migration_plan_duration(self):
        """Test migration plan duration calculation."""
        plan = MigrationPlan(
            source_algorithm=ClassicalAlgorithm.RSA_2048,
            target_algorithm=PQCAlgorithm.KYBER_1024,
        )
        
        # No duration without start/end times
        assert plan.duration_seconds is None
        
        # Add start and end times
        plan.started_at = datetime.now()
        plan.completed_at = datetime.now()
        
        assert plan.duration_seconds >= 0


class TestMigrationSnapshot:
    """Test migration snapshots for rollback."""
    
    def test_snapshot_creation(self, tmp_path):
        """Test creating a migration snapshot."""
        snapshot = MigrationSnapshot(
            metadata={'source': 'RSA-2048', 'target': 'Kyber1024'}
        )
        
        assert snapshot.timestamp is not None
        assert snapshot.metadata['source'] == 'RSA-2048'
    
    def test_snapshot_save_and_load(self, tmp_path):
        """Test saving and loading snapshots."""
        snapshot_path = tmp_path / "snapshot.json"
        
        snapshot = MigrationSnapshot(
            metadata={'test': 'data', 'value': 123}
        )
        snapshot.save(snapshot_path)
        
        assert snapshot_path.exists()
        
        loaded = MigrationSnapshot.load(snapshot_path)
        assert loaded.metadata['test'] == 'data'
        assert loaded.metadata['value'] == 123


class TestKyberMigrator:
    """Test Kyber key encapsulation migrator."""
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_kyber_keypair_generation(self):
        """Test Kyber key pair generation."""
        migrator = KyberMigrator()
        public_key, secret_key = migrator.generate_kyber_keypair()
        
        assert isinstance(public_key, bytes)
        assert isinstance(secret_key, bytes)
        assert len(public_key) > 0
        assert len(secret_key) > 0
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_kyber_encapsulation(self):
        """Test Kyber encapsulation."""
        migrator = KyberMigrator()
        public_key, _ = migrator.generate_kyber_keypair()
        
        ciphertext, shared_secret = migrator.encapsulate(public_key)
        
        assert isinstance(ciphertext, bytes)
        assert isinstance(shared_secret, bytes)
        assert len(ciphertext) > 0
        assert len(shared_secret) > 0
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_kyber_decapsulation(self):
        """Test Kyber decapsulation."""
        migrator = KyberMigrator()
        public_key, secret_key = migrator.generate_kyber_keypair()
        ciphertext, original_secret = migrator.encapsulate(public_key)
        
        recovered_secret = migrator.decapsulate(secret_key, ciphertext)
        
        assert recovered_secret == original_secret
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_kyber_benchmarking(self):
        """Test Kyber performance benchmarking."""
        migrator = KyberMigrator()
        benchmarks = migrator.benchmark_operations(iterations=5)
        
        assert len(benchmarks) >= 3
        assert all(isinstance(b, PerformanceBenchmark) for b in benchmarks)
        assert all(b.duration_ms > 0 for b in benchmarks)


class TestDilithiumMigrator:
    """Test Dilithium signature migrator."""
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_dilithium_keypair_generation(self):
        """Test Dilithium key pair generation."""
        migrator = DilithiumMigrator()
        public_key, secret_key = migrator.generate_signing_keypair()
        
        assert isinstance(public_key, bytes)
        assert isinstance(secret_key, bytes)
        assert len(public_key) > 0
        assert len(secret_key) > 0
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_dilithium_signing(self):
        """Test Dilithium signing."""
        migrator = DilithiumMigrator()
        _, secret_key = migrator.generate_signing_keypair()
        
        message = b"Test message for signing"
        signature = migrator.sign(secret_key, message)
        
        assert isinstance(signature, bytes)
        assert len(signature) > 0
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_dilithium_verification(self):
        """Test Dilithium signature verification."""
        migrator = DilithiumMigrator()
        public_key, secret_key = migrator.generate_signing_keypair()
        
        message = b"Test message for signing"
        signature = migrator.sign(secret_key, message)
        
        assert migrator.verify(public_key, message, signature)
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_dilithium_verification_fails_with_wrong_message(self):
        """Test that verification fails with tampered message."""
        migrator = DilithiumMigrator()
        public_key, secret_key = migrator.generate_signing_keypair()
        
        message = b"Original message"
        signature = migrator.sign(secret_key, message)
        
        tampered_message = b"Tampered message"
        assert not migrator.verify(public_key, tampered_message, signature)
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_dilithium_benchmarking(self):
        """Test Dilithium performance benchmarking."""
        migrator = DilithiumMigrator()
        benchmarks = migrator.benchmark_operations(iterations=5)
        
        assert len(benchmarks) >= 3
        assert all(isinstance(b, PerformanceBenchmark) for b in benchmarks)
        assert all(b.duration_ms > 0 for b in benchmarks)


class TestHybridMigrator:
    """Test hybrid classical + PQC migration."""
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_hybrid_kex_keypair_creation(self):
        """Test creating hybrid key exchange keypair."""
        migrator = HybridMigrator()
        classical_public = b"mock_rsa_public_key"
        
        hybrid_keypair = migrator.create_hybrid_kex_keypair(classical_public)
        
        assert 'classical_public_key' in hybrid_keypair
        assert 'kyber_public_key' in hybrid_keypair
        assert 'kyber_secret_key' in hybrid_keypair
        assert hybrid_keypair['classical_public_key'] == classical_public
        assert hybrid_keypair['classical_key_type'] == "RSA-2048"
    
    @pytest.mark.skipif(True, reason="liboqs not available in test environment")
    def test_hybrid_signature_keypair_creation(self):
        """Test creating hybrid signature keypair."""
        migrator = HybridMigrator()
        classical_public = b"mock_ecdsa_public_key"
        
        hybrid_keypair = migrator.create_hybrid_signature_keypair(classical_public)
        
        assert 'classical_public_key' in hybrid_keypair
        assert 'dilithium_public_key' in hybrid_keypair
        assert 'dilithium_secret_key' in hybrid_keypair
        assert hybrid_keypair['classical_public_key'] == classical_public
        assert hybrid_keypair['classical_key_type'] == "ECDSA-P256"


class TestMigrationRollback:
    """Test rollback mechanism for failed migrations."""
    
    def test_rollback_snapshot_creation(self, tmp_path):
        """Test creating rollback snapshots."""
        rollback = MigrationRollback(tmp_path)
        
        snapshot = MigrationSnapshot(
            metadata={'algorithm': 'RSA-2048'}
        )
        
        path = rollback.create_snapshot("test-migration-1", snapshot)
        
        assert path.exists()
        assert path.name == "test-migration-1.json"
    
    def test_rollback_snapshot_restoration(self, tmp_path):
        """Test restoring rollback snapshots."""
        rollback = MigrationRollback(tmp_path)
        
        original = MigrationSnapshot(
            metadata={'algorithm': 'RSA-2048', 'keys': 5}
        )
        rollback.create_snapshot("test-migration-1", original)
        
        restored = rollback.restore_snapshot("test-migration-1")
        
        assert restored.metadata['algorithm'] == 'RSA-2048'
        assert restored.metadata['keys'] == 5
    
    def test_rollback_list_snapshots(self, tmp_path):
        """Test listing available snapshots."""
        rollback = MigrationRollback(tmp_path)
        
        snapshot1 = MigrationSnapshot(metadata={'test': 1})
        snapshot2 = MigrationSnapshot(metadata={'test': 2})
        
        rollback.create_snapshot("migration-1", snapshot1)
        rollback.create_snapshot("migration-2", snapshot2)
        
        snapshots = rollback.list_snapshots()
        
        assert len(snapshots) == 2
        assert "migration-1" in snapshots
        assert "migration-2" in snapshots
    
    def test_rollback_delete_snapshot(self, tmp_path):
        """Test deleting snapshots."""
        rollback = MigrationRollback(tmp_path)
        
        snapshot = MigrationSnapshot(metadata={'test': 'data'})
        rollback.create_snapshot("test-migration", snapshot)
        
        assert len(rollback.list_snapshots()) == 1
        
        deleted = rollback.delete_snapshot("test-migration")
        
        assert deleted is True
        assert len(rollback.list_snapshots()) == 0


class TestPQCMigrator:
    """Test main PQC migration engine."""
    
    def test_pqc_migrator_initialization(self):
        """Test initializing PQC migrator."""
        migrator = PQCMigrator()
        
        assert migrator.rollback is not None
        assert isinstance(migrator.migrations, dict)
    
    def test_create_migration_plan(self):
        """Test creating a migration plan."""
        migrator = PQCMigrator()
        
        plan = migrator.create_migration_plan(
            source=ClassicalAlgorithm.RSA_2048,
            target=PQCAlgorithm.KYBER_1024,
            mode=MigrationMode.FULL_REPLACEMENT,
        )
        
        assert plan.id in migrator.migrations
        assert plan.source_algorithm == ClassicalAlgorithm.RSA_2048
        assert plan.target_algorithm == PQCAlgorithm.KYBER_1024
    
    def test_start_migration(self):
        """Test starting a migration."""
        migrator = PQCMigrator()
        
        plan = migrator.create_migration_plan(
            source=ClassicalAlgorithm.RSA_2048,
            target=PQCAlgorithm.KYBER_1024,
        )
        
        success = migrator.start_migration(plan.id)
        
        assert success is True
        assert plan.status == MigrationStatus.IN_PROGRESS
        assert plan.started_at is not None
    
    def test_complete_migration(self):
        """Test completing a migration."""
        migrator = PQCMigrator()
        
        plan = migrator.create_migration_plan(
            source=ClassicalAlgorithm.RSA_2048,
            target=PQCAlgorithm.KYBER_1024,
        )
        migrator.start_migration(plan.id)
        
        success = migrator.complete_migration(plan.id)
        
        assert success is True
        assert plan.status == MigrationStatus.COMPLETED
        assert plan.completed_at is not None
    
    def test_fail_migration(self):
        """Test marking migration as failed."""
        migrator = PQCMigrator()
        
        plan = migrator.create_migration_plan(
            source=ClassicalAlgorithm.RSA_2048,
            target=PQCAlgorithm.KYBER_1024,
        )
        
        success = migrator.fail_migration(plan.id, "Test failure")
        
        assert success is True
        assert plan.status == MigrationStatus.FAILED
        assert "Test failure" in plan.errors
    
    def test_rollback_migration(self, tmp_path):
        """Test rolling back a migration."""
        migrator = PQCMigrator(tmp_path)
        
        plan = migrator.create_migration_plan(
            source=ClassicalAlgorithm.RSA_2048,
            target=PQCAlgorithm.KYBER_1024,
        )
        migrator.start_migration(plan.id)
        
        success = migrator.rollback_migration(plan.id)
        
        assert success is True
        assert plan.status == MigrationStatus.ROLLED_BACK
    
    def test_get_migration_report(self):
        """Test generating migration report."""
        migrator = PQCMigrator()
        
        plan = migrator.create_migration_plan(
            source=ClassicalAlgorithm.RSA_2048,
            target=PQCAlgorithm.KYBER_1024,
        )
        
        report = migrator.get_migration_report(plan.id)
        
        assert 'migration_id' in report
        assert 'plan' in report
        assert report['plan']['source_algorithm'] == 'RSA-2048'
    
    def test_list_migrations(self):
        """Test listing all migrations."""
        migrator = PQCMigrator()
        
        plan1 = migrator.create_migration_plan(
            source=ClassicalAlgorithm.RSA_2048,
            target=PQCAlgorithm.KYBER_1024,
        )
        plan2 = migrator.create_migration_plan(
            source=ClassicalAlgorithm.ECDSA_P256,
            target=PQCAlgorithm.DILITHIUM3,
        )
        
        migrations = migrator.list_migrations()
        
        assert len(migrations) == 2
        assert any(m['source_algorithm'] == 'RSA-2048' for m in migrations)
        assert any(m['source_algorithm'] == 'ECDSA-P256' for m in migrations)


class TestMigrationFeasibility:
    """Test migration feasibility analysis."""
    
    def test_rsa_to_kyber_feasibility(self):
        """Test RSA to Kyber migration is feasible."""
        result = analyze_migration_feasibility(
            ClassicalAlgorithm.RSA_2048,
            PQCAlgorithm.KYBER_1024,
        )
        
        assert result['feasible'] is True
        assert result['compatibility'] == 'good'
        assert result['source'] == 'RSA-2048'
        assert result['target'] == 'Kyber1024'
    
    def test_ecdsa_to_dilithium_feasibility(self):
        """Test ECDSA to Dilithium migration is feasible."""
        result = analyze_migration_feasibility(
            ClassicalAlgorithm.ECDSA_P256,
            PQCAlgorithm.DILITHIUM3,
        )
        
        assert result['feasible'] is True
        assert result['compatibility'] == 'good'
        assert result['source'] == 'ECDSA-P256'
        assert result['target'] == 'Dilithium3'
    
    def test_incompatible_migration(self):
        """Test that incompatible migrations are marked as infeasible."""
        result = analyze_migration_feasibility(
            ClassicalAlgorithm.RSA_2048,
            PQCAlgorithm.DILITHIUM3,
        )
        
        assert result['feasible'] is False
        assert 'reason' in result
    
    def test_migration_effort_estimates(self):
        """Test that effort estimates are provided."""
        result = analyze_migration_feasibility(
            ClassicalAlgorithm.RSA_2048,
            PQCAlgorithm.KYBER_1024,
        )
        
        assert 'effort' in result
        assert 'timeline_weeks' in result
        assert result['timeline_weeks'] > 0


class TestPerformanceBenchmark:
    """Test performance benchmarking."""
    
    def test_benchmark_creation(self):
        """Test creating a performance benchmark."""
        bench = PerformanceBenchmark(
            metric=PerformanceMetric.KEY_GENERATION,
            algorithm="Kyber1024",
            operation="keypair_generation",
            duration_ms=2.5,
            memory_mb=10.2,
        )
        
        assert bench.metric == PerformanceMetric.KEY_GENERATION
        assert bench.algorithm == "Kyber1024"
        assert bench.duration_ms == 2.5
        assert bench.memory_mb == 10.2
    
    def test_benchmark_to_dict(self):
        """Test converting benchmark to dictionary."""
        bench = PerformanceBenchmark(
            metric=PerformanceMetric.ENCRYPTION,
            algorithm="Kyber1024",
            operation="encapsulation",
            duration_ms=1.5,
            memory_mb=8.0,
        )
        
        bench_dict = bench.to_dict()
        
        assert bench_dict['metric'] == 'encryption'
        assert bench_dict['algorithm'] == 'Kyber1024'
        assert bench_dict['duration_ms'] == 1.5


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