"""SQL backend for UsageLoggerMiddleware.

Persists :class:`~agentbyte.middleware.usage_logger.UsageEvent` rows to any
SQLAlchemy-supported database using SQLModel for table definitions and
SQLAlchemy's async session for writes.

Requires the ``sql`` optional extra::

    uv sync --extra sql

Supports both PostgreSQL (via asyncpg) and SQLite (via aiosqlite).

Usage::

    from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
    from sqlalchemy.orm import sessionmaker
    from agentbyte.middleware.sql_usage import SqlUsageBackend, create_usage_tables

    engine = create_async_engine("postgresql+asyncpg://user:pass@host/db")
    AsyncSessionLocal = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)

    async def get_session():
        async with AsyncSessionLocal() as session:
            yield session

    await create_usage_tables(engine)
    backend = SqlUsageBackend(session_factory=get_session)
"""

from __future__ import annotations

from collections.abc import AsyncGenerator, Callable
from datetime import datetime, timezone
from typing import Optional

from sqlalchemy.ext.asyncio import AsyncEngine, AsyncSession
from sqlmodel import Field, SQLModel

from .usage_logger import UsageEvent, UsagePersistenceBackend


class LlmUsageEvent(SQLModel, table=True):
    """SQLModel table that stores one row per LLM call.

    This class is internal to :class:`SqlUsageBackend`.  Callers who need
    direct DB access can import it for read-only queries.
    """

    __tablename__ = "llm_usage_events"

    id: Optional[int] = Field(default=None, primary_key=True)
    run_id: Optional[str] = Field(default=None, index=True)
    session_id: Optional[str] = Field(default=None, index=True)
    user_id: Optional[str] = Field(default=None)
    identifier: Optional[str] = Field(default=None, index=True)
    agent_name: str
    model: str
    duration_ms: int = Field(default=0)
    llm_calls: int = Field(default=0)
    tokens_input: int = Field(default=0)
    tokens_output: int = Field(default=0)
    tool_calls: int = Field(default=0)
    memory_operations: int = Field(default=0)
    cost_estimate: Optional[float] = Field(default=None)
    recorded_at: datetime = Field(
        default_factory=lambda: datetime.now(timezone.utc),
    )


async def create_usage_tables(engine: AsyncEngine) -> None:
    """Create the ``llm_usage_events`` table if it does not exist.

    Call this once at application startup before any agents run::

        await create_usage_tables(engine)
    """
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)


class SqlUsageBackend(UsagePersistenceBackend):
    """Persists :class:`~agentbyte.middleware.usage_logger.UsageEvent` rows
    to a SQL database via an async SQLAlchemy session.

    The caller is responsible for creating the engine and session factory.
    This backend holds only a reference to the factory — it does not own the
    connection pool.

    Args:
        session_factory: Zero-argument async generator factory that yields one
            :class:`~sqlalchemy.ext.asyncio.AsyncSession` per call.  The
            factory should handle its own cleanup (e.g. via ``async with``).

    Example::

        from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
        from sqlalchemy.orm import sessionmaker

        engine = create_async_engine("postgresql+asyncpg://user:pass@host/db")
        AsyncSessionLocal = sessionmaker(
            engine, class_=AsyncSession, expire_on_commit=False
        )

        async def get_session():
            async with AsyncSessionLocal() as session:
                yield session

        backend = SqlUsageBackend(session_factory=get_session)
    """

    def __init__(
        self,
        session_factory: Callable[[], AsyncGenerator[AsyncSession, None]],
    ) -> None:
        self._session_factory = session_factory

    async def write(self, event: UsageEvent) -> None:
        gen = self._session_factory()
        session: AsyncSession = await gen.__anext__()
        try:
            row = LlmUsageEvent(
                run_id=event.run_id,
                session_id=event.session_id,
                user_id=event.user_id,
                identifier=event.identifier,
                agent_name=event.agent_name,
                model=event.model,
                duration_ms=event.duration_ms,
                llm_calls=event.llm_calls,
                tokens_input=event.tokens_input,
                tokens_output=event.tokens_output,
                tool_calls=event.tool_calls,
                memory_operations=event.memory_operations,
                cost_estimate=event.cost_estimate,
                recorded_at=event.recorded_at,
            )
            session.add(row)
            await session.commit()
        finally:
            await gen.aclose()
