"""aiomysql replay patches for async MySQL.

Replays captured aiomysql database operations from transcripts.

Usage:
    from autodebug_replay.aiomysql_replay import install_aiomysql
    install_aiomysql(transcript_index)

    # Now all aiomysql calls return recorded data
    async with conn.cursor() as cur:
        await cur.execute("SELECT * FROM users")
        rows = await cur.fetchall()  # Returns recorded rows
"""

from __future__ import annotations

import asyncio
from typing import Any, Callable, Dict, List, Optional

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

# Original aiomysql functions for uninstall
_original_create_pool: Optional[Callable[..., Any]] = None
_original_connect: Optional[Callable[..., Any]] = None


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


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

    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


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


def _expect_sql(index: TranscriptIndex, sql: str, tx_id: str, seq_in_tx: int) -> Dict[str, Any]:
    """Get expected SQL record from transcript."""
    record = index.expect("SQL")

    # Verify SQL matches
    expected_sql = str(record.get("sql_canonical") or record.get("query") or "")
    actual_sql = _normalize_sql(sql)
    expected_fp = _normalize_sql(expected_sql)
    if expected_fp != actual_sql:
        raise ReplayDivergence("SQL_MISMATCH", {
            "expected_sql": expected_sql,
            "actual_sql": actual_sql,
            "expected_sql_fp": expected_fp,
            "actual_sql_fp": actual_sql,
        })

    # Verify transaction state
    if record.get("tx_id") != tx_id or record.get("seq_in_tx") != seq_in_tx:
        raise ReplayDivergence("SQL_ORDER", {
            "expected_tx": record.get("tx_id"),
            "expected_seq": record.get("seq_in_tx"),
            "actual_tx": tx_id,
            "actual_seq": seq_in_tx,
        })

    return record


class _ReplayPool:
    """Replay pool that returns replay connections."""

    def __init__(self, index: TranscriptIndex):
        self._index = index

    async def acquire(self):
        return await asyncio.sleep(0, result=_ReplayConnection(self._index, _next_conn_id()))

    def release(self, conn):
        pass

    async def __aenter__(self):
        await asyncio.sleep(0)
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await asyncio.sleep(0)
        return False

    def close(self):
        pass

    async def wait_closed(self):
        await asyncio.sleep(0)


class _ReplayConnection:
    """Replay connection for aiomysql."""

    def __init__(self, index: TranscriptIndex, conn_id: str):
        self._index = index
        self._tx_state = _TxState(conn_id)

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

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

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

    async def begin(self):
        pass

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass


class _ReplayCursorContextManager:
    """Context manager for replay cursor."""

    def __init__(self, index: TranscriptIndex, tx_state: _TxState):
        self._index = index
        self._tx_state = tx_state

    async def __aenter__(self):
        return _ReplayCursor(self._index, self._tx_state)

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        pass


class _ReplayCursor:
    """Replay cursor for aiomysql."""

    def __init__(self, index: TranscriptIndex, tx_state: _TxState):
        self._index = index
        self._tx_state = tx_state
        self._rows: List[tuple] = []
        self._row_index = 0
        self.rowcount = -1
        self.description = None
        self.lastrowid = None

    async def execute(self, sql: str, args=None):
        """Execute query and load recorded results."""
        seq_in_tx = self._tx_state.next_seq()
        record = _expect_sql(self._index, sql, self._tx_state.tx_id, seq_in_tx)

        # Check for recorded exception
        if record.get("exception"):
            exc_info = record["exception"]
            raise ReplayDivergence("SQL_EXCEPTION", {
                "type": exc_info.get("type"),
                "message": exc_info.get("message"),
            })

        # Load recorded results
        self.rowcount = record.get("rowcount", -1)
        self.description = record.get("description")
        self.lastrowid = record.get("lastrowid")

        # Convert rows to tuples
        rows = record.get("rows") or []
        self._rows = [tuple(row) for row in rows]
        self._row_index = 0

        return self.rowcount

    async def executemany(self, sql: str, args_list):
        """Execute many - replay as single execute."""
        return await self.execute(sql, args_list)

    async def fetchone(self):
        """Fetch one row from recorded results."""
        if self._row_index >= len(self._rows):
            return None
        row = self._rows[self._row_index]
        self._row_index += 1
        return row

    async def fetchmany(self, size=None):
        """Fetch many rows from recorded results."""
        if size is None:
            size = 1
        rows = self._rows[self._row_index:self._row_index + size]
        self._row_index += len(rows)
        return rows

    async def fetchall(self):
        """Fetch all remaining rows from recorded results."""
        rows = self._rows[self._row_index:]
        self._row_index = len(self._rows)
        return rows

    async def callproc(self, procname: str, args=()):
        """Call stored procedure - replay from transcript."""
        seq_in_tx = self._tx_state.next_seq()
        record = _expect_sql(self._index, f"CALL {procname}", self._tx_state.tx_id, seq_in_tx)

        if record.get("exception"):
            exc_info = record["exception"]
            raise ReplayDivergence("SQL_EXCEPTION", {
                "type": exc_info.get("type"),
                "message": exc_info.get("message"),
            })

        return args

    async def close(self):
        pass


def install_aiomysql(index: TranscriptIndex) -> None:
    """Install aiomysql replay patches."""
    global _original_create_pool, _original_connect

    try:
        import aiomysql
    except ImportError:
        return

    if _original_create_pool is not None or _original_connect is not None:
        return  # Already patched

    async def replay_create_pool(*args, **kwargs):
        return _ReplayPool(index)

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

    _original_create_pool = aiomysql.create_pool
    _original_connect = aiomysql.connect
    aiomysql.create_pool = replay_create_pool
    aiomysql.connect = replay_connect


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


def uninstall() -> None:
    """Restore original aiomysql functions."""
    global _original_create_pool, _original_connect

    try:
        import aiomysql
    except ImportError:
        return

    if _original_create_pool is not None:
        aiomysql.create_pool = _original_create_pool
        _original_create_pool = None

    if _original_connect is not None:
        aiomysql.connect = _original_connect
        _original_connect = None
