"""Data models for the CostWise MCP SDK."""

from dataclasses import dataclass, field
from enum import Enum
from typing import Optional
import time
import uuid


class TaskComplexity(str, Enum):
    """Task complexity levels used for model routing and token limiting."""
    SIMPLE = "simple"
    MEDIUM = "medium"
    COMPLEX = "complex"


class Provider(str, Enum):
    """Supported LLM providers."""
    OPENAI = "openai"
    ANTHROPIC = "anthropic"
    GOOGLE = "google"
    MISTRAL = "mistral"
    COHERE = "cohere"


@dataclass
class Decision:
    """
    The optimization decision returned by CostPolicy.optimize().

    Use this to call your LLM provider:
        decision = policy.optimize(prompt="...", model="gpt-4")
        response = openai.chat.completions.create(
            model=decision.recommended_model,
            max_tokens=decision.max_tokens,
        )
    """
    # What you should use
    recommended_model: str
    max_tokens: int
    estimated_cost: float

    # What was requested
    original_model: str
    input_tokens: int
    complexity: TaskComplexity

    # Metadata
    allowed: bool = True
    savings_pct: float = 0.0
    message: str = ""
    request_id: str = field(default_factory=lambda: uuid.uuid4().hex[:12])
    timestamp: float = field(default_factory=time.time)

    # Cost breakdown
    original_estimated_cost: float = 0.0
    estimated_savings: float = 0.0

    def to_dict(self) -> dict:
        return {
            "request_id": self.request_id,
            "allowed": self.allowed,
            "recommended_model": self.recommended_model,
            "original_model": self.original_model,
            "max_tokens": self.max_tokens,
            "input_tokens": self.input_tokens,
            "estimated_cost": round(self.estimated_cost, 8),
            "original_estimated_cost": round(self.original_estimated_cost, 8),
            "estimated_savings": round(self.estimated_savings, 8),
            "savings_pct": round(self.savings_pct, 2),
            "complexity": self.complexity.value,
            "message": self.message,
            "timestamp": self.timestamp,
        }


@dataclass
class UsageReport:
    """Report actual usage after LLM call completes (for backend analytics)."""
    request_id: str
    model: str
    input_tokens: int
    output_tokens: int
    total_tokens: int
    actual_cost: Optional[float] = None
    latency_ms: Optional[float] = None
    timestamp: float = field(default_factory=time.time)

    def to_dict(self) -> dict:
        return {
            "request_id": self.request_id,
            "model": self.model,
            "input_tokens": self.input_tokens,
            "output_tokens": self.output_tokens,
            "total_tokens": self.total_tokens,
            "actual_cost": self.actual_cost,
            "latency_ms": self.latency_ms,
            "timestamp": self.timestamp,
        }


@dataclass
class PolicyConfig:
    """
    Configuration for the cost policy engine.
    Override defaults to customize behavior per project/team.
    """
    # Token limits per complexity
    max_tokens_simple: int = 256
    max_tokens_medium: int = 1024
    max_tokens_complex: int = 4096

    # Budget enforcement (set to None to disable)
    daily_budget_usd: Optional[float] = None
    monthly_budget_usd: Optional[float] = None

    # Model routing
    auto_downgrade: bool = True          # Auto-select cheaper models for simple tasks
    allow_model_override: bool = False   # Let SDK change model even if user specifies one
    blocked_models: list = field(default_factory=list)  # Models to never use

    # Telemetry
    telemetry_enabled: bool = True
    telemetry_batch_size: int = 50       # Send after N decisions
    telemetry_flush_interval: float = 30.0  # Send every N seconds

    def get_max_tokens(self, complexity: TaskComplexity) -> int:
        return {
            TaskComplexity.SIMPLE: self.max_tokens_simple,
            TaskComplexity.MEDIUM: self.max_tokens_medium,
            TaskComplexity.COMPLEX: self.max_tokens_complex,
        }[complexity]
