"""Postgres budget ledger reconciliation.

Persists token usage to ``gateway.budget_usage`` (the durable source of truth)
via an idempotent upsert keyed by (key_id, period, period_start). The Redis
counter (``counter.py``) is the fast path; this ledger is what survives a Redis
flush and what ``show-usage`` reports against.
"""

from __future__ import annotations

import uuid

from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.ext.asyncio import AsyncSession

from neuronetz_gateway.budget.counter import period_start
from neuronetz_gateway.db.models import BudgetPeriod, BudgetUsage


class BudgetLedger:
    """Source-of-truth budget accounting in Postgres."""

    def __init__(self, session: AsyncSession) -> None:
        self._session = session

    async def record_usage(
        self, key_id: str, period: BudgetPeriod, tokens_in: int, tokens_out: int
    ) -> None:
        """Upsert usage into ``gateway.budget_usage`` for the active period.

        Uses an ``ON CONFLICT`` upsert so concurrent writers accumulate rather
        than clobber. ``requests`` increments by one per recorded request.
        """
        start = period_start(period)
        stmt = pg_insert(BudgetUsage).values(
            key_id=uuid.UUID(key_id) if isinstance(key_id, str) else key_id,
            period=period,
            period_start=start,
            tokens_in=tokens_in,
            tokens_out=tokens_out,
            requests=1,
        )
        stmt = stmt.on_conflict_do_update(
            index_elements=[
                BudgetUsage.key_id,
                BudgetUsage.period,
                BudgetUsage.period_start,
            ],
            set_={
                "tokens_in": BudgetUsage.tokens_in + stmt.excluded.tokens_in,
                "tokens_out": BudgetUsage.tokens_out + stmt.excluded.tokens_out,
                "requests": BudgetUsage.requests + stmt.excluded.requests,
            },
        )
        await self._session.execute(stmt)


__all__ = ["BudgetLedger"]