"""
test_api_config.py — Tests for API mode configuration and cost estimation.

Validates that cost estimates are accurate and API mode configuration
correctly disables expensive attacks.
"""

import pytest
from sichgate.api_config import (
    APIRunConfig,
    LocalRunConfig,
    APIQueryEstimate,
    estimate_api_queries,
    get_default_api_config,
    get_default_local_config,
    ATTACK_QUERY_COUNTS,
)
from sichgate.cost_estimator import CostEstimator


class TestAPIQueryEstimate:
    """Test query estimate calculations."""
    
    def test_query_estimate_initialization(self):
        """Test APIQueryEstimate initialization."""
        estimate = APIQueryEstimate(
            attack_name="test_attack",
            query_count=10,
            estimated_tokens=1000,
        )
        assert estimate.attack_name == "test_attack"
        assert estimate.query_count == 10
        assert estimate.estimated_tokens == 1000
    
    def test_cost_at_gpt4o(self):
        """Test GPT-4o cost calculation."""
        estimate = APIQueryEstimate(
            attack_name="test",
            query_count=27,
            estimated_tokens=21600,
        )
        cost = estimate.cost_at_gpt4o()
        # 27 queries × 0.003 ≈ 0.081
        assert 0.05 < cost < 0.15
    
    def test_cost_at_gpt35(self):
        """Test GPT-3.5 cost calculation."""
        estimate = APIQueryEstimate(
            attack_name="test",
            query_count=27,
            estimated_tokens=21600,
        )
        cost = estimate.cost_at_gpt35()
        # Cheaper than GPT-4o
        assert cost < estimate.cost_at_gpt4o()
    
    def test_cost_at_claude_haiku(self):
        """Test Claude Haiku cost calculation."""
        estimate = APIQueryEstimate(
            attack_name="test",
            query_count=27,
            estimated_tokens=21600,
        )
        cost = estimate.cost_at_claude_haiku()
        # Cheapest option
        assert cost < estimate.cost_at_gpt35()


class TestAPIRunConfig:
    """Test API mode configuration."""
    
    def test_api_mode_defaults(self):
        """Test that API mode disables expensive attacks by default."""
        config = APIRunConfig(mode="api")
        
        # High-signal attacks enabled
        assert config.run_prompt_injection is True
        assert config.run_instruction_hijacking is True
        assert config.run_safe_messaging is True
        assert config.run_hallucination is True
        
        # Expensive attacks disabled
        assert config.run_bias_fairness is False
        assert config.run_consistency is False
        assert config.run_crescendo is False
        assert config.run_membership_inference is False
        assert config.run_data_extraction is False
    
    def test_api_mode_timeout(self):
        """Test that API mode uses shorter timeouts."""
        config = APIRunConfig(mode="api")
        assert config.timeout_seconds == 30
    
    def test_api_mode_query_limit(self):
        """Test that API mode enforces query limit."""
        config = APIRunConfig(mode="api")
        assert config.max_total_queries == 30


class TestLocalRunConfig:
    """Test local mode configuration."""
    
    def test_local_mode_enables_all(self):
        """Test that local mode enables all attacks."""
        config = LocalRunConfig(mode="local")
        
        # All attacks enabled
        assert config.run_prompt_injection is True
        assert config.run_instruction_hijacking is True
        assert config.run_bias_fairness is True
        assert config.run_consistency is True
        assert config.run_crescendo is True
        assert config.run_membership_inference is True
        assert config.run_data_extraction is True
        assert config.run_model_extraction is True
    
    def test_local_mode_longer_timeout(self):
        """Test that local mode uses longer timeouts."""
        config = LocalRunConfig(mode="local")
        assert config.timeout_seconds == 120


class TestEstimateAPIQueries:
    """Test query estimation function."""
    
    def test_default_api_query_estimate(self):
        """Test estimate for default API configuration."""
        config = APIRunConfig()
        estimate = estimate_api_queries(config)
        
        # Should total around 27 queries for default API setup
        assert 20 < estimate.query_count < 35
        assert estimate.estimated_tokens > 0
    
    def test_custom_api_config_estimate(self):
        """Test estimate with custom API configuration."""
        config = APIRunConfig(
            run_prompt_injection=True,
            run_instruction_hijacking=False,  # Disable this
            run_safe_messaging=True,
            run_hallucination=True,
        )
        estimate = estimate_api_queries(config)
        
        # Should be less than default (8 fewer queries)
        assert estimate.query_count < 27
    
    def test_all_attacks_disabled(self):
        """Test estimate when all attacks disabled."""
        config = APIRunConfig(
            run_prompt_injection=False,
            run_instruction_hijacking=False,
            run_safe_messaging=False,
            run_hallucination=False,
            run_token_smuggling=False,
            run_pii_leakage=False,
        )
        estimate = estimate_api_queries(config)
        assert estimate.query_count == 0


class TestCostEstimator:
    """Test the cost estimator."""
    
    def test_cost_estimator_initialization(self):
        """Test CostEstimator initialization."""
        estimator = CostEstimator(mode="api")
        assert estimator.mode == "api"
    
    def test_pricing_data_exists(self):
        """Test that pricing data is populated."""
        estimator = CostEstimator()
        
        # All expected models should have pricing
        assert "gpt-4o" in estimator.PRICING
        assert "gpt-3.5" in estimator.PRICING
        assert "claude-3-haiku" in estimator.PRICING
        assert "claude-3-sonnet" in estimator.PRICING
    
    def test_rate_limits_exist(self):
        """Test that rate limit data is populated."""
        estimator = CostEstimator()
        
        # All expected models should have rate limits
        assert "gpt-4o" in estimator.RATE_LIMITS
        assert "gpt-3.5" in estimator.RATE_LIMITS
        assert "claude-3-haiku" in estimator.RATE_LIMITS
    
    def test_estimate_run(self):
        """Test full estimate run calculation."""
        estimator = CostEstimator(mode="api")
        config = APIRunConfig()
        estimate = estimator.estimate_run(config, "gpt-4o")
        
        # Should have all required fields
        assert "mode" in estimate
        assert "total_queries" in estimate
        assert "estimated_tokens" in estimate
        assert "costs" in estimate
        assert "rate_limit_check" in estimate
        
        # Costs should exist for all providers
        assert "gpt-4o" in estimate["costs"]
        assert "gpt-3.5" in estimate["costs"]
        assert "claude-3-haiku" in estimate["costs"]
    
    def test_cost_ordering(self):
        """Test that costs are cheaper for cheaper providers."""
        estimator = CostEstimator(mode="api")
        config = APIRunConfig()
        estimate = estimator.estimate_run(config, "gpt-4o")
        
        costs = estimate["costs"]
        # Claude Haiku should be cheapest
        assert costs["claude-3-haiku"] < costs["gpt-3.5"]
        # GPT-3.5 should be cheaper than GPT-4
        assert costs["gpt-3.5"] < costs["gpt-4o"]
    
    def test_format_estimate_for_cli(self):
        """Test CLI formatting of estimates."""
        estimator = CostEstimator(mode="api")
        config = APIRunConfig()
        estimate = estimator.estimate_run(config, "gpt-4o")
        
        formatted = estimator.format_estimate_for_cli(estimate)
        
        # Should contain key information
        assert "API COST ESTIMATE" in formatted
        assert "Total API Calls" in formatted
        assert "OpenAI GPT-4o" in formatted
        assert "Claude 3 Haiku" in formatted


class TestGetDefaultConfigs:
    """Test default configuration factories."""
    
    def test_get_default_api_config(self):
        """Test getting default API configuration."""
        config = get_default_api_config()
        
        assert isinstance(config, APIRunConfig)
        assert config.mode == "api"
        assert config.run_prompt_injection is True
        assert config.run_bias_fairness is False
    
    def test_get_default_local_config(self):
        """Test getting default local configuration."""
        config = get_default_local_config()
        
        assert isinstance(config, LocalRunConfig)
        assert config.mode == "local"
        assert config.run_bias_fairness is True


class TestAttackQueryCounts:
    """Test attack query count estimates."""
    
    def test_query_count_constants_exist(self):
        """Test that query counts are defined for expected attacks."""
        expected_attacks = [
            "prompt_injection_direct",
            "instruction_hijacking",
            "safe_messaging",
            "hallucination",
            "bias_fairness",
            "consistency",
            "crescendo",
        ]
        
        for attack in expected_attacks:
            assert attack in ATTACK_QUERY_COUNTS
            assert ATTACK_QUERY_COUNTS[attack] > 0
    
    def test_expensive_attacks_have_high_counts(self):
        """Test that expensive attacks are marked as such."""
        assert ATTACK_QUERY_COUNTS["bias_fairness"] > 40
        assert ATTACK_QUERY_COUNTS["consistency"] > 15
        assert ATTACK_QUERY_COUNTS["crescendo"] > 10


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