"""Persistent spend ledger for runtime budget-aware model routing."""

from __future__ import annotations

from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
import sqlite3
from typing import Any, Protocol
from uuid import uuid4

from logos_adk.persistence import (
    get_postgresql_schema_version,
    get_sqlite_schema_version,
    set_postgresql_schema_version,
    set_sqlite_schema_version,
)
from logos_adk.postgres_utils import import_psycopg, resolve_database_url


SPEND_LEDGER_SCHEMA_KIND = "spend_ledger"
SPEND_LEDGER_SCHEMA_VERSION = 1


def _segment(value: str | None) -> str:
    return value or ""


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


@dataclass
class SpendRecord:
    agent_name: str
    model: str
    amount_usd: float
    workspace_id: str | None = None
    session_id: str | None = None
    trace_id: str | None = None
    created_at: str = field(default_factory=_utcnow)


@dataclass
class SpendBudgetStatus:
    scope: str
    agent_name: str
    spent_usd: float
    soft_limit_usd: float
    ratio: float
    near_limit: bool
    limit_key: str | None = None


class SpendLedger(Protocol):
    def record(self, record: SpendRecord) -> None:
        ...

    def spent(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
    ) -> float:
        ...

    def list_records(self, limit: int = 100) -> list[SpendRecord]:
        ...

    def budget_status(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
        workspace_soft_limit_usd: float | None = None,
        session_soft_limit_usd: float | None = None,
        threshold_ratio: float = 0.8,
    ) -> dict[str, SpendBudgetStatus]:
        ...


class InMemorySpendLedger:
    def __init__(self) -> None:
        self._records: list[SpendRecord] = []

    def record(self, record: SpendRecord) -> None:
        self._records.append(
            SpendRecord(
                agent_name=record.agent_name,
                model=record.model,
                amount_usd=max(0.0, float(record.amount_usd)),
                workspace_id=record.workspace_id,
                session_id=record.session_id,
                trace_id=record.trace_id,
                created_at=record.created_at,
            )
        )

    def spent(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
    ) -> float:
        total = 0.0
        for item in self._records:
            if item.agent_name != agent_name:
                continue
            if workspace_id is not None and item.workspace_id != workspace_id:
                continue
            if session_id is not None and item.session_id != session_id:
                continue
            total += item.amount_usd
        return round(total, 8)

    def list_records(self, limit: int = 100) -> list[SpendRecord]:
        return list(reversed(self._records[-limit:]))

    def budget_status(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
        workspace_soft_limit_usd: float | None = None,
        session_soft_limit_usd: float | None = None,
        threshold_ratio: float = 0.8,
    ) -> dict[str, SpendBudgetStatus]:
        return _build_budget_status(
            workspace_spent=self.spent(agent_name=agent_name, workspace_id=workspace_id),
            session_spent=self.spent(
                agent_name=agent_name,
                workspace_id=workspace_id,
                session_id=session_id,
            ) if session_id is not None else 0.0,
            agent_name=agent_name,
            workspace_id=workspace_id,
            session_id=session_id,
            workspace_soft_limit_usd=workspace_soft_limit_usd,
            session_soft_limit_usd=session_soft_limit_usd,
            threshold_ratio=threshold_ratio,
        )


class SQLiteSpendLedger:
    def __init__(self, path: str = ".logos_adk/spend_ledger.db") -> None:
        self.path = path

    def _connect(self) -> sqlite3.Connection:
        db_path = Path(self.path)
        if db_path.parent and not db_path.parent.exists():
            db_path.parent.mkdir(parents=True, exist_ok=True)
        connection = sqlite3.connect(self.path)
        connection.row_factory = sqlite3.Row
        self._ensure_schema(connection)
        return connection

    def _ensure_schema(self, connection: sqlite3.Connection) -> None:
        current = get_sqlite_schema_version(connection, SPEND_LEDGER_SCHEMA_KIND) or 0
        connection.execute(
            """
            CREATE TABLE IF NOT EXISTS spend_ledger_records (
                id TEXT PRIMARY KEY,
                agent_name TEXT NOT NULL,
                workspace_id TEXT NOT NULL,
                session_id TEXT NOT NULL,
                model TEXT NOT NULL,
                amount_usd REAL NOT NULL,
                trace_id TEXT,
                created_at TEXT NOT NULL
            )
            """
        )
        connection.execute(
            """
            CREATE INDEX IF NOT EXISTS idx_spend_ledger_records_scope
            ON spend_ledger_records (agent_name, workspace_id, session_id, created_at DESC)
            """
        )
        if current < SPEND_LEDGER_SCHEMA_VERSION:
            set_sqlite_schema_version(
                connection,
                SPEND_LEDGER_SCHEMA_KIND,
                SPEND_LEDGER_SCHEMA_VERSION,
            )
        connection.commit()

    def record(self, record: SpendRecord) -> None:
        with self._connect() as connection:
            connection.execute(
                """
                INSERT INTO spend_ledger_records
                (id, agent_name, workspace_id, session_id, model, amount_usd, trace_id, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    uuid4().hex,
                    record.agent_name,
                    _segment(record.workspace_id),
                    _segment(record.session_id),
                    record.model,
                    max(0.0, float(record.amount_usd)),
                    record.trace_id,
                    record.created_at,
                ),
            )

    def spent(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
    ) -> float:
        clauses = ["agent_name = ?"]
        params: list[object] = [agent_name]
        if workspace_id is not None:
            clauses.append("workspace_id = ?")
            params.append(_segment(workspace_id))
        if session_id is not None:
            clauses.append("session_id = ?")
            params.append(_segment(session_id))
        query = f"""
            SELECT COALESCE(SUM(amount_usd), 0) AS total
            FROM spend_ledger_records
            WHERE {' AND '.join(clauses)}
        """
        with self._connect() as connection:
            row = connection.execute(query, tuple(params)).fetchone()
        return round(float(row["total"] if row is not None else 0.0), 8)

    def list_records(self, limit: int = 100) -> list[SpendRecord]:
        query = f"SELECT * FROM spend_ledger_records ORDER BY created_at DESC LIMIT {limit}"
        with self._connect() as connection:
            rows = connection.execute(query).fetchall()
        return [
            SpendRecord(
                agent_name=row["agent_name"],
                model=row["model"],
                amount_usd=row["amount_usd"],
                workspace_id=row["workspace_id"],
                session_id=row["session_id"],
                trace_id=row["trace_id"],
                created_at=row["created_at"],
            )
            for row in rows
        ]

    def budget_status(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
        workspace_soft_limit_usd: float | None = None,
        session_soft_limit_usd: float | None = None,
        threshold_ratio: float = 0.8,
    ) -> dict[str, SpendBudgetStatus]:
        return _build_budget_status(
            workspace_spent=self.spent(agent_name=agent_name, workspace_id=workspace_id),
            session_spent=self.spent(
                agent_name=agent_name,
                workspace_id=workspace_id,
                session_id=session_id,
            ) if session_id is not None else 0.0,
            agent_name=agent_name,
            workspace_id=workspace_id,
            session_id=session_id,
            workspace_soft_limit_usd=workspace_soft_limit_usd,
            session_soft_limit_usd=session_soft_limit_usd,
            threshold_ratio=threshold_ratio,
        )


class PostgreSQLSpendLedger:
    def __init__(
        self,
        dsn: str | None = None,
        table_name: str = "spend_ledger_records",
    ) -> None:
        self.dsn = resolve_database_url(dsn)
        if not self.dsn:
            raise ValueError(
                "PostgreSQL spend ledger requires a DSN or DATABASE_URL/LOGOS_ADK_DATABASE_URL"
            )
        self.table_name = table_name

    def _connect(self) -> Any:
        return import_psycopg().connect(self.dsn)

    def _ensure_schema(self, connection: Any) -> None:
        current = get_postgresql_schema_version(connection, SPEND_LEDGER_SCHEMA_KIND) or 0
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id TEXT PRIMARY KEY,
                    agent_name TEXT NOT NULL,
                    workspace_id TEXT NOT NULL,
                    session_id TEXT NOT NULL,
                    model TEXT NOT NULL,
                    amount_usd DOUBLE PRECISION NOT NULL,
                    trace_id TEXT,
                    created_at TEXT NOT NULL
                )
                """
            )
            cursor.execute(
                f"""
                CREATE INDEX IF NOT EXISTS idx_{self.table_name}_scope
                ON {self.table_name} (agent_name, workspace_id, session_id, created_at DESC)
                """
            )
        if current < SPEND_LEDGER_SCHEMA_VERSION:
            set_postgresql_schema_version(
                connection,
                SPEND_LEDGER_SCHEMA_KIND,
                SPEND_LEDGER_SCHEMA_VERSION,
            )
        connection.commit()

    def record(self, record: SpendRecord) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    INSERT INTO {self.table_name}
                    (id, agent_name, workspace_id, session_id, model, amount_usd, trace_id, created_at)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
                    """,
                    (
                        uuid4().hex,
                        record.agent_name,
                        _segment(record.workspace_id),
                        _segment(record.session_id),
                        record.model,
                        max(0.0, float(record.amount_usd)),
                        record.trace_id,
                        record.created_at,
                    ),
                )
            connection.commit()

    def spent(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
    ) -> float:
        clauses = ["agent_name = %s"]
        params: list[object] = [agent_name]
        if workspace_id is not None:
            clauses.append("workspace_id = %s")
            params.append(_segment(workspace_id))
        if session_id is not None:
            clauses.append("session_id = %s")
            params.append(_segment(session_id))
        query = f"""
            SELECT COALESCE(SUM(amount_usd), 0) AS total
            FROM {self.table_name}
            WHERE {' AND '.join(clauses)}
        """
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(query, tuple(params))
                row = cursor.fetchone()
        return round(float(_row_value(row, "total", 0) if row is not None else 0.0), 8)

    def list_records(self, limit: int = 100) -> list[SpendRecord]:
        query = f"SELECT * FROM {self.table_name} ORDER BY created_at DESC LIMIT %s"
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(query, (limit,))
                rows = cursor.fetchall()
        return [
            SpendRecord(
                agent_name=_row_value(row, "agent_name", 1),
                model=_row_value(row, "model", 4),
                amount_usd=float(_row_value(row, "amount_usd", 5)),
                workspace_id=_row_value(row, "workspace_id", 2),
                session_id=_row_value(row, "session_id", 3),
                trace_id=_row_value(row, "trace_id", 6),
                created_at=_row_value(row, "created_at", 7),
            )
            for row in rows
        ]

    def budget_status(
        self,
        *,
        agent_name: str,
        workspace_id: str | None = None,
        session_id: str | None = None,
        workspace_soft_limit_usd: float | None = None,
        session_soft_limit_usd: float | None = None,
        threshold_ratio: float = 0.8,
    ) -> dict[str, SpendBudgetStatus]:
        return _build_budget_status(
            workspace_spent=self.spent(agent_name=agent_name, workspace_id=workspace_id),
            session_spent=self.spent(
                agent_name=agent_name,
                workspace_id=workspace_id,
                session_id=session_id,
            ) if session_id is not None else 0.0,
            agent_name=agent_name,
            workspace_id=workspace_id,
            session_id=session_id,
            workspace_soft_limit_usd=workspace_soft_limit_usd,
            session_soft_limit_usd=session_soft_limit_usd,
            threshold_ratio=threshold_ratio,
        )


def _row_value(row: Any, name: str, index: int) -> Any:
    if row is None:
        return None
    try:
        return row[name]
    except (KeyError, TypeError, IndexError):
        return row[index]


def _build_budget_status(
    *,
    workspace_spent: float,
    session_spent: float,
    agent_name: str,
    workspace_id: str | None,
    session_id: str | None,
    workspace_soft_limit_usd: float | None,
    session_soft_limit_usd: float | None,
    threshold_ratio: float,
) -> dict[str, SpendBudgetStatus]:
    payload: dict[str, SpendBudgetStatus] = {}
    if workspace_soft_limit_usd is not None and workspace_soft_limit_usd > 0:
        ratio = workspace_spent / float(workspace_soft_limit_usd)
        payload["workspace"] = SpendBudgetStatus(
            scope="workspace",
            agent_name=agent_name,
            spent_usd=round(workspace_spent, 8),
            soft_limit_usd=float(workspace_soft_limit_usd),
            ratio=round(ratio, 6),
            near_limit=ratio >= threshold_ratio,
            limit_key=workspace_id,
        )
    if session_id is not None and session_soft_limit_usd is not None and session_soft_limit_usd > 0:
        ratio = session_spent / float(session_soft_limit_usd)
        payload["session"] = SpendBudgetStatus(
            scope="session",
            agent_name=agent_name,
            spent_usd=round(session_spent, 8),
            soft_limit_usd=float(session_soft_limit_usd),
            ratio=round(ratio, 6),
            near_limit=ratio >= threshold_ratio,
            limit_key=session_id,
        )
    return payload


__all__ = [
    "SPEND_LEDGER_SCHEMA_KIND",
    "SPEND_LEDGER_SCHEMA_VERSION",
    "SpendBudgetStatus",
    "SpendLedger",
    "SpendRecord",
    "InMemorySpendLedger",
    "PostgreSQLSpendLedger",
    "SQLiteSpendLedger",
]
