"""``cost_usd`` / ``model_usage`` are per-turn deltas of the SDK's running totals."""

import tempfile
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

from sqlmodel import col, select

from beaver_gateway.storage import Database, Usage, append_usage, backfill_usage_deltas


def _mu(cost: float, out: int, cache_read: int) -> dict[str, Any]:
    return {
        "claude-opus-5": {
            "costUSD": cost,
            "inputTokens": 1,
            "outputTokens": out,
            "cacheReadInputTokens": cache_read,
            "cacheCreationInputTokens": 0,
            "webSearchRequests": 0,
            "contextWindow": 1_000_000,
            "provider": "firstParty",
        }
    }


def _row(
    ts: datetime, sid: str | None, cost: float, out: int, cache_read: int
) -> Usage:
    return Usage(
        ts=ts,
        agent_name="d",
        session_id=sid,
        model="claude-opus-5",
        output_tokens=out,
        cache_read_tokens=cache_read,
        cost_usd=cost,
        model_usage=_mu(cost, out, cache_read),
    )


async def _db() -> Database:
    path = Path(tempfile.mkdtemp(prefix="beaver-usage-")) / "u.db"
    db = Database(f"sqlite:///{path}")
    await db.create_all()
    return db


async def _all(db: Database) -> list[Usage]:
    async with db.session() as session:
        return list((await session.exec(select(Usage).order_by(col(Usage.id)))).all())


async def test_append_usage_stores_deltas_and_keeps_totals() -> None:
    db = await _db()
    t0 = datetime.now(UTC).replace(tzinfo=None)
    async with db.session() as session:
        await append_usage(session, _row(t0, "s1", 0.5, 100, 1000))
        await append_usage(
            session, _row(t0 + timedelta(minutes=1), "s1", 0.8, 160, 2500)
        )
        # the claude process was respawned: the counters start over
        await append_usage(session, _row(t0 + timedelta(minutes=2), "s1", 0.2, 30, 400))
        # another session is a separate counter
        await append_usage(session, _row(t0 + timedelta(minutes=3), "s2", 0.7, 50, 900))
        # no session id - nothing to diff against
        await append_usage(session, _row(t0 + timedelta(minutes=4), None, 0.3, 10, 100))
    rows = await _all(db)
    assert [r.cost_usd for r in rows] == [0.5, 0.3, 0.2, 0.7, 0.3]
    assert [r.cost_usd_total for r in rows] == [0.5, 0.8, 0.2, 0.7, 0.3]
    second = rows[1].model_usage
    assert second is not None
    opus = second["claude-opus-5"]
    assert (opus["costUSD"], opus["outputTokens"], opus["cacheReadInputTokens"]) == (
        0.3,
        60,
        1500,
    )
    assert (opus["contextWindow"], opus["provider"]) == (1_000_000, "firstParty")
    assert rows[1].model_usage_total == _mu(0.8, 160, 2500)
    assert rows[2].model_usage is not None
    assert rows[2].model_usage["claude-opus-5"]["outputTokens"] == 30
    await db.dispose()


async def test_backfill_rewrites_legacy_cumulative_rows_once() -> None:
    db = await _db()
    t0 = datetime.now(UTC).replace(tzinfo=None)
    async with db.session() as session:
        # legacy rows: the raw running total sits in cost_usd / model_usage
        for i, cost in enumerate((1.0, 1.5, 2.5, 0.4)):
            session.add(_row(t0 + timedelta(minutes=i), "s1", cost, 10 * (i + 1), 100))
        session.add(_row(t0 + timedelta(minutes=9), "s9", 0.9, 5, 50))
        await session.commit()
    async with db.session() as session:
        assert await backfill_usage_deltas(session) == 5
    rows = await _all(db)
    assert [r.cost_usd for r in rows] == [1.0, 0.5, 1.0, 0.4, 0.9]
    assert [r.cost_usd_total for r in rows] == [1.0, 1.5, 2.5, 0.4, 0.9]
    assert rows[2].model_usage is not None
    assert rows[2].model_usage["claude-opus-5"]["outputTokens"] == 10
    async with db.session() as session:
        assert await backfill_usage_deltas(session) == 0
    await db.dispose()