from __future__ import annotations

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

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


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(LlmBillingExecution.status == 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())
    ids = [item.id for item in executions]
    calls_by_execution: dict[str, list[LlmCallAttempt]] = {item_id: [] for item_id in ids}
    if ids:
        call_result = await db.execute(
            select(LlmCallAttempt)
            .where(LlmCallAttempt.billing_execution_id.in_(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 = []
    for execution in executions:
        data = {
            column.name: getattr(execution, column.name)
            for column in LlmBillingExecution.__table__.columns
        }
        data["calls"] = calls_by_execution.get(execution.id, [])
        items.append(data)
    return items, total