"""Asyncpg replay patches for async PostgreSQL.

Provides deterministic replay of asyncpg database operations
using captured transcript records.
"""

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 asyncpg functions for uninstall
_original_connect: Optional[Callable[..., Any]] = None
_original_create_pool: Optional[Callable[..., Any]] = None


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


def _current_task_id() -> Optional[str]:
    """Get current async task ID for ordering validation."""
    try:
        import asyncio
        task = asyncio.current_task()
        if task is None:
            return None
        return str(id(task))
    except Exception:
        return None


def _match_query(expected_sql: str, actual_sql: str, fuzzy: bool = True) -> bool:
    """Check if queries match, optionally with fuzzy matching.

    Args:
        expected_sql: SQL from transcript
        actual_sql: SQL being executed
        fuzzy: If True, ignore whitespace differences

    Returns:
        True if queries match
    """
    if fuzzy:
        return _normalize_sql(expected_sql) == _normalize_sql(actual_sql)
    return expected_sql == actual_sql


class AsyncPGReplayConnection:
    """Replay connection that returns recorded results for asyncpg.

    This replaces actual asyncpg connections during replay, returning
    pre-recorded query results from the transcript.
    """

    def __init__(self, index: TranscriptIndex, conn_id: str):
        """Initialize replay connection.

        Args:
            index: Transcript index for looking up records
            conn_id: Connection identifier
        """
        self._index = index
        self._conn_id = conn_id
        self._tx_counter = 0
        self._closed = False

    @property
    def _tx_id(self) -> str:
        return f"{self._conn_id}-{self._tx_counter}"

    async def fetch(self, query: str, *args, timeout: float = None) -> List[Dict[str, Any]]:
        """Replay fetch operation.

        Args:
            query: SQL query
            *args: Query arguments
            timeout: Query timeout (ignored in replay)

        Returns:
            List of record dicts from transcript

        Raises:
            ReplayDivergence: If query doesn't match transcript
        """
        record = self._expect_sql(query, list(args))

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

        # Return rows as list of dicts (asyncpg Record-like)
        rows = record.get("rows") or []
        description = record.get("description") or []

        # Convert list rows to dicts using column names
        if description and rows:
            col_names = [col.get("name", f"col{i}") for i, col in enumerate(description)]
            return [dict(zip(col_names, row)) for row in rows]

        return rows

    async def fetchrow(self, query: str, *args, timeout: float = None) -> Optional[Dict[str, Any]]:
        """Replay fetchrow operation.

        Args:
            query: SQL query
            *args: Query arguments
            timeout: Query timeout (ignored in replay)

        Returns:
            Single record dict or None
        """
        rows = await self.fetch(query, *args, timeout=timeout)
        return rows[0] if rows else None

    async def fetchval(self, query: str, *args, column: int = 0, timeout: float = None) -> Any:
        """Replay fetchval operation.

        Args:
            query: SQL query
            *args: Query arguments
            column: Column index to return
            timeout: Query timeout (ignored in replay)

        Returns:
            Single value from the specified column
        """
        row = await self.fetchrow(query, *args, timeout=timeout)
        if row is None:
            return None

        # Get value by column index
        values = list(row.values())
        if column < len(values):
            return values[column]
        return None

    async def execute(self, query: str, *args, timeout: float = None) -> str:
        """Replay execute operation.

        Args:
            query: SQL query
            *args: Query arguments
            timeout: Query timeout (ignored in replay)

        Returns:
            Status string like "INSERT 0 1"
        """
        record = self._expect_sql(query, list(args))

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

        rowcount = record.get("rowcount", 0)
        return f"UPDATE {rowcount}"

    async def executemany(self, query: str, args: List, *, timeout: float = None) -> None:
        """Replay executemany operation.

        Args:
            query: SQL query
            args: List of argument tuples
            timeout: Query timeout (ignored in replay)
        """
        # Execute each set of args
        for arg_set in args:
            await self.execute(query, *arg_set, timeout=timeout)

    def _expect_sql(self, query: str, params: List[Any]) -> Dict[str, Any]:
        """Get expected SQL record from transcript.

        Args:
            query: SQL query
            params: Query parameters

        Returns:
            Transcript record

        Raises:
            ReplayDivergence: If query doesn't match or no more records
        """
        record = self._index.expect("SQL")

        expected_sql = record.get("query") or record.get("sql_canonical", "")
        if not _match_query(expected_sql, query):
            raise ReplayDivergence(
                "SQL_MISMATCH",
                {
                    "expected": expected_sql,
                    "actual": query,
                }
            )

        # Validate async task ordering if available
        expected_task = record.get("task_id")
        current_task = _current_task_id()
        if expected_task and current_task and expected_task != current_task:
            raise ReplayDivergence(
                "ASYNC_ORDER",
                {"expected_task": expected_task, "actual_task": current_task}
            )

        return record

    def transaction(self, *, isolation: str = None, readonly: bool = False, deferrable: bool = False):
        """Create a replay transaction context manager."""
        return AsyncPGReplayTransaction(self)

    async def close(self) -> None:
        """Close the connection."""
        await asyncio.sleep(0)
        self._closed = True

    def is_closed(self) -> bool:
        """Check if connection is closed."""
        return self._closed


class AsyncPGReplayTransaction:
    """Replay transaction context manager for asyncpg."""

    def __init__(self, conn: AsyncPGReplayConnection):
        self._conn = conn

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if exc_type is None:
            # Commit
            self._conn._tx_counter += 1
        else:
            # Rollback
            self._conn._tx_counter += 1
        return False


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

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

    def _next_conn_id(self) -> str:
        self._conn_counter += 1
        return f"asyncpg-{self._conn_counter}"

    async def acquire(self, *, timeout: float = None) -> AsyncPGReplayConnection:
        """Acquire a replay connection from the pool."""
        return await asyncio.sleep(0, result=AsyncPGReplayConnection(self._index, self._next_conn_id()))

    async def release(self, conn: AsyncPGReplayConnection) -> None:
        """Release a connection back to the pool."""
        await asyncio.sleep(0)

    def acquire_context(self, *, timeout: float = None):
        """Return context manager for acquiring connection."""
        return _PoolAcquireContext(self, timeout)

    async def close(self) -> None:
        """Close the pool."""
        await asyncio.sleep(0)


class _PoolAcquireContext:
    """Context manager for pool.acquire()."""

    def __init__(self, pool: AsyncPGReplayPool, timeout: float):
        self._pool = pool
        self._timeout = timeout
        self._conn = None

    async def __aenter__(self) -> AsyncPGReplayConnection:
        self._conn = await self._pool.acquire(timeout=self._timeout)
        return self._conn

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self._conn:
            await self._pool.release(self._conn)
        return False


def install_asyncpg(index: TranscriptIndex) -> None:
    """Install asyncpg replay patches.

    Args:
        index: Transcript index containing SQL records
    """
    global _original_connect, _original_create_pool

    try:
        import asyncpg
    except ImportError:
        return

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

    _conn_counter = 0

    def _next_conn_id() -> str:
        nonlocal _conn_counter
        _conn_counter += 1
        return f"asyncpg-{_conn_counter}"

    async def replay_connect(*args, **kwargs) -> AsyncPGReplayConnection:
        """Replacement for asyncpg.connect()."""
        return AsyncPGReplayConnection(index, _next_conn_id())

    async def replay_create_pool(*args, **kwargs) -> AsyncPGReplayPool:
        """Replacement for asyncpg.create_pool()."""
        return AsyncPGReplayPool(index)

    # Patch asyncpg module
    _original_connect = asyncpg.connect
    _original_create_pool = asyncpg.create_pool
    asyncpg.connect = replay_connect  # type: ignore[assignment]
    asyncpg.create_pool = replay_create_pool  # type: ignore[assignment]


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

    try:
        import asyncpg
    except ImportError:
        return

    if _original_connect is not None:
        asyncpg.connect = _original_connect  # type: ignore[assignment]
        _original_connect = None

    if _original_create_pool is not None:
        asyncpg.create_pool = _original_create_pool  # type: ignore[assignment]
        _original_create_pool = None
