"""
Test module for quantum risk calculation engine
"""

import pytest
import json
import csv
from datetime import datetime, timedelta
from pathlib import Path
import tempfile
import numpy as np

from chronos.core.analyze.risk_calculator import (
    QuantumRiskCalculator,
    QDayProbabilityModel,
    DataAsset,
    DataAssetType,
    DataSensitivity,
    CryptographyType,
    QuantumThreatAssessment,
    DataValueDecay,
    MigrationScenario,
    RiskMetrics,
)


class TestQDayProbabilityModel:
    """Test Q-Day probability modeling."""
    
    def test_qday_model_initialization(self):
        """Test Q-Day model initialization."""
        model = QDayProbabilityModel(
            target_year=2035,
            standard_deviation=5,
        )
        assert model.target_year == 2035
        assert model.standard_deviation == 5
    
    def test_probability_at_year(self):
        """Test probability calculation at specific years."""
        model = QDayProbabilityModel(target_year=2035, standard_deviation=5)
        
        # Before target year (should be increasing)
        prob_2020 = model.get_probability_at_year(2020)
        prob_2025 = model.get_probability_at_year(2025)
        prob_2030 = model.get_probability_at_year(2030)
        assert prob_2020 <= prob_2025 <= prob_2030
        
        # Around target year
        prob_2035 = model.get_probability_at_year(2035)
        assert 0.15 < prob_2035 < 0.8  # Peak should be somewhere in range
        
        # After target year
        prob_2040 = model.get_probability_at_year(2040)
        prob_2045 = model.get_probability_at_year(2045)
        assert prob_2035 < prob_2040 < prob_2045
        
        # Far future
        prob_2100 = model.get_probability_at_year(2100)
        assert prob_2100 > 0.95
    
    def test_year_at_probability(self):
        """Test inverse calculation (year at probability)."""
        model = QDayProbabilityModel(target_year=2035)
        
        # Test monotonic increase with probability
        year_10 = model.get_year_at_probability(0.1)
        year_50 = model.get_year_at_probability(0.5)
        year_90 = model.get_year_at_probability(0.9)
        
        # Higher probabilities should give later years
        assert year_10 < year_50 < year_90
        
        # All should be reasonable years
        assert 1970 < year_10 < 2100
        assert 1970 < year_50 < 2100
        assert 1970 < year_90 < 2100
    
    def test_probability_bounds(self):
        """Test probability bounds (0-1)."""
        model = QDayProbabilityModel()
        
        for year in range(2000, 2100, 5):
            prob = model.get_probability_at_year(year)
            assert 0.0 <= prob <= 1.0


class TestDataAsset:
    """Test data asset models."""
    
    def test_asset_creation(self):
        """Test data asset creation."""
        asset = DataAsset(
            name="Customer Database",
            asset_type=DataAssetType.DATABASE,
            sensitivity=DataSensitivity.RESTRICTED,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
        )
        
        assert asset.name == "Customer Database"
        assert asset.sensitivity == DataSensitivity.RESTRICTED
    
    def test_asset_value_score(self):
        """Test asset value scoring."""
        # Restricted asset
        asset_restricted = DataAsset(
            name="High Value",
            asset_type=DataAssetType.VAULT,
            sensitivity=DataSensitivity.RESTRICTED,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
            business_impact=10,
        )
        
        # Public asset
        asset_public = DataAsset(
            name="Public Data",
            asset_type=DataAssetType.FILE_STORAGE,
            sensitivity=DataSensitivity.PUBLIC,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
            business_impact=1,
        )
        
        assert asset_restricted.get_value_score() > asset_public.get_value_score()


class TestQuantumRiskCalculator:
    """Test quantum risk calculator."""
    
    def test_calculator_initialization(self):
        """Test calculator initialization."""
        calc = QuantumRiskCalculator(verbose=False)
        assert calc.verbose is False
        assert len(calc.assets) == 0
    
    def test_add_asset(self):
        """Test adding assets."""
        calc = QuantumRiskCalculator()
        asset = DataAsset(
            name="Test Asset",
            asset_type=DataAssetType.DATABASE,
            sensitivity=DataSensitivity.CONFIDENTIAL,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
        )
        
        calc.add_asset(asset)
        assert len(calc.assets) == 1
    
    def test_threat_level_calculation(self):
        """Test threat level calculation."""
        calc = QuantumRiskCalculator()
        
        # Critical threat
        critical_asset = DataAsset(
            name="Weak RSA",
            asset_type=DataAssetType.CERTIFICATE,
            sensitivity=DataSensitivity.RESTRICTED,
            crypto_type=CryptographyType.RSA,
            key_size=1024,
        )
        
        threat = calc._calculate_threat_level(critical_asset)
        assert threat == "critical"
        
        # Low threat
        safe_asset = DataAsset(
            name="Strong RSA",
            asset_type=DataAssetType.DATABASE,
            sensitivity=DataSensitivity.PUBLIC,
            crypto_type=CryptographyType.RSA,
            key_size=4096,
        )
        
        threat = calc._calculate_threat_level(safe_asset)
        assert threat == "low"
    
    def test_harvest_now_risk(self):
        """Test harvest-now-decrypt-later risk calculation."""
        calc = QuantumRiskCalculator()
        
        # High risk asset
        asset_high = DataAsset(
            name="Sensitive",
            asset_type=DataAssetType.VAULT,
            sensitivity=DataSensitivity.RESTRICTED,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
            metadata={"data_lifetime": 20},
        )
        
        risk_high = calc._calculate_harvest_now_risk(asset_high)
        assert risk_high > 0.7
        
        # Low risk asset
        asset_low = DataAsset(
            name="Public",
            asset_type=DataAssetType.FILE_STORAGE,
            sensitivity=DataSensitivity.PUBLIC,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
        )
        
        risk_low = calc._calculate_harvest_now_risk(asset_low)
        assert risk_low < 0.3
    
    def test_assess_quantum_threat(self):
        """Test quantum threat assessment."""
        calc = QuantumRiskCalculator()
        asset = DataAsset(
            name="Test Asset",
            asset_type=DataAssetType.DATABASE,
            sensitivity=DataSensitivity.CONFIDENTIAL,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
            metadata={"data_lifetime": 15},
        )
        
        assessment = calc.assess_quantum_threat(asset)
        
        assert assessment.asset_name == "Test Asset"
        assert assessment.crypto_type == CryptographyType.RSA
        assert assessment.threat_level in ["critical", "high", "medium", "low"]
        assert 0.0 <= assessment.harvest_now_risk <= 1.0
        assert 0.0 <= assessment.urgency_score <= 100.0
    
    def test_assess_all_assets(self):
        """Test assessing all assets."""
        calc = QuantumRiskCalculator()
        
        # Add multiple assets
        for i in range(5):
            asset = DataAsset(
                name=f"Asset {i}",
                asset_type=DataAssetType.DATABASE,
                sensitivity=DataSensitivity.CONFIDENTIAL,
                crypto_type=CryptographyType.RSA,
                key_size=2048 + (i * 512),
            )
            calc.add_asset(asset)
        
        assessments = calc.assess_all_assets()
        assert len(assessments) == 5
    
    def test_migration_scenario_roi(self):
        """Test migration ROI calculation."""
        calc = QuantumRiskCalculator()
        
        # Add some assets
        for i in range(10):
            asset = DataAsset(
                name=f"Asset {i}",
                asset_type=DataAssetType.DATABASE,
                sensitivity=DataSensitivity.CONFIDENTIAL,
                crypto_type=CryptographyType.RSA,
                key_size=2048,
                business_impact=5,
            )
            calc.add_asset(asset)
        
        # Create migration scenario
        scenario = MigrationScenario(
            name="PQC Migration 2025",
            migration_year=2025,
            pqc_algorithm="ML-KEM",
            cost_per_asset=5000,
        )
        
        roi = calc.calculate_migration_roi(scenario)
        
        assert "total_cost" in roi
        assert "expected_loss_avoided" in roi
        assert "roi_percentage" in roi
        assert roi["total_cost"] > 0
    
    def test_risk_metrics(self):
        """Test risk metrics calculation."""
        calc = QuantumRiskCalculator()
        
        # Add mixed assets
        for i in range(3):
            asset = DataAsset(
                name=f"Critical {i}",
                asset_type=DataAssetType.VAULT,
                sensitivity=DataSensitivity.RESTRICTED,
                crypto_type=CryptographyType.RSA,
                key_size=1024,  # Weak
            )
            calc.add_asset(asset)
        
        for i in range(7):
            asset = DataAsset(
                name=f"Safe {i}",
                asset_type=DataAssetType.DATABASE,
                sensitivity=DataSensitivity.INTERNAL,
                crypto_type=CryptographyType.RSA,
                key_size=4096,
            )
            calc.add_asset(asset)
        
        metrics = calc.get_risk_metrics()
        
        assert isinstance(metrics, RiskMetrics)
        assert metrics.critical_threat_assets > 0
        assert metrics.total_quantum_risk > 0
        assert metrics.timeline_to_resilience > 0
    
    def test_data_value_decay(self):
        """Test data value decay calculation."""
        calc = QuantumRiskCalculator()
        asset = DataAsset(
            name="Decaying Data",
            asset_type=DataAssetType.DATABASE,
            sensitivity=DataSensitivity.CONFIDENTIAL,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
        )
        
        decay = calc.calculate_data_value_decay(asset, retention_years=10)
        
        assert decay.asset_name == "Decaying Data"
        
        # Value should decrease over time
        value_now = decay.get_value_at_year(datetime.now().year)
        value_future = decay.get_value_at_year(datetime.now().year + 5)
        
        assert value_now > value_future
        assert value_future > 0


class TestDataValueDecay:
    """Test data value decay model."""
    
    def test_decay_initialization(self):
        """Test decay model initialization."""
        decay = DataValueDecay(
            asset_name="Test",
            initial_value=1000,
            sensitivity=DataSensitivity.CONFIDENTIAL,
            creation_date=datetime.now(),
            retention_years=10,
            decay_rate=0.15,
        )
        
        assert decay.initial_value == 1000
        assert decay.retention_years == 10
    
    def test_value_at_year(self):
        """Test value calculation at specific year."""
        now = datetime.now()
        decay = DataValueDecay(
            asset_name="Test",
            initial_value=100.0,
            sensitivity=DataSensitivity.CONFIDENTIAL,
            creation_date=now,
            retention_years=10,
        )
        
        # At creation year
        value_now = decay.get_value_at_year(now.year)
        assert value_now == 100.0
        
        # After retention period
        value_expired = decay.get_value_at_year(now.year + 15)
        assert value_expired == 0.0
        
        # Mid-period
        value_mid = decay.get_value_at_year(now.year + 5)
        assert 0.0 < value_mid < 100.0
    
    def test_decay_curve(self):
        """Test decay curve generation."""
        now = datetime.now()
        decay = DataValueDecay(
            asset_name="Test",
            initial_value=100.0,
            sensitivity=DataSensitivity.CONFIDENTIAL,
            creation_date=now,
            retention_years=10,
        )
        
        curve = decay.get_value_curve(now.year, now.year + 12, step=2)
        
        # Should have values for each year
        assert len(curve) > 0
        
        # Values should be monotonically decreasing or constant
        values = list(curve.values())
        for i in range(len(values) - 1):
            assert values[i] >= values[i + 1]


class TestMigrationScenario:
    """Test migration scenarios."""
    
    def test_scenario_creation(self):
        """Test migration scenario creation."""
        scenario = MigrationScenario(
            name="PQC Migration",
            migration_year=2025,
            pqc_algorithm="ML-KEM",
            cost_per_asset=5000,
        )
        
        assert scenario.name == "PQC Migration"
        assert scenario.migration_year == 2025
    
    def test_total_cost_calculation(self):
        """Test migration cost calculation."""
        scenario = MigrationScenario(
            name="PQC Migration",
            migration_year=2025,
            pqc_algorithm="ML-KEM",
            cost_per_asset=5000,
        )
        
        total_cost = scenario.get_total_cost(num_assets=10)
        assert total_cost == 50000
    
    def test_completion_date(self):
        """Test completion date calculation."""
        scenario = MigrationScenario(
            name="PQC Migration",
            migration_year=2025,
            pqc_algorithm="ML-KEM",
            testing_months=6,
            rollout_months=12,
        )
        
        completion = scenario.get_completion_date()
        assert completion.year >= 2025


class TestExportFunctionality:
    """Test export functionality."""
    
    def test_export_json(self):
        """Test JSON export."""
        calc = QuantumRiskCalculator()
        
        # Add assets
        asset = DataAsset(
            name="Test Asset",
            asset_type=DataAssetType.DATABASE,
            sensitivity=DataSensitivity.CONFIDENTIAL,
            crypto_type=CryptographyType.RSA,
            key_size=2048,
        )
        calc.add_asset(asset)
        
        # Export to JSON
        with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
            filepath = f.name
        
        calc.export_risk_report_json(filepath)
        
        # Verify file exists and is valid JSON
        assert Path(filepath).exists()
        
        with open(filepath) as f:
            data = json.load(f)
            assert "summary" in data
            assert "assessments" in data
        
        # Cleanup
        Path(filepath).unlink()
    
    def test_export_csv(self):
        """Test CSV export."""
        calc = QuantumRiskCalculator()
        
        # Add assets
        for i in range(3):
            asset = DataAsset(
                name=f"Asset {i}",
                asset_type=DataAssetType.DATABASE,
                sensitivity=DataSensitivity.CONFIDENTIAL,
                crypto_type=CryptographyType.RSA,
                key_size=2048 + (i * 256),
            )
            calc.add_asset(asset)
        
        # Export to CSV
        with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
            filepath = f.name
        
        calc.export_risk_report_csv(filepath)
        
        # Verify file exists and is valid CSV
        assert Path(filepath).exists()
        
        with open(filepath) as f:
            reader = csv.DictReader(f)
            rows = list(reader)
            assert len(rows) == 3
            assert "asset_name" in rows[0]
        
        # Cleanup
        Path(filepath).unlink()


class TestMonteCarloSimulation:
    """Test Monte Carlo simulation."""
    
    def test_monte_carlo_simulation(self):
        """Test Monte Carlo simulation runs."""
        calc = QuantumRiskCalculator()
        
        # Add some assets
        for i in range(5):
            asset = DataAsset(
                name=f"Asset {i}",
                asset_type=DataAssetType.DATABASE,
                sensitivity=DataSensitivity.CONFIDENTIAL,
                crypto_type=CryptographyType.RSA,
                key_size=2048,
            )
            calc.add_asset(asset)
        
        # Run simulation
        results = calc.monte_carlo_simulation(
            num_simulations=100, years=5
        )
        
        # Verify results
        assert len(results) > 0
        
        # Results should be increasing over time (risk increases)
        years = sorted(results.keys())
        values = [results[y] for y in years]
        
        # General trend should be increasing
        assert values[-1] >= values[0]


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