"""Postgres-backed persistence for cumulative AI model usage counters.

Keeps ModelUsageTracker (an in-memory, process-lifetime-only registry defined
in app/shared/model_usage_tracker.py) from losing its counters on every
backend restart. This store only ever persists the *current cumulative
snapshot* per provider+model — not a historical time-series log — matching
the "durable counters" scope decided in
docs/superpowers/specs/2026-07-23-status-model-usage-hardening-design.md.
"""

from __future__ import annotations

from contextlib import contextmanager

import psycopg2
import psycopg2.extras
from psycopg2.pool import ThreadedConnectionPool

from app.config.settings import settings
from app.shared.model_usage_tracker import ModelUsageEntry

# Table creation follows the same CREATE TABLE IF NOT EXISTS idiom used by
# every other Postgres store in this codebase — no migration framework.
_CREATE_TABLE = """
CREATE TABLE IF NOT EXISTS model_usage_stats (
    provider          VARCHAR(64)  NOT NULL,
    model             VARCHAR(128) NOT NULL,
    total_tokens      BIGINT       NOT NULL DEFAULT 0,
    prompt_tokens     BIGINT       NOT NULL DEFAULT 0,
    completion_tokens BIGINT       NOT NULL DEFAULT 0,
    call_count_ok     BIGINT       NOT NULL DEFAULT 0,
    call_count_error  BIGINT       NOT NULL DEFAULT 0,
    last_called_at    TIMESTAMPTZ,
    last_latency_ms   INTEGER,
    last_error        TEXT,
    updated_at        TIMESTAMPTZ  NOT NULL DEFAULT NOW(),
    PRIMARY KEY (provider, model)
);
"""

_UPSERT = """
INSERT INTO model_usage_stats
    (provider, model, total_tokens, prompt_tokens, completion_tokens,
     call_count_ok, call_count_error, last_called_at, last_latency_ms, last_error, updated_at)
VALUES
    (%(provider)s, %(model)s, %(total_tokens)s, %(prompt_tokens)s, %(completion_tokens)s,
     %(call_count_ok)s, %(call_count_error)s, %(last_called_at)s, %(last_latency_ms)s, %(last_error)s, NOW())
ON CONFLICT (provider, model) DO UPDATE SET
    total_tokens = EXCLUDED.total_tokens,
    prompt_tokens = EXCLUDED.prompt_tokens,
    completion_tokens = EXCLUDED.completion_tokens,
    call_count_ok = EXCLUDED.call_count_ok,
    call_count_error = EXCLUDED.call_count_error,
    last_called_at = EXCLUDED.last_called_at,
    last_latency_ms = EXCLUDED.last_latency_ms,
    last_error = EXCLUDED.last_error,
    updated_at = NOW();
"""


class PostgresModelUsageStore:
    """Load and flush ModelUsageTracker snapshots to/from a Postgres table."""

    def __init__(self) -> None:
        """Open a small connection pool and ensure the table exists."""
        self._pool = ThreadedConnectionPool(
            minconn=1,
            maxconn=3,
            host=settings.postgres_host,
            port=settings.postgres_port,
            user=settings.postgres_user,
            password=settings.postgres_password,
            dbname=settings.postgres_db,
        )
        self._ensure_schema()

    def _ensure_schema(self) -> None:
        """Create the model_usage_stats table if it does not already exist."""
        with self._conn() as conn:
            with conn.cursor() as cur:
                cur.execute(_CREATE_TABLE)
            conn.commit()

    @contextmanager
    def _conn(self):
        """Borrow a pooled connection and always return it, even on error."""
        conn = self._pool.getconn()
        try:
            yield conn
        finally:
            self._pool.putconn(conn)

    def load_all(self) -> dict[str, ModelUsageEntry]:
        """Return every persisted row as {"provider:model": ModelUsageEntry}."""
        with self._conn() as conn:
            with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
                cur.execute("SELECT * FROM model_usage_stats")
                rows = cur.fetchall()
        entries: dict[str, ModelUsageEntry] = {}
        for row in rows:
            entry = ModelUsageEntry(
                provider=row["provider"],
                model=row["model"],
                total_tokens=row["total_tokens"],
                prompt_tokens=row["prompt_tokens"],
                completion_tokens=row["completion_tokens"],
                call_count_ok=row["call_count_ok"],
                call_count_error=row["call_count_error"],
                last_called_at=row["last_called_at"],
                last_latency_ms=row["last_latency_ms"],
                last_error=row["last_error"],
            )
            entries[f"{entry.provider}:{entry.model}"] = entry
        return entries

    def flush(self, entries: dict[str, ModelUsageEntry]) -> None:
        """Upsert the current cumulative snapshot of every tracked entry.

        A no-op for an empty snapshot — avoids opening a connection for nothing
        (e.g. before any LLM/embedding/reranker call has happened yet).
        """
        if not entries:
            return
        with self._conn() as conn:
            with conn.cursor() as cur:
                for entry in entries.values():
                    cur.execute(
                        _UPSERT,
                        {
                            "provider": entry.provider,
                            "model": entry.model,
                            "total_tokens": entry.total_tokens,
                            "prompt_tokens": entry.prompt_tokens,
                            "completion_tokens": entry.completion_tokens,
                            "call_count_ok": entry.call_count_ok,
                            "call_count_error": entry.call_count_error,
                            "last_called_at": entry.last_called_at,
                            "last_latency_ms": entry.last_latency_ms,
                            "last_error": entry.last_error,
                        },
                    )
            conn.commit()
