"""SQL patching for psycopg2 and psycopg (v3)."""

from __future__ import annotations

import base64
import threading
from typing import Any, Iterable, List, Optional

from autodebug_probe.exception_patch import format_stack
from autodebug_probe.recorder import get_recorder


def _normalize_sql(sql: str) -> str:
    return " ".join(sql.split())


def _serialize_value(value: Any) -> Any:
    if isinstance(value, bytes):
        return {"__type__": "bytes", "base64": base64.b64encode(value).decode("ascii")}
    if isinstance(value, memoryview):
        raw = value.tobytes()
        return {"__type__": "bytes", "base64": base64.b64encode(raw).decode("ascii")}
    if isinstance(value, (list, tuple)):
        return [_serialize_value(item) for item in value]
    if isinstance(value, dict):
        return {str(k): _serialize_value(v) for k, v in value.items()}
    try:
        import json  # noqa: F401
        _ = value  # keep value
        return value
    except TypeError:
        return {"__type__": "repr", "value": repr(value)}


def _serialize_rows(rows: Iterable[Iterable[Any]]) -> List[List[Any]]:
    return [[_serialize_value(item) for item in row] for row in rows]


def _serialize_description(description: Optional[Iterable]) -> Optional[List[dict]]:
    if description is None:
        return None
    serialized = []
    for col in description:
        serialized.append(
            {
                "name": col[0],
                "type_code": col[1],
                "display_size": col[2],
                "internal_size": col[3],
                "precision": col[4],
                "scale": col[5],
                "null_ok": col[6],
            }
        )
    return serialized


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 _CursorProxy:
    def __init__(self, inner, tx_state: _TxState) -> None:
        self._inner = inner
        self._tx_state = tx_state

    def _record_error(self, sql: str, params: Any, exc: Exception) -> None:
        """Record a SQL error payload and an EXCEPTION event."""
        payload = {
            "sql_canonical": _normalize_sql(str(sql)),
            "params_canonical": _serialize_value(params),
            "tx_id": self._tx_state.tx_id,
            "seq_in_tx": self._tx_state.next_seq(),
            "rowcount": getattr(self._inner, "rowcount", None),
            "description": _serialize_description(
                getattr(self._inner, "description", None)
            ),
            "rows": None,
            "exception": {
                "type": exc.__class__.__name__,
                "message": str(exc),
                "sqlstate": (
                    getattr(exc, "pgcode", None)
                    or getattr(exc, "sqlstate", None)
                ),
            },
        }
        get_recorder().record_sql(payload)
        get_recorder().record_event(
            {
                "type": "EXCEPTION",
                "exc_type": exc.__class__.__name__,
                "message": str(exc),
                "stack": format_stack(exc.__traceback__),
                "boundary": "SQL",
            }
        )

    def execute(self, sql, params=None):
        try:
            result = self._inner.execute(sql, params)
        except Exception as exc:  # pylint: disable=broad-except
            self._record_error(sql, params, exc)
            raise

        # Record the query immediately without eagerly fetching rows.
        # This avoids buffering large result sets in memory and preserves
        # streaming / server-side cursor semantics.  Rows can still be
        # obtained later via fetchone/fetchmany/fetchall which delegate
        # directly to the inner cursor.
        description = getattr(self._inner, "description", None)
        payload = {
            "sql_canonical": _normalize_sql(str(sql)),
            "params_canonical": _serialize_value(params),
            "tx_id": self._tx_state.tx_id,
            "seq_in_tx": self._tx_state.next_seq(),
            "rowcount": getattr(self._inner, "rowcount", None),
            "description": _serialize_description(description),
            "rows": None,
            "exception": None,
        }
        get_recorder().record_sql(payload)
        return result

    def executemany(self, sql, params_list):
        """Execute a SQL statement against all parameter sequences.

        Delegates to the inner cursor's executemany() rather than calling
        execute() in a loop, preserving the native batching behaviour.
        """
        try:
            result = self._inner.executemany(sql, params_list)
        except Exception as exc:  # pylint: disable=broad-except
            self._record_error(sql, params_list, exc)
            raise

        payload = {
            "sql_canonical": _normalize_sql(str(sql)),
            "params_canonical": _serialize_value(params_list),
            "tx_id": self._tx_state.tx_id,
            "seq_in_tx": self._tx_state.next_seq(),
            "rowcount": getattr(self._inner, "rowcount", None),
            "description": _serialize_description(
                getattr(self._inner, "description", None)
            ),
            "rows": None,
            "exception": None,
        }
        get_recorder().record_sql(payload)
        return result

    def fetchone(self):
        return self._inner.fetchone()

    def fetchmany(self, size=None):
        if size is None:
            return self._inner.fetchmany()
        return self._inner.fetchmany(size)

    def fetchall(self):
        return self._inner.fetchall()

    def __iter__(self):
        return self

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

    def __getattr__(self, name: str):
        return getattr(self._inner, name)


class _ConnectionProxy:
    def __init__(self, inner, conn_id: str) -> None:
        self._inner = inner
        self._tx_state = _TxState(conn_id)

    def cursor(self, *args, **kwargs):
        inner_cursor = self._inner.cursor(*args, **kwargs)
        return _CursorProxy(inner_cursor, self._tx_state)

    def commit(self):
        result = self._inner.commit()
        self._tx_state.rotate()
        return result

    def rollback(self):
        result = self._inner.rollback()
        self._tx_state.rotate()
        return result

    def __getattr__(self, name: str):
        return getattr(self._inner, name)


_proxy_counter = 0
_proxy_counter_lock = threading.Lock()


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


def install_psycopg2() -> None:
    try:
        import psycopg2
    except ImportError:
        return

    original_connect = psycopg2.connect

    def wrapped_connect(*args, **kwargs):
        conn = original_connect(*args, **kwargs)
        return _ConnectionProxy(conn, _next_conn_id())

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


def install_psycopg() -> None:
    try:
        import psycopg
    except ImportError:
        return

    original_connect = psycopg.connect

    def wrapped_connect(*args, **kwargs):
        conn = original_connect(*args, **kwargs)
        return _ConnectionProxy(conn, _next_conn_id())

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