from datetime import date
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field
from typing_extensions import TypedDict


class GroupByDimension(str, Enum):
    DATE = "date"
    MODEL = "model"
    API_KEY = "api_key"
    TEAM = "team"
    ORGANIZATION = "organization"
    MODEL_GROUP = "model_group"
    PROVIDER = "custom_llm_provider"


class SpendMetrics(BaseModel):
    spend: float = Field(default=0.0)
    flat_cost: float = Field(default=0.0)
    prompt_tokens: int = Field(default=0)
    completion_tokens: int = Field(default=0)
    cache_read_input_tokens: int = Field(default=0)
    cache_creation_input_tokens: int = Field(default=0)
    compression_saved_tokens: int = Field(default=0)
    compression_savings_spend: float = Field(default=0.0)
    prompt_caching_savings_spend: float = Field(default=0.0)
    gateway_injected_caching_savings_spend: float = Field(default=0.0)
    autorouter_savings_spend: float = Field(default=0.0)
    total_tokens: int = Field(default=0)
    successful_requests: int = Field(default=0)
    failed_requests: int = Field(default=0)
    api_requests: int = Field(default=0)


class MetricBase(BaseModel):
    metrics: SpendMetrics


class KeyMetadata(BaseModel):
    """Metadata for a key"""

    key_alias: str | None = None
    team_id: str | None = None


class KeyMetricWithMetadata(MetricBase):
    """Base class for metrics with additional metadata"""

    metadata: KeyMetadata = Field(default_factory=KeyMetadata)


class MetricWithMetadata(MetricBase):
    metadata: dict[str, Any] = Field(default_factory=dict)
    # API key breakdown for this metric (e.g., which API keys are using this MCP server)
    api_key_breakdown: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict)  # api_key -> {metrics, metadata}


class BreakdownMetrics(BaseModel):
    """Breakdown of spend by different dimensions"""

    mcp_servers: dict[str, MetricWithMetadata] = Field(default_factory=dict)  # mcp_server -> {metrics, metadata}
    models: dict[str, MetricWithMetadata] = Field(default_factory=dict)  # model -> {metrics, metadata}
    model_groups: dict[str, MetricWithMetadata] = Field(default_factory=dict)  # model_group -> {metrics, metadata}
    providers: dict[str, MetricWithMetadata] = Field(default_factory=dict)  # provider -> {metrics, metadata}
    endpoints: dict[str, MetricWithMetadata] = Field(default_factory=dict)  # endpoint -> {metrics, metadata}
    api_keys: dict[str, KeyMetricWithMetadata] = Field(default_factory=dict)  # api_key -> {metrics, metadata}
    entities: dict[str, MetricWithMetadata] = Field(default_factory=dict)  # entity -> {metrics, metadata}


class DailySpendData(BaseModel):
    date: date
    metrics: SpendMetrics
    breakdown: BreakdownMetrics = Field(default_factory=BreakdownMetrics)


class DailySpendMetadata(BaseModel):
    total_spend: float = Field(default=0.0)
    total_flat_cost: float = Field(default=0.0)
    total_prompt_tokens: int = Field(default=0)
    total_completion_tokens: int = Field(default=0)
    total_tokens: int = Field(default=0)
    total_api_requests: int = Field(default=0)
    total_successful_requests: int = Field(default=0)
    total_failed_requests: int = Field(default=0)
    total_cache_read_input_tokens: int = Field(default=0)
    total_cache_creation_input_tokens: int = Field(default=0)
    total_compression_saved_tokens: int = Field(default=0)
    total_compression_savings_spend: float = Field(default=0.0)
    total_prompt_caching_savings_spend: float = Field(default=0.0)
    total_gateway_injected_caching_savings_spend: float = Field(default=0.0)
    total_autorouter_savings_spend: float = Field(default=0.0)
    page: int = Field(default=1)
    total_pages: int = Field(default=1)
    has_more: bool = Field(default=False)


class SpendAnalyticsPaginatedResponse(BaseModel):
    results: list[DailySpendData]
    metadata: DailySpendMetadata = Field(default_factory=DailySpendMetadata)


class LiteLLM_DailyUserSpend(BaseModel):
    id: str
    user_id: str
    date: str
    api_key: str
    mcp_server_id: str | None = None
    model: str | None = None
    model_group: str | None = None
    custom_llm_provider: str | None = None
    prompt_tokens: int = 0
    completion_tokens: int = 0
    cache_read_input_tokens: int = 0
    cache_creation_input_tokens: int = 0
    compression_saved_tokens: int = 0
    compression_savings_spend: float = 0.0
    prompt_caching_savings_spend: float = 0.0
    gateway_injected_caching_savings_spend: float = 0.0
    autorouter_savings_spend: float = 0.0
    spend: float = 0.0
    api_requests: int = 0
    successful_requests: int = 0
    failed_requests: int = 0


class GroupedData(TypedDict):
    metrics: SpendMetrics
    breakdown: BreakdownMetrics
