"""PostgreSQL storage backend — traces in database tables.

Tables are auto-created on init if they don't exist.
Uses psycopg (v3, sync mode) for connections.
"""

from __future__ import annotations

import json
from datetime import datetime
from typing import Any, Iterator

from ..models import Trace
from .base import StoreBase, sanitize_project_name
from .normalize import normalize_trace
from .summary import summarize_trace

_CREATE_TRACES = """\
CREATE TABLE IF NOT EXISTS bandito_traces (
    trace_id TEXT PRIMARY KEY,
    project TEXT NOT NULL,
    data JSONB NOT NULL,
    started_at TIMESTAMPTZ,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_bandito_traces_project ON bandito_traces(project);
CREATE INDEX IF NOT EXISTS idx_bandito_traces_started_at ON bandito_traces(started_at);
CREATE INDEX IF NOT EXISTS idx_bandito_traces_project_started ON bandito_traces(project, started_at DESC);
"""



class PostgresStore(StoreBase):
    """PostgreSQL storage backend."""

    def __init__(self, conninfo: str, auto_create: bool = True) -> None:
        import psycopg

        self._conninfo = conninfo
        self._conn = psycopg.connect(conninfo, autocommit=True)

        if auto_create:
            self._create_tables()

    def validate(self) -> str | None:
        try:
            with self._conn.cursor() as cur:
                cur.execute("SELECT 1")
            return None
        except Exception as e:
            return f"Cannot connect to PostgreSQL: {e}"

    def _create_tables(self) -> None:
        with self._conn.cursor() as cur:
            cur.execute(_CREATE_TRACES)

    def close(self) -> None:
        self._conn.close()

    def __enter__(self) -> "PostgresStore":
        return self

    def __exit__(self, *exc: object) -> None:
        self.close()

    def __del__(self) -> None:
        if self._conn and not self._conn.closed:
            self._conn.close()

    # --- Trace operations ---

    def write_trace(self, trace: Trace) -> None:
        data = trace.to_dict()
        started_at = data.get("started_at")
        safe_project = sanitize_project_name(trace.project)
        with self._conn.cursor() as cur:
            cur.execute(
                """
                INSERT INTO bandito_traces (trace_id, project, data, started_at)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (trace_id) DO UPDATE SET
                    data = EXCLUDED.data,
                    started_at = EXCLUDED.started_at
                """,
                (trace.trace_id, safe_project, json.dumps(data), started_at),
            )

    def _read_trace(self, trace_id: str) -> dict[str, Any] | None:
        with self._conn.cursor() as cur:
            cur.execute(
                "SELECT data FROM bandito_traces WHERE trace_id = %s",
                (trace_id,),
            )
            row = cur.fetchone()
            if row is None:
                return None
            return _ensure_dict(row[0])

    def iter_traces(self, project: str | None = None) -> Iterator[dict[str, Any]]:
        safe_project = sanitize_project_name(project) if project else None
        with self._conn.cursor() as cur:
            if safe_project is not None:
                cur.execute(
                    "SELECT data FROM bandito_traces WHERE project = %s ORDER BY started_at DESC",
                    (safe_project,),
                )
            else:
                cur.execute("SELECT data FROM bandito_traces ORDER BY started_at DESC")
            for row in cur:
                d = _ensure_dict(row[0])
                if d:
                    yield d

    def list_projects(self) -> list[str]:
        with self._conn.cursor() as cur:
            cur.execute("SELECT DISTINCT project FROM bandito_traces ORDER BY project")
            return [row[0] for row in cur]

    # --- SQL pushdown for pagination ---

    def _query_traces(
        self,
        project: str | None = None,
        since: datetime | None = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        """Execute a paginated trace query with SQL-level filtering."""
        query = "SELECT data FROM bandito_traces"
        conditions: list[str] = []
        params: list[Any] = []

        if project is not None:
            safe_project = sanitize_project_name(project)
            conditions.append("project = %s")
            params.append(safe_project)
        if since is not None:
            conditions.append("started_at >= %s")
            params.append(since)

        if conditions:
            query += " WHERE " + " AND ".join(conditions)
        query += " ORDER BY started_at DESC"

        if offset:
            query += " OFFSET %s"
            params.append(offset)
        if limit is not None and limit > 0:
            query += " LIMIT %s"
            params.append(limit)

        with self._conn.cursor() as cur:
            cur.execute(query, params)
            rows = []
            for row in cur:
                d = _ensure_dict(row[0])
                if d:
                    rows.append(d)
            return rows

    def list_traces_filtered(
        self,
        project: str | None = None,
        tag: str | None = None,
        since: datetime | None = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        """Override with SQL pushdown for since/limit/offset."""
        if tag is not None:
            # Tags live in JSONB — over-fetch then filter in-memory
            sql_limit = (limit * 3) if limit else None
            rows = self._query_traces(project, since, sql_limit, offset=0)
            traces = []
            for raw in rows:
                trace = normalize_trace(raw)
                if tag not in trace.get("tags", []):
                    continue
                traces.append(trace)
            if offset:
                traces = traces[offset:]
            if limit is not None and limit > 0:
                traces = traces[:limit]
            if len(traces) < (limit or 0) and sql_limit and len(rows) >= sql_limit:
                # Over-fetch wasn't enough — fall back to base class
                return super().list_traces_filtered(project, tag, since, limit, offset)
            return traces

        rows = self._query_traces(project, since, limit, offset)
        return [normalize_trace(raw) for raw in rows]

    def list_trace_summaries(
        self,
        project: str | None = None,
        tag: str | None = None,
        since: datetime | None = None,
        limit: int | None = None,
        offset: int = 0,
    ) -> list[dict[str, Any]]:
        """Override with SQL pushdown for since/limit/offset."""
        if tag is not None:
            sql_limit = (limit * 3) if limit else None
            rows = self._query_traces(project, since, sql_limit, offset=0)
            summaries = []
            for raw in rows:
                trace = normalize_trace(raw)
                if tag not in trace.get("tags", []):
                    continue
                summaries.append(summarize_trace(trace))
            if offset:
                summaries = summaries[offset:]
            if limit is not None and limit > 0:
                summaries = summaries[:limit]
            if len(summaries) < (limit or 0) and sql_limit and len(rows) >= sql_limit:
                return super().list_trace_summaries(project, tag, since, limit, offset)
            return summaries

        rows = self._query_traces(project, since, limit, offset)
        return [summarize_trace(normalize_trace(raw)) for raw in rows]


def _ensure_dict(value: Any) -> dict[str, Any] | None:
    """Ensure value is a dict — psycopg may return JSONB as dict or str."""
    if isinstance(value, dict):
        return value
    if isinstance(value, str):
        try:
            return json.loads(value)
        except (json.JSONDecodeError, TypeError):
            return None
    return None
