"""AsyncPG patching for async PostgreSQL capture.

Captures all asyncpg database operations for replay.

Supported operations:
- fetch, fetchrow, fetchval
- execute, executemany
- prepare (prepared statements)
- copy_from_query, copy_to_table

Usage:
    from autodebug_probe.asyncpg_patch import install_asyncpg
    install_asyncpg()

    # Now all asyncpg calls are captured
    rows = await conn.fetch("SELECT * FROM users")
"""

from __future__ import annotations

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

from autodebug_probe.recorder import get_recorder


def _serialize_row(row) -> Dict[str, Any]:
    """Serialize an asyncpg Record to a dict."""
    if row is None:
        return None
    if hasattr(row, "items"):
        return dict(row.items())
    if hasattr(row, "_asdict"):
        return row._asdict()
    return {"_value": str(row)}


def _serialize_result(result) -> Any:
    """Serialize query result for recording."""
    if result is None:
        return None
    if isinstance(result, list):
        return [_serialize_row(r) for r in result]
    if isinstance(result, str):
        return result  # Status string like "INSERT 1"
    return _serialize_row(result)


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

    # Patch Connection.fetch
    original_fetch = asyncpg.Connection.fetch

    async def patched_fetch(self, query: str, *args, timeout: float = None):
        result = await original_fetch(self, query, *args, timeout=timeout)
        get_recorder().record_sql({
            "driver": "asyncpg",
            "method": "fetch",
            "query": query,
            "params": list(args),
            "result": _serialize_result(result),
            "rowcount": len(result) if result else 0,
        })
        return result

    asyncpg.Connection.fetch = patched_fetch

    # Patch Connection.fetchrow
    original_fetchrow = asyncpg.Connection.fetchrow

    async def patched_fetchrow(self, query: str, *args, timeout: float = None):
        result = await original_fetchrow(self, query, *args, timeout=timeout)
        get_recorder().record_sql({
            "driver": "asyncpg",
            "method": "fetchrow",
            "query": query,
            "params": list(args),
            "result": _serialize_row(result),
            "rowcount": 1 if result else 0,
        })
        return result

    asyncpg.Connection.fetchrow = patched_fetchrow

    # Patch Connection.fetchval
    original_fetchval = asyncpg.Connection.fetchval

    async def patched_fetchval(self, query: str, *args, column: int = 0, timeout: float = None):
        result = await original_fetchval(self, query, *args, column=column, timeout=timeout)
        get_recorder().record_sql({
            "driver": "asyncpg",
            "method": "fetchval",
            "query": query,
            "params": list(args),
            "result": str(result) if result is not None else None,
            "rowcount": 1 if result is not None else 0,
        })
        return result

    asyncpg.Connection.fetchval = patched_fetchval

    # Patch Connection.execute
    original_execute = asyncpg.Connection.execute

    async def patched_execute(self, query: str, *args, timeout: float = None):
        result = await original_execute(self, query, *args, timeout=timeout)
        get_recorder().record_sql({
            "driver": "asyncpg",
            "method": "execute",
            "query": query,
            "params": list(args),
            "result": result,  # Status string like "INSERT 1"
            "rowcount": _parse_rowcount(result),
        })
        return result

    asyncpg.Connection.execute = patched_execute

    # Patch Connection.executemany
    original_executemany = asyncpg.Connection.executemany

    async def patched_executemany(self, query: str, args, *, timeout: float = None):
        result = await original_executemany(self, query, args, timeout=timeout)
        args_list = list(args) if args else []
        get_recorder().record_sql({
            "driver": "asyncpg",
            "method": "executemany",
            "query": query,
            "params": [list(a) for a in args_list[:100]],  # Limit for large batches
            "batch_size": len(args_list),
            "result": result,
        })
        return result

    asyncpg.Connection.executemany = patched_executemany


def _parse_rowcount(status: str) -> int:
    """Parse rowcount from status string like 'INSERT 1' or 'UPDATE 5'."""
    if not status:
        return 0
    parts = status.split()
    if len(parts) >= 2 and parts[-1].isdigit():
        return int(parts[-1])
    return 0
