from __future__ import annotations

from decimal import Decimal
from typing import Any

from sqlalchemy import case, func, select

from supply_infra.db.models.llm_usage_event import LlmUsageEvent
from supply_infra.db.repositories.base import BaseRepository


class LlmUsageRepository(BaseRepository[LlmUsageEvent]):
    model = LlmUsageEvent

    def create_event(
        self,
        *,
        biz_dt: str,
        provider: str,
        model: str | None,
        agent_name: str | None,
        run_id: str | None,
        iteration: int | None,
        prompt_tokens: int,
        completion_tokens: int,
        total_tokens: int,
        cost_usd: Decimal | None,
    ) -> LlmUsageEvent:
        entity = LlmUsageEvent(
            biz_dt=biz_dt,
            provider=provider,
            model=model,
            agent_name=agent_name,
            run_id=run_id,
            iteration=iteration,
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            total_tokens=total_tokens,
            cost_usd=cost_usd,
        )
        return self.add(entity)

    def summarize_daily(
        self,
        *,
        start_dt: str,
        end_dt: str,
    ) -> list[dict[str, Any]]:
        """按业务日汇总费用与调用量。"""
        cost_sum = func.coalesce(func.sum(LlmUsageEvent.cost_usd), 0)
        stmt = (
            select(
                LlmUsageEvent.biz_dt,
                func.count(LlmUsageEvent.id).label("call_count"),
                func.coalesce(func.sum(LlmUsageEvent.prompt_tokens), 0).label("prompt_tokens"),
                func.coalesce(func.sum(LlmUsageEvent.completion_tokens), 0).label(
                    "completion_tokens"
                ),
                func.coalesce(func.sum(LlmUsageEvent.total_tokens), 0).label("total_tokens"),
                cost_sum.label("total_cost_usd"),
                func.sum(case((LlmUsageEvent.cost_usd.is_(None), 1), else_=0)).label(
                    "missing_cost_count"
                ),
            )
            .where(
                LlmUsageEvent.biz_dt >= start_dt,
                LlmUsageEvent.biz_dt <= end_dt,
            )
            .group_by(LlmUsageEvent.biz_dt)
            .order_by(LlmUsageEvent.biz_dt.asc())
        )
        rows = self.session.execute(stmt).all()
        return [
            {
                "biz_dt": row.biz_dt,
                "call_count": int(row.call_count or 0),
                "prompt_tokens": int(row.prompt_tokens or 0),
                "completion_tokens": int(row.completion_tokens or 0),
                "total_tokens": int(row.total_tokens or 0),
                "total_cost_usd": float(row.total_cost_usd or 0),
                "missing_cost_count": int(row.missing_cost_count or 0),
            }
            for row in rows
        ]

    def breakdown_by_agent(
        self,
        *,
        start_dt: str,
        end_dt: str,
    ) -> list[dict[str, Any]]:
        cost_sum = func.coalesce(func.sum(LlmUsageEvent.cost_usd), 0)
        agent_label = func.coalesce(LlmUsageEvent.agent_name, "unknown")
        stmt = (
            select(
                agent_label.label("agent_name"),
                func.count(LlmUsageEvent.id).label("call_count"),
                func.coalesce(func.sum(LlmUsageEvent.total_tokens), 0).label("total_tokens"),
                cost_sum.label("total_cost_usd"),
            )
            .where(
                LlmUsageEvent.biz_dt >= start_dt,
                LlmUsageEvent.biz_dt <= end_dt,
            )
            .group_by(agent_label)
            .order_by(cost_sum.desc(), func.count(LlmUsageEvent.id).desc())
        )
        rows = self.session.execute(stmt).all()
        return [
            {
                "agent_name": row.agent_name,
                "call_count": int(row.call_count or 0),
                "total_tokens": int(row.total_tokens or 0),
                "total_cost_usd": float(row.total_cost_usd or 0),
            }
            for row in rows
        ]

    def breakdown_by_model(
        self,
        *,
        start_dt: str,
        end_dt: str,
    ) -> list[dict[str, Any]]:
        cost_sum = func.coalesce(func.sum(LlmUsageEvent.cost_usd), 0)
        model_label = func.coalesce(LlmUsageEvent.model, "unknown")
        stmt = (
            select(
                model_label.label("model"),
                func.count(LlmUsageEvent.id).label("call_count"),
                func.coalesce(func.sum(LlmUsageEvent.total_tokens), 0).label("total_tokens"),
                cost_sum.label("total_cost_usd"),
            )
            .where(
                LlmUsageEvent.biz_dt >= start_dt,
                LlmUsageEvent.biz_dt <= end_dt,
            )
            .group_by(model_label)
            .order_by(cost_sum.desc(), func.count(LlmUsageEvent.id).desc())
        )
        rows = self.session.execute(stmt).all()
        return [
            {
                "model": row.model,
                "call_count": int(row.call_count or 0),
                "total_tokens": int(row.total_tokens or 0),
                "total_cost_usd": float(row.total_cost_usd or 0),
            }
            for row in rows
        ]

    def list_recent_events(
        self,
        *,
        start_dt: str | None = None,
        end_dt: str | None = None,
        limit: int = 50,
    ) -> list[LlmUsageEvent]:
        stmt = select(LlmUsageEvent)
        if start_dt:
            stmt = stmt.where(LlmUsageEvent.biz_dt >= start_dt)
        if end_dt:
            stmt = stmt.where(LlmUsageEvent.biz_dt <= end_dt)
        stmt = stmt.order_by(LlmUsageEvent.id.desc()).limit(limit)
        return list(self.session.scalars(stmt).all())
