from __future__ import annotations

from sqlalchemy import func, or_, select
from sqlalchemy.ext.asyncio import AsyncSession

from app.enums.llm_billing import LlmBillingExecutionStatus
from app.models.llm_billing.call_attempt import LlmCallAttempt
from app.models.llm_billing.execution import LlmBillingExecution


def _normalized_execution_status(status: str) -> str:
    if status == LlmBillingExecutionStatus.PRE_DEDUCTED.value:
        return LlmBillingExecutionStatus.CHARGED.value
    return status


def _execution_status_filter(status: str):
    # 新接口只暴露 charged；查询时同时兼容历史 pre_deducted 数据。
    if status in {
        LlmBillingExecutionStatus.CHARGED.value,
        LlmBillingExecutionStatus.PRE_DEDUCTED.value,
    }:
        return or_(
            LlmBillingExecution.status == LlmBillingExecutionStatus.CHARGED.value,
            LlmBillingExecution.status == LlmBillingExecutionStatus.PRE_DEDUCTED.value,
        )
    return LlmBillingExecution.status == status


async def list_executions_with_calls(
    db: AsyncSession,
    *,
    page: int = 1,
    page_size: int = 20,
    scene_code: str | None = None,
    status: str | None = None,
    user_id: str | None = None,
) -> tuple[list[dict], int]:
    filters = []
    if scene_code:
        filters.append(LlmBillingExecution.scene_code == scene_code)
    if status:
        filters.append(_execution_status_filter(status))
    if user_id:
        filters.append(LlmBillingExecution.user_id == user_id)

    total_result = await db.execute(select(func.count(LlmBillingExecution.id)).where(*filters))
    total = int(total_result.scalar_one() or 0)
    result = await db.execute(
        select(LlmBillingExecution)
        .where(*filters)
        .order_by(LlmBillingExecution.created_at.desc(), LlmBillingExecution.id.desc())
        .offset((page - 1) * page_size)
        .limit(page_size)
    )
    executions = list(result.scalars().all())

    # 先收集 execution ids，再一次批查所有 CallAttempt 并按 execution_id 回填，避免 N+1。
    execution_ids = [item.id for item in executions]
    calls_by_execution: dict[str, list[LlmCallAttempt]] = {item_id: [] for item_id in execution_ids}
    if execution_ids:
        call_result = await db.execute(
            select(LlmCallAttempt)
            .where(LlmCallAttempt.billing_execution_id.in_(execution_ids))
            .order_by(LlmCallAttempt.billing_execution_id, LlmCallAttempt.call_sequence)
        )
        for call in call_result.scalars().all():
            calls_by_execution.setdefault(call.billing_execution_id, []).append(call)

    items: list[dict] = []
    for execution in executions:
        # 数据库物理字段 pre_deduct_credits 为历史命名；Admin/API 统一只暴露 charge_credits。
        data = {
            column.name: getattr(execution, column.name)
            for column in LlmBillingExecution.__table__.columns
            if column.name != "pre_deduct_credits"
        }
        data["charge_credits"] = execution.pre_deduct_credits
        data["status"] = _normalized_execution_status(execution.status)
        data["calls"] = calls_by_execution.get(execution.id, [])
        items.append(data)
    return items, total