"""Cost tracking for evaluation runs — pricing lookup and usage-based cost.

成本口径：评测自身消耗的 LLM token（judge/generator/embedding/moderation，
经 ModelGateway 归集进 RunSummary.eval_usage_by_purpose），按岗位所用模型的
单价计费。被评智能体的通道用量不在口径内（tutu-api 通道不返回用量）。

单价解析顺序：``data/model_pricing.json``（部署侧覆盖）→ ``DEFAULT_PRICING``。
"""

import json
from pathlib import Path
from typing import Any, Optional

from pydantic import BaseModel, Field

# data/model_pricing.json：与 DB 同目录，部署时挂载覆盖，缺省文件视为无覆盖
_PRICING_FILE = Path(__file__).resolve().parents[3] / "data" / "model_pricing.json"


class ModelPricing(BaseModel):
    """Pricing for a model (per 1M tokens)."""

    model_id: str
    prompt_cost_per_1m: float = Field(description="Cost per 1M prompt tokens in USD")
    completion_cost_per_1m: float = Field(description="Cost per 1M completion tokens in USD")


class CostBreakdown(BaseModel):
    """Cost breakdown for one usage bucket (e.g. one purpose or a whole run)."""

    prompt_tokens: int = 0
    completion_tokens: int = 0
    total_tokens: int = 0
    cost_usd: float = 0.0


# Default pricing for common models (per 1M tokens in USD)
DEFAULT_PRICING: dict[str, ModelPricing] = {
    "gpt-4o": ModelPricing(model_id="gpt-4o", prompt_cost_per_1m=5.0, completion_cost_per_1m=15.0),
    "gpt-4o-mini": ModelPricing(model_id="gpt-4o-mini", prompt_cost_per_1m=0.15, completion_cost_per_1m=0.60),
    "gpt-4-turbo": ModelPricing(model_id="gpt-4-turbo", prompt_cost_per_1m=10.0, completion_cost_per_1m=30.0),
    "gpt-3.5-turbo": ModelPricing(model_id="gpt-3.5-turbo", prompt_cost_per_1m=0.50, completion_cost_per_1m=1.50),
    "claude-3-5-sonnet": ModelPricing(model_id="claude-3-5-sonnet", prompt_cost_per_1m=3.0, completion_cost_per_1m=15.0),
    "claude-3-haiku": ModelPricing(model_id="claude-3-haiku", prompt_cost_per_1m=0.25, completion_cost_per_1m=1.25),
}

_pricing_overrides: Optional[dict[str, ModelPricing]] = None


def _load_overrides() -> dict[str, ModelPricing]:
    """Parse data/model_pricing.json once; a missing/broken file means no override."""
    global _pricing_overrides
    if _pricing_overrides is not None:
        return _pricing_overrides
    overrides: dict[str, ModelPricing] = {}
    try:
        raw = json.loads(_PRICING_FILE.read_text(encoding="utf-8"))
        if isinstance(raw, dict):
            for name, entry in raw.items():
                try:
                    overrides[str(name)] = ModelPricing(model_id=str(name), **(entry or {}))
                except Exception:
                    continue
    except (OSError, ValueError):
        pass
    _pricing_overrides = overrides
    return overrides


def reload_pricing_overrides() -> None:
    """Drop the cached overrides so the next get_pricing() re-reads the file."""
    global _pricing_overrides
    _pricing_overrides = None


def get_pricing(model_name: Optional[str]) -> Optional[ModelPricing]:
    """Resolve pricing for a model name; None when the model is unknown."""
    if not model_name:
        return None
    return _load_overrides().get(model_name) or DEFAULT_PRICING.get(model_name)


def calculate_cost(
    prompt_tokens: int,
    completion_tokens: int,
    pricing: ModelPricing,
) -> float:
    """Calculate cost in USD for given token usage and pricing."""
    prompt_cost = (prompt_tokens / 1_000_000) * pricing.prompt_cost_per_1m
    completion_cost = (completion_tokens / 1_000_000) * pricing.completion_cost_per_1m
    return prompt_cost + completion_cost


def usage_breakdown(usage: dict[str, Any], pricing: ModelPricing) -> CostBreakdown:
    """Turn one recorded usage dict (prompt/completion tokens) into a CostBreakdown."""
    prompt_tokens = int(usage.get("prompt_tokens") or 0)
    completion_tokens = int(usage.get("completion_tokens") or 0)
    return CostBreakdown(
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=prompt_tokens + completion_tokens,
        cost_usd=calculate_cost(prompt_tokens, completion_tokens, pricing),
    )


def build_eval_cost_section(
    usage_by_purpose: Optional[dict[str, dict[str, Any]]],
    model_configs: Optional[dict[str, Any]],
) -> Optional[dict[str, Any]]:
    """Per-purpose eval-LLM cost from recorded run usage; None when no usage.

    ``model_configs`` is the run summary's purpose → model snapshot map; a purpose
    whose model has no known pricing gets ``cost_usd: None`` (tokens still shown).
    """
    if not usage_by_purpose:
        return None
    model_configs = model_configs or {}
    items: list[dict[str, Any]] = []
    total_cost = 0.0
    has_any_cost = False
    for purpose, usage in sorted(usage_by_purpose.items()):
        model_name = (model_configs.get(purpose) or {}).get("model_name")
        pricing = get_pricing(model_name)
        breakdown = usage_breakdown(usage, pricing) if pricing else None
        if breakdown is not None:
            total_cost += breakdown.cost_usd
            has_any_cost = True
        items.append(
            {
                "purpose": purpose,
                "model_name": model_name,
                "prompt_tokens": breakdown.prompt_tokens if breakdown else int(usage.get("prompt_tokens") or 0),
                "completion_tokens": breakdown.completion_tokens if breakdown else int(usage.get("completion_tokens") or 0),
                "total_tokens": breakdown.total_tokens if breakdown else int(usage.get("total_tokens") or 0),
                "cost_usd": breakdown.cost_usd if breakdown else None,
            }
        )
    return {
        "by_purpose": items,
        "total_tokens": sum(i["total_tokens"] for i in items),
        "total_cost_usd": round(total_cost, 6) if has_any_cost else None,
    }