"""aiosqlite patching for async SQLite capture.

Captures all aiosqlite database operations for replay.

Supported operations:
- execute, executemany, executescript
- fetchone, fetchmany, fetchall
- commit, rollback

Usage:
    from autodebug_probe.aiosqlite_patch import install_aiosqlite
    install_aiosqlite()

    # Now all aiosqlite calls are captured
    async with aiosqlite.connect("db.sqlite") as db:
        async with db.execute("SELECT * FROM users") as cursor:
            rows = await cursor.fetchall()
"""

from __future__ import annotations

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

from autodebug_probe.recorder import get_recorder
from autodebug_probe.correlation import get_async_task_id


def _serialize_value(value: Any) -> Any:
    """Serialize a value for JSON recording."""
    if value is None:
        return None
    if isinstance(value, (str, int, float, bool)):
        return value
    if isinstance(value, bytes):
        return {"__type__": "bytes", "base64": base64.b64encode(value).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()}
    # Handle datetime
    if hasattr(value, "isoformat"):
        return {"__type__": "datetime", "value": value.isoformat()}
    return str(value)


def _serialize_rows(rows: List[tuple]) -> List[List[Any]]:
    """Serialize query result rows."""
    if not rows:
        return []
    return [[_serialize_value(item) for item in row] for row in rows]


def _normalize_sql(sql: str) -> str:
    """Normalize SQL for comparison."""
    return " ".join(sql.split())


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


class _TxState:
    """Transaction state tracker."""

    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


_proxy_counter = 0
_proxy_counter_lock = threading.Lock()


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


def install_aiosqlite() -> None:
    """Install aiosqlite patches for capture."""
    try:
        import aiosqlite
    except ImportError:
        return

    original_connect = aiosqlite.connect

    def patched_connect(database, **kwargs):
        """Wrap aiosqlite.connect to return capturing connection."""
        return _ConnectionContextManager(original_connect, database, kwargs)

    aiosqlite.connect = patched_connect


class _ConnectionContextManager:
    """Context manager that wraps aiosqlite connection."""

    def __init__(self, original_connect, database, kwargs):
        self._original_connect = original_connect
        self._database = database
        self._kwargs = kwargs
        self._conn_cm = None  # The original context manager
        self._conn = None  # The actual connection
        self._wrapper = None

    async def __aenter__(self):
        # Get the context manager from original connect
        self._conn_cm = self._original_connect(self._database, **self._kwargs)
        # Enter the context manager to get the actual connection
        self._conn = await self._conn_cm.__aenter__()
        self._wrapper = _ConnectionWrapper(self._conn, _next_conn_id())
        return self._wrapper

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self._conn_cm is not None:
            return await self._conn_cm.__aexit__(exc_type, exc_val, exc_tb)
        return None

    def __await__(self):
        return self._await_impl().__await__()

    async def _await_impl(self):
        # For direct await (not context manager), get the connection directly
        self._conn_cm = self._original_connect(self._database, **self._kwargs)
        self._conn = await self._conn_cm
        self._wrapper = _ConnectionWrapper(self._conn, _next_conn_id())
        return self._wrapper


class _ConnectionWrapper:
    """Wrapper around aiosqlite Connection that captures queries."""

    def __init__(self, conn, conn_id: str):
        self._conn = conn
        self._tx_state = _TxState(conn_id)

    async def execute(self, sql: str, parameters=None):
        """Execute a query and capture it."""
        return await self._execute_and_capture(sql, parameters, "execute")

    async def executemany(self, sql: str, parameters):
        """Execute many and capture."""
        return await self._execute_and_capture(sql, parameters, "executemany")

    async def executescript(self, sql_script: str):
        """Execute script and capture."""
        return await self._execute_and_capture(sql_script, None, "executescript")

    async def _execute_and_capture(self, sql: str, parameters, method: str):
        """Execute and capture the query."""
        try:
            if method == "executemany":
                cursor = await self._conn.executemany(sql, parameters)
            elif method == "executescript":
                cursor = await self._conn.executescript(sql)
            elif parameters is not None:
                cursor = await self._conn.execute(sql, parameters)
            else:
                cursor = await self._conn.execute(sql)
        except Exception as exc:
            # Record failed query
            get_recorder().record_sql({
                "driver": "aiosqlite",
                "sql_canonical": _normalize_sql(sql),
                "params_canonical": _serialize_value(parameters),
                "tx_id": self._tx_state.tx_id,
                "seq_in_tx": self._tx_state.next_seq(),
                "task_id": get_async_task_id(),
                "method": method,
                "rowcount": -1,
                "exception": {
                    "type": exc.__class__.__name__,
                    "message": str(exc),
                },
            })
            raise

        # Wrap cursor for capturing fetches
        wrapped = _CursorWrapper(cursor, self._tx_state, sql, parameters, method)
        return wrapped

    async def commit(self):
        await self._conn.commit()
        self._tx_state.rotate()

    async def rollback(self):
        await self._conn.rollback()
        self._tx_state.rotate()

    async def close(self):
        await self._conn.close()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    def __getattr__(self, name):
        return getattr(self._conn, name)


class _CursorWrapper:
    """Wrapper around aiosqlite Cursor that captures results."""

    def __init__(self, cursor, tx_state: _TxState, sql: str, parameters, method: str):
        self._cursor = cursor
        self._tx_state = tx_state
        self._sql = sql
        self._parameters = parameters
        self._method = method
        self._recorded = False
        self._fetched_rows: List[tuple] = []

    async def fetchone(self):
        """Fetch one row and accumulate for recording."""
        row = await self._cursor.fetchone()
        if row is not None:
            self._fetched_rows.append(row)
        return row

    async def fetchmany(self, size=None):
        """Fetch many rows and accumulate for recording."""
        rows = await self._cursor.fetchmany(size)
        self._fetched_rows.extend(rows)
        return rows

    async def fetchall(self):
        """Fetch all rows and record."""
        rows = await self._cursor.fetchall()
        self._fetched_rows.extend(rows)
        self._record_query()
        return rows

    def _record_query(self):
        """Record the query with all fetched results."""
        if self._recorded:
            return
        self._recorded = True

        get_recorder().record_sql({
            "driver": "aiosqlite",
            "sql_canonical": _normalize_sql(self._sql),
            "params_canonical": _serialize_value(self._parameters),
            "tx_id": self._tx_state.tx_id,
            "seq_in_tx": self._tx_state.next_seq(),
            "task_id": get_async_task_id(),
            "method": self._method,
            "rowcount": self._cursor.rowcount,
            "description": _serialize_description(self._cursor.description),
            "rows": _serialize_rows(self._fetched_rows),
            "exception": None,
        })

    @property
    def rowcount(self):
        return self._cursor.rowcount

    @property
    def description(self):
        return self._cursor.description

    @property
    def lastrowid(self):
        return self._cursor.lastrowid

    async def close(self):
        # Record if we haven't yet (for non-SELECT queries)
        if not self._recorded:
            self._record_query()
        await self._cursor.close()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    def __aiter__(self):
        return self

    async def __anext__(self):
        row = await self.fetchone()
        if row is None:
            self._record_query()
            raise StopAsyncIteration
        return row

    def __getattr__(self, name):
        return getattr(self._cursor, name)


def install() -> None:
    """Alias for install_aiosqlite."""
    install_aiosqlite()
