"""Async ``Database`` wrapper + the bare-minimum CRUD helpers.

Async to match the rest of the stack (aiohttp, uvicorn, claude-code-api).
psycopg3 has native async support — ``postgresql+psycopg://...`` works
with ``create_async_engine`` directly. SQLite goes through ``aiosqlite``
(``sqlite+aiosqlite://...``); user-facing config still uses the plain
``sqlite:///`` form and we normalise the URL here, so nothing leaks into
``.env`` / docker-compose.

No repository layer (PLAN §4.1 explicitly waives it). Helpers take an
``AsyncSession`` so callers can batch operations into one transaction
(e.g. touch ``last_used_at`` + write an audit line on the same request).
"""

from __future__ import annotations

import json
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

from sqlalchemy import inspect, text
from sqlalchemy.ext.asyncio import create_async_engine
from sqlmodel import SQLModel, col, select
from sqlmodel.ext.asyncio.session import AsyncSession

from beaver_gateway.storage.models import AuditLog, Token, Usage

if TYPE_CHECKING:
    from collections.abc import Sequence

    from sqlalchemy import Connection
    from sqlalchemy.ext.asyncio import AsyncEngine


def _utcnow() -> datetime:
    return datetime.now(UTC)


def _to_async_url(url: str) -> str:
    """Promote the user-facing sync URL to its async driver variant.

    Users write ``sqlite:///gw.db`` or ``postgresql://...`` in ``.env``;
    we translate to ``sqlite+aiosqlite://`` / ``postgresql+psycopg://``
    so they don't have to know which driver we use internally.
    """
    if url.startswith(("sqlite+aiosqlite://", "postgresql+psycopg://")):
        return url
    if url.startswith("sqlite://"):
        return "sqlite+aiosqlite://" + url[len("sqlite://") :]
    if url.startswith("postgresql://"):
        return "postgresql+psycopg://" + url[len("postgresql://") :]
    if url.startswith("postgres://"):
        return "postgresql+psycopg://" + url[len("postgres://") :]
    return url


class Database:
    """Owner of the async SQLAlchemy engine.

    Construct once in ``cli`` from ``settings.database_url``, ``await``
    :meth:`create_all` at startup, hand the instance to
    :class:`GatewayRuntime`. Frontends grab sessions via
    :meth:`session` (an async context manager).

    ``echo=False`` keeps SQL out of INFO logs (admin UI is the
    user-facing view); ``connect_args={"check_same_thread": False}``
    isn't needed under async sqlite — aiosqlite already runs each
    connection on a dedicated thread.
    """

    __slots__ = ("_engine",)

    def __init__(self, url: str) -> None:
        self._engine: AsyncEngine = create_async_engine(_to_async_url(url), echo=False)

    async def create_all(self) -> None:
        """Create missing tables, then add columns models gained since.

        Nullable columns only - that is the whole migration story until
        Alembic is worth it.
        """
        async with self._engine.begin() as conn:
            await conn.run_sync(SQLModel.metadata.create_all)
            await conn.run_sync(_add_missing_columns)

    def session(self) -> AsyncSession:
        """Open a fresh :class:`AsyncSession` (use as ``async with``).

        Callers commit explicitly; we don't auto-commit on exit so a
        request that fails mid-flight rolls back by simply not
        committing.
        """
        return AsyncSession(self._engine, expire_on_commit=False)

    async def dispose(self) -> None:
        """Close the engine's connection pool. Idempotent."""
        await self._engine.dispose()


def _add_missing_columns(conn: Connection) -> None:
    inspector = inspect(conn)
    for table in SQLModel.metadata.sorted_tables:
        existing = {c["name"] for c in inspector.get_columns(table.name)}
        for column in table.columns:
            if column.name in existing:
                continue
            kind = column.type.compile(conn.dialect)
            ddl = f"ALTER TABLE {table.name} ADD COLUMN {column.name} {kind}"
            default = _column_default(column)
            if default is not None:
                ddl += f" DEFAULT {default}"
            conn.execute(text(ddl))


def _column_default(column: Any) -> str | None:
    """Literal for ``ADD COLUMN ... DEFAULT`` so old rows get the model default."""
    server_default = getattr(column, "server_default", None)
    if server_default is not None:
        return str(server_default.arg.text)
    default = getattr(column, "default", None)
    if default is None or default.is_callable:
        return None
    value = default.arg
    if isinstance(value, bool):
        return "TRUE" if value else "FALSE"
    if isinstance(value, int | float):
        return repr(value)
    if isinstance(value, str):
        return "'" + value.replace("'", "''") + "'"
    return None


# ---- Token CRUD ---------------------------------------------------------


async def list_active_tokens(session: AsyncSession) -> Sequence[Token]:
    """Return every non-revoked token (Phase 4.2 seeds the cache from this)."""
    stmt = select(Token).where(Token.revoked_at.is_(None))  # ty: ignore[unresolved-attribute]
    result = await session.exec(stmt)
    return result.all()


async def list_tokens(
    session: AsyncSession, *, include_revoked: bool = False
) -> Sequence[Token]:
    """Return tokens ordered newest-first (Phase 4.3 admin table)."""
    stmt = select(Token).order_by(Token.created_at.desc())  # ty: ignore[unresolved-attribute]
    if not include_revoked:
        stmt = stmt.where(Token.revoked_at.is_(None))  # ty: ignore[unresolved-attribute]
    result = await session.exec(stmt)
    return result.all()


async def create_token(
    session: AsyncSession, *, name: str, scope: str, hashed_value: str
) -> Token:
    """Persist a new token. Caller hashes the plaintext before passing it in."""
    row = Token(name=name, scope=scope, hashed_value=hashed_value)
    session.add(row)
    await session.commit()
    await session.refresh(row)
    return row


async def revoke_token(session: AsyncSession, *, token_id: int) -> bool:
    """Mark a token revoked. Returns ``False`` if no such row."""
    row = await session.get(Token, token_id)
    if row is None or row.revoked_at is not None:
        return False
    row.revoked_at = _utcnow()
    session.add(row)
    await session.commit()
    return True


async def touch_token(session: AsyncSession, *, token_id: int) -> None:
    """Bump ``last_used_at``. Phase 4.2 batches these — not per-request."""
    row = await session.get(Token, token_id)
    if row is None:
        return
    row.last_used_at = _utcnow()
    session.add(row)
    await session.commit()


# ---- Audit --------------------------------------------------------------


async def append_audit(
    session: AsyncSession,
    *,
    actor: str,
    kind: str,
    agent_name: str | None = None,
    detail: dict[str, Any] | None = None,
) -> AuditLog:
    """Append-only insert. ``detail`` JSON-serialised here, not by callers."""
    row = AuditLog(
        actor=actor,
        kind=kind,
        agent_name=agent_name,
        detail_json=json.dumps(detail or {}, separators=(",", ":")),
    )
    session.add(row)
    await session.commit()
    await session.refresh(row)
    return row


async def list_audit_records(
    session: AsyncSession, *, limit: int = 50, before_id: int | None = None
) -> Sequence[AuditLog]:
    """Return audit entries newest-first, optionally paginated by id cursor.

    ``before_id`` is a forward-only cursor: pass the smallest id from the
    current page to fetch the next slice. Cheap because ``id`` is the
    primary key (ordered insert).
    """
    stmt = select(AuditLog).order_by(AuditLog.id.desc()).limit(limit)  # ty: ignore[unresolved-attribute]
    if before_id is not None:
        stmt = stmt.where(AuditLog.id < before_id)  # ty: ignore[unsupported-operator]
    result = await session.exec(stmt)
    return result.all()


# ---- Usage --------------------------------------------------------------


async def append_usage(session: AsyncSession, row: Usage) -> None:
    """Persist a turn's usage with ``cost_usd`` / ``model_usage`` as per-turn deltas.

    ``row`` arrives with the SDK's cumulative ``total_cost_usd`` and
    ``model_usage`` in ``cost_usd`` / ``model_usage``; they move to the
    ``*_total`` columns and the stored values become the difference against
    the previous row of the same session (see :func:`usage_deltas`).
    """
    prev = await _last_usage(session, row.session_id)
    usage_deltas(row, prev)
    session.add(row)
    await session.commit()


async def _last_usage(session: AsyncSession, session_id: str | None) -> Usage | None:
    if not session_id:
        return None
    stmt = (
        select(Usage)
        .where(col(Usage.session_id) == session_id)
        .order_by(col(Usage.ts).desc(), col(Usage.id).desc())
        .limit(1)
    )
    return (await session.exec(stmt)).first()


# Counters inside ``ResultMessage.model_usage[model]``; everything else there
# (provider, costBasis, contextWindow, maxOutputTokens, ...) is a constant.
_MODEL_USAGE_COUNTERS = (
    "inputTokens",
    "outputTokens",
    "cacheReadInputTokens",
    "cacheCreationInputTokens",
    "webSearchRequests",
    "costUSD",
)


def usage_deltas(row: Usage, prev: Usage | None) -> None:
    """Turn the SDK's cumulative cost counters on ``row`` into per-turn deltas.

    ``total_cost_usd`` and ``model_usage`` in a ``ResultMessage`` are running
    totals of the claude process: they grow with every turn of the session and
    start over when the process is respawned (rotation, resume after a
    restart). A row fresh from the SDK (``cost_usd_total is None``) first moves
    the raw values into ``cost_usd_total`` / ``model_usage_total``; then
    ``cost_usd`` / ``model_usage`` become ``current - previous`` against
    ``prev`` (the last row of the same session), or the raw value when there
    is no previous row or the counter went backwards - a fresh process.
    Idempotent on rows that already carry their totals, which is what the
    backfill relies on.
    """
    if row.cost_usd_total is None and row.model_usage_total is None:
        row.cost_usd_total = row.cost_usd
        row.model_usage_total = row.model_usage
    prev_cost = prev.cost_usd_total if prev is not None else None
    prev_models = prev.model_usage_total if prev is not None else None
    row.cost_usd = _counter_delta(row.cost_usd_total, prev_cost)
    row.model_usage = _model_usage_delta(row.model_usage_total, prev_models)


def _counter_delta(current: float | None, previous: float | None) -> float | None:
    if current is None:
        return None
    if previous is None or current < previous:
        return current
    return round(current - previous, 6)


def _model_usage_delta(
    current: dict[str, Any] | None, previous: dict[str, Any] | None
) -> dict[str, Any] | None:
    if current is None:
        return None
    out: dict[str, Any] = {}
    for model, fields in current.items():
        if not isinstance(fields, dict):
            out[model] = fields
            continue
        before = (previous or {}).get(model)
        before = before if isinstance(before, dict) else {}
        delta = dict(fields)
        for key in _MODEL_USAGE_COUNTERS:
            if key not in fields:
                continue
            cur, prev = fields.get(key), before.get(key)
            if not isinstance(cur, int | float) or isinstance(cur, bool):
                continue
            prev_num = (
                prev
                if isinstance(prev, int | float) and not isinstance(prev, bool)
                else None
            )
            value = _counter_delta(cur, prev_num)
            delta[key] = (
                round(value or 0, 6) if isinstance(cur, float) else int(value or 0)
            )
        out[model] = delta
    return out


async def backfill_usage_deltas(session: AsyncSession) -> int:
    """Rewrite ``cost_usd`` / ``model_usage`` of every row as per-turn deltas.

    Rows written before the ``*_total`` columns existed hold the SDK's
    cumulative values in ``cost_usd`` / ``model_usage``; they are moved to the
    totals first. Rows are replayed per session in time order. Idempotent:
    a second run rewrites nothing. Returns the number of rows changed.
    """
    stmt = select(Usage).order_by(col(Usage.session_id), col(Usage.ts), col(Usage.id))
    rows = (await session.exec(stmt)).all()
    last_by_session: dict[str, Usage] = {}
    changed = 0
    for row in rows:
        before = (
            row.cost_usd,
            row.model_usage,
            row.cost_usd_total,
            row.model_usage_total,
        )
        prev = last_by_session.get(row.session_id) if row.session_id else None
        usage_deltas(row, prev)
        after = (
            row.cost_usd,
            row.model_usage,
            row.cost_usd_total,
            row.model_usage_total,
        )
        if after != before:
            session.add(row)
            changed += 1
        if row.session_id:
            last_by_session[row.session_id] = row
    await session.commit()
    return changed


__all__ = [
    "Database",
    "append_audit",
    "append_usage",
    "backfill_usage_deltas",
    "create_token",
    "list_active_tokens",
    "list_audit_records",
    "list_tokens",
    "revoke_token",
    "touch_token",
    "usage_deltas",
]