"""aiosqlite replay patches for async SQLite.

Replays captured aiosqlite database operations from transcripts.

Usage:
    from autodebug_replay.aiosqlite_replay import install_aiosqlite
    install_aiosqlite(transcript_index)

    # Now all aiosqlite calls return recorded data
    async with aiosqlite.connect("db.sqlite") as db:
        async with db.execute("SELECT * FROM users") as cursor:
            rows = await cursor.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 aiosqlite functions for uninstall
_original_connect: Optional[Callable[..., Any]] = None
_active_index: Optional[TranscriptIndex] = 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"aiosqlite-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 _ReplayConnectionContextManager:
    """Context manager for replay connection."""

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

    async def __aenter__(self):
        self._conn = _ReplayConnection(self._index, _next_conn_id())
        return self._conn

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

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

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


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

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

    async def execute(self, sql: str, parameters=None):
        """Execute and return replay cursor."""
        return await self._execute_replay(sql, parameters)

    async def executemany(self, sql: str, parameters):
        """Execute many and return replay cursor."""
        return await self._execute_replay(sql, parameters)

    async def executescript(self, sql_script: str):
        """Execute script and return replay cursor."""
        return await self._execute_replay(sql_script, None)

    async def _execute_replay(self, sql: str, parameters):
        """Create replay cursor with recorded data."""
        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"),
            })

        return _ReplayCursor(record)

    async def commit(self):
        await asyncio.sleep(0)
        self._tx_state.rotate()

    async def rollback(self):
        await asyncio.sleep(0)
        self._tx_state.rotate()

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

    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


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

    def __init__(self, record: Dict[str, Any]):
        self._record = record
        rows = record.get("rows") or []
        self._rows = [tuple(row) for row in rows]
        self._row_index = 0
        self.rowcount = record.get("rowcount", -1)
        self.description = record.get("description")
        self.lastrowid = record.get("lastrowid")

    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 close(self):
        await asyncio.sleep(0)

    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 __aiter__(self):
        return self

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


def install_aiosqlite(index: TranscriptIndex) -> None:
    """Install aiosqlite replay patches."""
    global _original_connect, _active_index

    try:
        import aiosqlite
    except ImportError:
        return

    _active_index = index
    if _original_connect is not None:
        return  # Already patched; switch active transcript index.

    def replay_connect(database, **kwargs):
        if _active_index is None:
            raise RuntimeError("aiosqlite replay installed without active transcript index")
        return _ReplayConnectionContextManager(_active_index)

    _original_connect = aiosqlite.connect
    aiosqlite.connect = replay_connect


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


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

    try:
        import aiosqlite
    except ImportError:
        return

    if _original_connect is not None:
        aiosqlite.connect = _original_connect
        _original_connect = None
    _active_index = None
