"""Persist LLM usage events into MySQL for billing dashboards."""

from __future__ import annotations

import logging
from datetime import datetime
from decimal import Decimal, InvalidOperation
from typing import Any
from zoneinfo import ZoneInfo

from supply_infra.db.repositories.llm_usage_repo import LlmUsageRepository
from supply_infra.db.session import get_session

logger = logging.getLogger(__name__)
_SHANGHAI = ZoneInfo("Asia/Shanghai")


def _today_biz_dt() -> str:
    return datetime.now(_SHANGHAI).strftime("%Y%m%d")


def _as_int(value: Any) -> int:
    if value is None:
        return 0
    try:
        return int(value)
    except (TypeError, ValueError):
        return 0


def _as_cost(value: Any) -> Decimal | None:
    if value is None:
        return None
    try:
        return Decimal(str(value))
    except (InvalidOperation, ValueError, TypeError):
        return None


def record_llm_usage_to_db(payload: dict[str, Any]) -> None:
    """Write one LLM call's usage/cost into ``llm_usage_event``."""
    usage = payload.get("usage") or {}
    prompt_tokens = _as_int(usage.get("prompt_tokens"))
    completion_tokens = _as_int(usage.get("completion_tokens"))
    total_tokens = _as_int(usage.get("total_tokens"))
    if total_tokens <= 0:
        total_tokens = prompt_tokens + completion_tokens

    cost_usd = _as_cost(usage.get("cost"))
    if cost_usd is None and isinstance(usage.get("cost_details"), dict):
        # Some OpenRouter responses nest cost under cost_details.upstream_inference_cost
        details = usage["cost_details"]
        cost_usd = _as_cost(details.get("upstream_inference_cost") or details.get("total_cost"))

    # Skip empty records (e.g. stream path without usage).
    if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0 and cost_usd is None:
        return

    biz_dt = str(payload.get("biz_dt") or _today_biz_dt())
    with get_session() as session:
        repo = LlmUsageRepository(session)
        repo.create_event(
            biz_dt=biz_dt,
            provider=str(payload.get("provider") or "openrouter"),
            model=(str(payload["model"]) if payload.get("model") else None),
            agent_name=(str(payload["agent_name"]) if payload.get("agent_name") else None),
            run_id=(str(payload["run_id"]) if payload.get("run_id") else None),
            iteration=_as_int(payload.get("iteration")) if payload.get("iteration") is not None else None,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=total_tokens,
            cost_usd=cost_usd,
        )
    logger.debug(
        "recorded llm usage run_id=%s model=%s cost_usd=%s tokens=%s",
        payload.get("run_id"),
        payload.get("model"),
        cost_usd,
        total_tokens,
    )
