"""Tests for runtime spend ledger."""

from __future__ import annotations

from dataclasses import asdict

import pytest

from logos_adk.spend_ledger import (
    InMemorySpendLedger,
    PostgreSQLSpendLedger,
    SQLiteSpendLedger,
    SpendRecord,
    _build_budget_status,
    _row_value,
)


def test_in_memory_spend_ledger_tracks_workspace_and_session_spend():
    ledger = InMemorySpendLedger()
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s1",
            model="gpt-4o-mini",
            amount_usd=0.002,
        )
    )
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s2",
            model="claude-sonnet-4-20250514",
            amount_usd=0.004,
        )
    )

    assert ledger.spent(agent_name="router", workspace_id="acme") == 0.006
    assert ledger.spent(agent_name="router", workspace_id="acme", session_id="s1") == 0.002

    status = ledger.budget_status(
        agent_name="router",
        workspace_id="acme",
        session_id="s1",
        workspace_soft_limit_usd=0.005,
        session_soft_limit_usd=0.001,
        threshold_ratio=0.8,
    )
    assert status["workspace"].near_limit is True
    assert status["session"].near_limit is True


def test_in_memory_spend_ledger_clamps_negative_amounts_and_lists_latest_first():
    ledger = InMemorySpendLedger()
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s1",
            model="gpt-4o-mini",
            amount_usd=-5,
            created_at="2024-01-01T00:00:00+00:00",
        )
    )
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s2",
            model="gpt-4.1-mini",
            amount_usd=0.004,
            created_at="2024-01-02T00:00:00+00:00",
        )
    )
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="beta",
            session_id="s3",
            model="claude",
            amount_usd=0.003,
            created_at="2024-01-03T00:00:00+00:00",
        )
    )

    latest_two = ledger.list_records(limit=2)

    assert ledger.spent(agent_name="router", workspace_id="acme") == 0.004
    assert [record.session_id for record in latest_two] == ["s3", "s2"]
    assert latest_two[1].amount_usd == 0.004


def test_spend_record_uses_fresh_default_timestamps():
    first = SpendRecord(agent_name="router", model="gpt", amount_usd=0.1)
    second = SpendRecord(agent_name="router", model="gpt", amount_usd=0.2)

    assert first.created_at
    assert second.created_at
    assert asdict(first)["created_at"] != ""
    assert asdict(second)["created_at"] != ""


def test_sqlite_spend_ledger_survives_restart(tmp_path):
    path = tmp_path / "spend.db"
    ledger = SQLiteSpendLedger(str(path))
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s1",
            model="gpt-4o-mini",
            amount_usd=0.003,
        )
    )

    restarted = SQLiteSpendLedger(str(path))
    assert restarted.spent(agent_name="router", workspace_id="acme") == 0.003


def test_sqlite_spend_ledger_clamps_negative_amount_and_lists_records(tmp_path):
    path = tmp_path / "spend.db"
    ledger = SQLiteSpendLedger(str(path))
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s1",
            model="gpt-4o-mini",
            amount_usd=-1.0,
            created_at="2024-01-01T00:00:00+00:00",
        )
    )
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id=None,
            session_id=None,
            model="gpt-4.1-mini",
            amount_usd=0.25,
            trace_id="trace-2",
            created_at="2024-01-02T00:00:00+00:00",
        )
    )

    records = ledger.list_records(limit=10)

    assert ledger.spent(agent_name="router", workspace_id="acme") == 0.0
    assert len(records) == 2
    assert records[0].trace_id == "trace-2"
    assert records[0].workspace_id == ""
    assert records[0].session_id == ""
    assert records[1].amount_usd == 0.0


class _FakeCursor:
    def __init__(self, records):
        self.records = records
        self._fetchone = None
        self._fetchall = []

    def execute(self, query, params=None):
        normalized = " ".join(str(query).split())
        self._fetchone = None
        self._fetchall = []

        if normalized.startswith("INSERT INTO spend_ledger_records"):
            self.records.append(
                {
                    "id": params[0],
                    "agent_name": params[1],
                    "workspace_id": params[2],
                    "session_id": params[3],
                    "model": params[4],
                    "amount_usd": params[5],
                    "trace_id": params[6],
                    "created_at": params[7],
                }
            )
            return

        if "SELECT COALESCE(SUM(amount_usd), 0) AS total" in normalized:
            filtered = [record for record in self.records if record["agent_name"] == params[0]]
            param_index = 1
            if "workspace_id = %s" in normalized:
                filtered = [record for record in filtered if record["workspace_id"] == params[param_index]]
                param_index += 1
            if "session_id = %s" in normalized:
                filtered = [record for record in filtered if record["session_id"] == params[param_index]]
            self._fetchone = {"total": sum(record["amount_usd"] for record in filtered)}
            return

        if normalized.startswith("SELECT * FROM spend_ledger_records ORDER BY created_at DESC LIMIT %s"):
            limit = int(params[0])
            ordered = sorted(self.records, key=lambda item: item["created_at"], reverse=True)
            self._fetchall = ordered[:limit]

    def fetchone(self):
        return self._fetchone

    def fetchall(self):
        return self._fetchall

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


class _FakeConnection:
    def __init__(self):
        self.records = []
        self.commits = 0

    def cursor(self):
        return _FakeCursor(self.records)

    def commit(self):
        self.commits += 1

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False


def test_postgresql_spend_ledger_constructor_resolves_env_dsn(monkeypatch):
    monkeypatch.setenv("DATABASE_URL", "postgresql://env/spend")

    ledger = PostgreSQLSpendLedger()

    assert ledger.dsn == "postgresql://env/spend"


def test_postgresql_spend_ledger_requires_dsn_when_env_missing(monkeypatch):
    monkeypatch.delenv("DATABASE_URL", raising=False)
    monkeypatch.delenv("LOGOS_ADK_DATABASE_URL", raising=False)

    with pytest.raises(ValueError):
        PostgreSQLSpendLedger()


def test_postgresql_spend_ledger_tracks_spend_with_fake_backend(monkeypatch):
    connection = _FakeConnection()
    ledger = PostgreSQLSpendLedger("postgresql://user:pass@localhost/spend")
    ledger._connect = lambda: connection
    monkeypatch.setattr(PostgreSQLSpendLedger, "_ensure_schema", lambda self, conn: None)

    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s1",
            model="gpt-4o-mini",
            amount_usd=0.002,
        )
    )
    ledger.record(
        SpendRecord(
            agent_name="router",
            workspace_id="acme",
            session_id="s2",
            model="claude-sonnet-4-20250514",
            amount_usd=0.004,
        )
    )

    assert ledger.spent(agent_name="router", workspace_id="acme") == 0.006
    assert ledger.spent(agent_name="router", workspace_id="acme", session_id="s1") == 0.002
    records = ledger.list_records(limit=10)

    assert len(records) == 2
    status = ledger.budget_status(
        agent_name="router",
        workspace_id="acme",
        session_id="s1",
        workspace_soft_limit_usd=0.005,
        session_soft_limit_usd=0.001,
        threshold_ratio=0.8,
    )
    assert status["workspace"].near_limit is True
    assert status["session"].near_limit is True


def test_postgresql_helpers_cover_row_lookup_and_budget_status_without_limits():
    assert _row_value({"total": 1.5}, "total", 0) == 1.5
    assert _row_value(("a", "b", "c"), "missing", 2) == "c"

    empty = _build_budget_status(
        workspace_spent=0.0,
        session_spent=0.0,
        agent_name="router",
        workspace_id="acme",
        session_id="s1",
        workspace_soft_limit_usd=None,
        session_soft_limit_usd=0.0,
        threshold_ratio=0.8,
    )
    populated = _build_budget_status(
        workspace_spent=4.0,
        session_spent=1.0,
        agent_name="router",
        workspace_id="acme",
        session_id="s1",
        workspace_soft_limit_usd=5.0,
        session_soft_limit_usd=2.0,
        threshold_ratio=0.8,
    )

    assert empty == {}
    assert populated["workspace"].ratio == 0.8
    assert populated["workspace"].near_limit is True
    assert populated["workspace"].limit_key == "acme"
    assert populated["session"].ratio == 0.5
    assert populated["session"].near_limit is False
    assert populated["session"].limit_key == "s1"
