"""SQL replay patches for psycopg2 and psycopg."""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from autodebug_replay.canonicalize import sql_fingerprint
from autodebug_replay.transcripts import ReplayDivergence, TranscriptIndex


def _normalize_sql(sql: str) -> str:
    return sql_fingerprint(sql)


def _expect_sql(index: TranscriptIndex, sql: str, params: Any, tx_id: str, seq_in_tx: int) -> Dict[str, Any]:
    record = index.expect("SQL")
    expected_sql = str(record.get("sql_canonical") or record.get("query") or "")
    actual_sql = str(sql)
    expected_fp = _normalize_sql(expected_sql)
    actual_fp = _normalize_sql(actual_sql)
    if expected_fp != actual_fp:
        raise ReplayDivergence(
            "MISMATCH",
            {
                "expected_sql": expected_sql,
                "actual_sql": actual_sql,
                "expected_sql_fp": expected_fp,
                "actual_sql_fp": actual_fp,
            },
        )
    if record.get("tx_id") != tx_id or record.get("seq_in_tx") != seq_in_tx:
        raise ReplayDivergence(
            "SQL_DIVERGENCE",
            {"expected_stmt_index": record.get("seq_in_tx"), "observed_stmt_index": seq_in_tx},
        )
    return record


class _TxState:
    def __init__(self, conn_id: str) -> None:
        self._conn_id = conn_id
        self._tx_counter = 0
        self._seq_in_tx = 0

    @property
    def tx_id(self) -> str:
        return f"{self._conn_id}-{self._tx_counter}"

    def next_seq(self) -> int:
        self._seq_in_tx += 1
        return self._seq_in_tx

    def rotate(self) -> None:
        self._tx_counter += 1
        self._seq_in_tx = 0


class _ReplayCursor:
    def __init__(self, index: TranscriptIndex, tx_state: _TxState) -> None:
        self._index = index
        self._tx_state = tx_state
        self._rows: Optional[List[List[Any]]] = None
        self.rowcount = None
        self.description = None

    def execute(self, sql, params=None):
        seq_in_tx = self._tx_state.next_seq()
        record = _expect_sql(self._index, sql, params, self._tx_state.tx_id, seq_in_tx)
        if record.get("exception"):
            exc = record["exception"]
            raise ReplayDivergence("SQL_DIVERGENCE", {"exception": exc})
        self.rowcount = record.get("rowcount")
        self.description = record.get("description")
        self._rows = record.get("rows") or []
        return None

    def executemany(self, sql, params_list):
        return self.execute(sql, params_list)

    def fetchone(self):
        if self._rows is None:
            raise ReplayDivergence("SQL_DIVERGENCE", {"detail": "fetchone without rows"})
        if not self._rows:
            return None
        return self._rows.pop(0)

    def fetchmany(self, size=None):
        if self._rows is None:
            raise ReplayDivergence("SQL_DIVERGENCE", {"detail": "fetchmany without rows"})
        if size is None:
            size = 1
        rows = self._rows[:size]
        self._rows = self._rows[size:]
        return rows

    def fetchall(self):
        if self._rows is None:
            raise ReplayDivergence("SQL_DIVERGENCE", {"detail": "fetchall without rows"})
        rows = list(self._rows)
        self._rows = []
        return rows

    def __iter__(self):
        return self

    def __next__(self):
        row = self.fetchone()
        if row is None:
            raise StopIteration
        return row


class _ReplayConnection:
    def __init__(self, index: TranscriptIndex, conn_id: str) -> None:
        self._index = index
        self._tx_state = _TxState(conn_id)

    def cursor(self, *args, **kwargs):
        return _ReplayCursor(self._index, self._tx_state)

    def commit(self):
        self._tx_state.rotate()

    def rollback(self):
        self._tx_state.rotate()


_proxy_counter = 0

# Store originals for uninstall
_orig_psycopg2_connect = None
_orig_psycopg_connect = None


def _next_conn_id() -> str:
    global _proxy_counter
    _proxy_counter += 1
    return f"conn-{_proxy_counter}"


def install_psycopg2(index: TranscriptIndex) -> None:
    global _orig_psycopg2_connect
    try:
        import psycopg2
    except ImportError:
        return

    _orig_psycopg2_connect = psycopg2.connect

    def wrapped_connect(*args, **kwargs):
        return _ReplayConnection(index, _next_conn_id())

    psycopg2.connect = wrapped_connect  # type: ignore[assignment]


def install_psycopg(index: TranscriptIndex) -> None:
    global _orig_psycopg_connect
    try:
        import psycopg
    except ImportError:
        return

    _orig_psycopg_connect = psycopg.connect

    def wrapped_connect(*args, **kwargs):
        return _ReplayConnection(index, _next_conn_id())

    psycopg.connect = wrapped_connect  # type: ignore[assignment]


def uninstall() -> None:
    """Restore original SQL functions."""
    global _orig_psycopg2_connect, _orig_psycopg_connect, _proxy_counter

    if _orig_psycopg2_connect is not None:
        try:
            import psycopg2
            psycopg2.connect = _orig_psycopg2_connect
        except ImportError:
            pass
        _orig_psycopg2_connect = None

    if _orig_psycopg_connect is not None:
        try:
            import psycopg
            psycopg.connect = _orig_psycopg_connect
        except ImportError:
            pass
        _orig_psycopg_connect = None

    _proxy_counter = 0
