"""
cost_estimator.py — Real-time cost & query estimation for API deployments.

Shows users exactly what they're about to pay before they incur the cost.
This is critical for customer trust and preventing bill shock.
"""

from typing import Literal
from sichgate.api_config import APIQueryEstimate, estimate_api_queries, APIRunConfig


class CostEstimator:
    """
    Estimates API costs and query counts before execution.
    
    Customers deploying via OpenAI/Anthropic endpoints need to know:
    1. How many queries will hit their account
    2. Estimated cost at their pricing tier
    3. Whether they'll hit rate limits
    """
    
    # Pricing (as of Feb 2025, update as needed)
    PRICING = {
        "gpt-4o": {
            "name": "OpenAI GPT-4o",
            "input_per_1m": 5.00,      # $5 per 1M input tokens
            "output_per_1m": 15.00,    # $15 per 1M output tokens
        },
        "gpt-3.5": {
            "name": "OpenAI GPT-3.5 Turbo",
            "input_per_1m": 0.50,      # $0.50 per 1M input tokens
            "output_per_1m": 1.50,     # $1.50 per 1M output tokens
        },
        "claude-3-haiku": {
            "name": "Claude 3 Haiku",
            "input_per_1m": 0.80,      # $0.80 per 1M input tokens
            "output_per_1m": 4.00,     # $4.00 per 1M output tokens
        },
        "claude-3-sonnet": {
            "name": "Claude 3 Sonnet",
            "input_per_1m": 3.00,
            "output_per_1m": 15.00,
        },
    }
    
    # Typical rate limits per provider
    RATE_LIMITS = {
        "gpt-4o": {"requests_per_minute": 10000, "tokens_per_minute": 2000000},
        "gpt-3.5": {"requests_per_minute": 10000, "tokens_per_minute": 2000000},
        "claude-3-haiku": {"requests_per_minute": 10000, "tokens_per_minute": 40000},
        "claude-3-sonnet": {"requests_per_minute": 10000, "tokens_per_minute": 40000},
    }
    
    def __init__(self, mode: Literal["api", "local"] = "api"):
        self.mode = mode
    
    def estimate_run(
        self,
        config: APIRunConfig,
        model_provider: str = "gpt-4o",
    ) -> dict:
        """
        Estimate total cost and query count for a run.
        
        Args:
            config: APIRunConfig or LocalRunConfig
            model_provider: One of "gpt-4o", "gpt-3.5", "claude-3-haiku", "claude-3-sonnet"
        
        Returns:
            Dictionary with query count, estimated costs at different tiers, rate limit info
        """
        estimate = estimate_api_queries(config)
        
        return {
            "mode": self.mode,
            "total_queries": estimate.query_count,
            "estimated_tokens": estimate.estimated_tokens,
            "costs": {
                "gpt-4o": self._estimate_cost("gpt-4o", estimate),
                "gpt-3.5": self._estimate_cost("gpt-3.5", estimate),
                "claude-3-haiku": self._estimate_cost("claude-3-haiku", estimate),
                "claude-3-sonnet": self._estimate_cost("claude-3-sonnet", estimate),
            },
            "rate_limit_check": self._check_rate_limits("gpt-4o", estimate),  # default
        }
    
    def _estimate_cost(self, provider: str, estimate: APIQueryEstimate) -> float:
        """Calculate actual cost for a specific provider."""
        # Simplified: assume 80% of tokens are input, 20% output
        # Adjust these ratios based on your actual attack patterns
        input_tokens = int(estimate.estimated_tokens * 0.8)
        output_tokens = int(estimate.estimated_tokens * 0.2)
        
        pricing = self.PRICING.get(provider, {})
        input_per_1m = pricing.get("input_per_1m", 0)
        output_per_1m = pricing.get("output_per_1m", 0)
        
        input_cost = (input_tokens / 1_000_000) * input_per_1m
        output_cost = (output_tokens / 1_000_000) * output_per_1m
        
        return input_cost + output_cost
    
    def _check_rate_limits(self, provider: str, estimate: APIQueryEstimate) -> dict:
        """Check if query count would hit typical rate limits."""
        limits = self.RATE_LIMITS.get(provider, {})
        tokens_per_min = limits.get("tokens_per_minute", 0)
        requests_per_min = limits.get("requests_per_minute", 0)
        
        # Rough heuristic: if we do all queries in 5 minutes, would we exceed?
        tokens_per_min_needed = estimate.estimated_tokens / 5
        queries_per_min_needed = estimate.query_count / 5
        
        return {
            "provider": provider,
            "estimated_queries": estimate.query_count,
            "estimated_tokens": estimate.estimated_tokens,
            "rate_limit_tokens_per_min": tokens_per_min,
            "rate_limit_requests_per_min": requests_per_min,
            "would_exceed_token_limit": tokens_per_min_needed > tokens_per_min,
            "would_exceed_request_limit": queries_per_min_needed > requests_per_min,
            "recommended_delay_seconds": max([
                0,
                int((tokens_per_min_needed / tokens_per_min - 1) * 300) if tokens_per_min > 0 else 0,
            ]),
        }
    
    def format_estimate_for_cli(self, estimate_dict: dict) -> str:
        """
        Format estimate nicely for CLI display.
        
        Returns a formatted string showing costs, queries, and warnings.
        """
        lines = []
        lines.append("")
        lines.append("═" * 70)
        lines.append("                    API COST ESTIMATE")
        lines.append("═" * 70)
        lines.append("")
        
        # Query count
        lines.append(f"Total API Calls:              {estimate_dict['total_queries']} queries")
        lines.append(f"Estimated Tokens:            ~{estimate_dict['estimated_tokens']:,} tokens")
        lines.append("")
        
        # Costs by model
        lines.append("Estimated Cost by Model:")
        lines.append("─" * 70)
        costs = estimate_dict["costs"]
        for provider, cost in costs.items():
            pricing = self.PRICING.get(provider, {})
            name = pricing.get("name", provider)
            if cost < 0.01:
                cost_str = f"< $0.01"
            else:
                cost_str = f"${cost:.4f}"
            lines.append(f"  {name:<30} {cost_str:>10}")
        
        lines.append("")
        lines.append("═" * 70)
        
        # Rate limit check
        check = estimate_dict["rate_limit_check"]
        if check["would_exceed_token_limit"] or check["would_exceed_request_limit"]:
            lines.append("")
            lines.append("⚠️  WARNING: May hit rate limits!")
            lines.append(f"Recommended delay between calls: {check['recommended_delay_seconds']}s")
        else:
            lines.append("✓ Within typical rate limits")
        
        lines.append("═" * 70)
        lines.append("")
        
        return "\n".join(lines)
