from __future__ import annotations

import json
import uuid
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Any


@dataclass(frozen=True)
class RetrievalTrace:
    trace_id: str
    query: str
    subject_scope: str
    intent: str
    decision: str
    reason: str
    resolved_drug_id: str | None
    citations: tuple[dict[str, Any], ...]
    created_at: datetime | None = None


class PostgresTraceRepository:
    """Opens a new connection per call — no pooling (F-09: a real pool, with
    startup-time lifecycle, is a further improvement not made here).

    `connect_timeout` matters more than it looks: found live 2026-08-06 that
    an unreachable Postgres (packets dropped, not actively refused) makes a
    bare `psycopg.connect()` hang on the OS-level TCP timeout — tens of
    seconds, not immediate — which defeats a caller's try/except fail-open
    around `save()` just as effectively as no try/except at all, since the
    exception it's waiting for never arrives in time. `routers/rag.py`
    wraps `save()` to keep a trace outage from failing an already-computed
    answer; this bounds how long that protection can take to kick in.
    """

    def __init__(self, dsn: str) -> None:
        self._dsn = dsn

    def migrate(self, migration_path: Path) -> None:
        import psycopg

        statement = migration_path.read_text(encoding="utf-8")
        with psycopg.connect(self._dsn, connect_timeout=5) as connection:
            connection.execute(statement)

    def save(
        self,
        *,
        query: str,
        subject_scope: str,
        intent: str,
        decision: str,
        reason: str,
        resolved_drug_id: str | None,
        citations: tuple[dict[str, Any], ...],
    ) -> str:
        import psycopg

        trace_id = str(uuid.uuid4())
        with psycopg.connect(self._dsn, connect_timeout=5) as connection:
            connection.execute(
                """
                INSERT INTO rag_retrieval_trace (
                    trace_id, query_text, subject_scope, query_intent,
                    decision, reason, resolved_drug_id, citations
                ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s::jsonb)
                """,
                (
                    trace_id, query, subject_scope, intent, decision, reason,
                    resolved_drug_id, json.dumps(citations, ensure_ascii=False),
                ),
            )
        return trace_id

    def get(self, trace_id: str) -> RetrievalTrace | None:
        import psycopg

        with psycopg.connect(self._dsn, connect_timeout=5) as connection:
            row = connection.execute(
                """
                SELECT trace_id::text, query_text, subject_scope, query_intent,
                       decision, reason, resolved_drug_id, citations, created_at
                FROM rag_retrieval_trace WHERE trace_id = %s
                """,
                (trace_id,),
            ).fetchone()
        if row is None:
            return None
        return RetrievalTrace(
            trace_id=row[0], query=row[1], subject_scope=row[2], intent=row[3],
            decision=row[4], reason=row[5], resolved_drug_id=row[6],
            citations=tuple(row[7]), created_at=row[8],
        )