"""aiomysql patching for async MySQL capture.

Captures all aiomysql database operations for replay.

Supported operations:
- execute, executemany
- fetchone, fetchmany, fetchall
- callproc (stored procedures)

Usage:
    from autodebug_probe.aiomysql_patch import install_aiomysql
    install_aiomysql()

    # Now all aiomysql calls are captured
    async with conn.cursor() as cur:
        await cur.execute("SELECT * FROM users")
        rows = await cur.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()}
    # Handle Decimal
    if hasattr(value, "__float__"):
        return float(value)
    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())


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"aiomysql-conn-{conn_id}"


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

    # Store original pool create function
    original_create_pool = aiomysql.create_pool

    async def patched_create_pool(*args, **kwargs):
        """Wrap pool to return wrapped connections."""
        pool = await original_create_pool(*args, **kwargs)
        return _PoolWrapper(pool)

    aiomysql.create_pool = patched_create_pool

    # Also patch direct connection creation
    original_connect = aiomysql.connect

    async def patched_connect(*args, **kwargs):
        """Wrap direct connection."""
        conn = await original_connect(*args, **kwargs)
        return _ConnectionWrapper(conn, _next_conn_id())

    aiomysql.connect = patched_connect


class _PoolWrapper:
    """Wrapper around aiomysql Pool that wraps acquired connections."""

    def __init__(self, pool):
        self._pool = pool

    async def acquire(self):
        conn = await self._pool.acquire()
        return _ConnectionWrapper(conn, _next_conn_id())

    def release(self, conn):
        if isinstance(conn, _ConnectionWrapper):
            self._pool.release(conn._conn)
        else:
            self._pool.release(conn)

    async def __aenter__(self):
        return self

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

    def close(self):
        self._pool.close()

    async def wait_closed(self):
        await self._pool.wait_closed()

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


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

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

    def cursor(self, *args, **kwargs):
        """Return a wrapped cursor."""
        return _CursorContextManager(self._conn, self._tx_state, args, kwargs)

    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 begin(self):
        await self._conn.begin()

    async def __aenter__(self):
        return self

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

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


class _CursorContextManager:
    """Context manager for cursor that wraps the actual cursor."""

    def __init__(self, conn, tx_state: _TxState, args, kwargs):
        self._conn = conn
        self._tx_state = tx_state
        self._args = args
        self._kwargs = kwargs
        self._cursor = None
        self._wrapper = None

    async def __aenter__(self):
        self._cursor = await self._conn.cursor(*self._args, **self._kwargs)
        self._wrapper = _CursorWrapper(self._cursor, self._tx_state)
        return self._wrapper

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


class _CursorWrapper:
    """Wrapper around aiomysql Cursor that captures queries."""

    def __init__(self, cursor, tx_state: _TxState):
        self._cursor = cursor
        self._tx_state = tx_state
        self._last_sql: Optional[str] = None
        self._last_params: Optional[Any] = None

    async def execute(self, sql: str, args=None):
        """Execute a query and capture it."""
        self._last_sql = sql
        self._last_params = args

        try:
            result = await self._cursor.execute(sql, args)
        except Exception as exc:
            # Record failed query
            get_recorder().record_sql({
                "driver": "aiomysql",
                "sql_canonical": _normalize_sql(sql),
                "params_canonical": _serialize_value(args),
                "tx_id": self._tx_state.tx_id,
                "seq_in_tx": self._tx_state.next_seq(),
                "task_id": get_async_task_id(),
                "rowcount": -1,
                "exception": {
                    "type": exc.__class__.__name__,
                    "message": str(exc),
                },
            })
            raise

        # Record successful query (rows captured on fetch)
        get_recorder().record_sql({
            "driver": "aiomysql",
            "sql_canonical": _normalize_sql(sql),
            "params_canonical": _serialize_value(args),
            "tx_id": self._tx_state.tx_id,
            "seq_in_tx": self._tx_state.next_seq(),
            "task_id": get_async_task_id(),
            "rowcount": self._cursor.rowcount,
            "description": self._serialize_description(),
            "exception": None,
        })

        return result

    async def executemany(self, sql: str, args_list):
        """Execute many queries and capture them."""
        self._last_sql = sql

        try:
            result = await self._cursor.executemany(sql, args_list)
        except Exception as exc:
            args_sample = list(args_list)[:10] if args_list else []
            get_recorder().record_sql({
                "driver": "aiomysql",
                "sql_canonical": _normalize_sql(sql),
                "params_canonical": _serialize_value(args_sample),
                "tx_id": self._tx_state.tx_id,
                "seq_in_tx": self._tx_state.next_seq(),
                "task_id": get_async_task_id(),
                "rowcount": -1,
                "batch_size": len(args_list) if args_list else 0,
                "exception": {
                    "type": exc.__class__.__name__,
                    "message": str(exc),
                },
            })
            raise

        args_sample = list(args_list)[:10] if args_list else []
        get_recorder().record_sql({
            "driver": "aiomysql",
            "sql_canonical": _normalize_sql(sql),
            "params_canonical": _serialize_value(args_sample),
            "tx_id": self._tx_state.tx_id,
            "seq_in_tx": self._tx_state.next_seq(),
            "task_id": get_async_task_id(),
            "rowcount": self._cursor.rowcount,
            "batch_size": len(args_list) if args_list else 0,
            "exception": None,
        })

        return result

    async def fetchone(self):
        """Fetch one row."""
        return await self._cursor.fetchone()

    async def fetchmany(self, size=None):
        """Fetch many rows."""
        if size is None:
            size = self._cursor.arraysize
        return await self._cursor.fetchmany(size)

    async def fetchall(self):
        """Fetch all rows."""
        return await self._cursor.fetchall()

    async def callproc(self, procname: str, args=()):
        """Call a stored procedure."""
        try:
            result = await self._cursor.callproc(procname, args)
        except Exception as exc:
            get_recorder().record_sql({
                "driver": "aiomysql",
                "sql_canonical": f"CALL {procname}",
                "params_canonical": _serialize_value(args),
                "tx_id": self._tx_state.tx_id,
                "seq_in_tx": self._tx_state.next_seq(),
                "task_id": get_async_task_id(),
                "stored_procedure": procname,
                "exception": {
                    "type": exc.__class__.__name__,
                    "message": str(exc),
                },
            })
            raise

        get_recorder().record_sql({
            "driver": "aiomysql",
            "sql_canonical": f"CALL {procname}",
            "params_canonical": _serialize_value(args),
            "tx_id": self._tx_state.tx_id,
            "seq_in_tx": self._tx_state.next_seq(),
            "task_id": get_async_task_id(),
            "stored_procedure": procname,
            "exception": None,
        })

        return result

    def _serialize_description(self) -> Optional[List[dict]]:
        """Serialize cursor description."""
        desc = self._cursor.description
        if desc is None:
            return None
        return [
            {
                "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],
            }
            for col in desc
        ]

    @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):
        await self._cursor.close()

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


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