"""Async detection and patching for AutoDebug.

This module provides:
1. Detection of async runtimes (asyncio, anyio, trio)
2. Patching of asyncio.sleep for time capture
3. Patching of httpx AsyncClient for HTTP capture
4. AsyncDeterminismGuard for detecting concurrent I/O
5. Task graph capture for asyncio.create_task / TaskGroup
"""

from __future__ import annotations

import asyncio
import sys
import threading
from typing import Any, Callable, Dict, List, Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from autodebug_probe.recorder import Recorder

# Original functions
_original_asyncio_sleep: Optional[Callable] = None
_original_httpx_async_request: Optional[Callable] = None
_original_httpx_async_send: Optional[Callable] = None
_original_create_task: Optional[Callable] = None

# Global recorder reference
_recorder: Optional["Recorder"] = None

# ---------------------------------------------------------------------------
# Task graph capture
# ---------------------------------------------------------------------------

_task_counter: int = 0
_task_counter_lock = threading.Lock()
_task_id_map: Dict[int, int] = {}  # id(asyncio.Task) -> monotonic task_id
_concurrent_tasks: int = 0
_max_concurrent_tasks: int = 0


def _next_task_id() -> int:
    global _task_counter
    with _task_counter_lock:
        _task_counter += 1
        return _task_counter


def _get_task_id(task: asyncio.Task) -> int:
    """Get or assign a monotonic task ID for an asyncio.Task."""
    obj_id = id(task)
    if obj_id not in _task_id_map:
        _task_id_map[obj_id] = _next_task_id()
    return _task_id_map[obj_id]


def _get_current_task_id() -> Optional[int]:
    """Get the task ID of the currently running asyncio task."""
    try:
        task = asyncio.current_task()
        if task is not None:
            return _get_task_id(task)
    except RuntimeError:
        pass
    return None


def _record_task_event(event_type: str, task_id: int, **extra: Any) -> None:
    """Record a task lifecycle event to the task graph stream."""
    recorder = _get_recorder()
    recorder.record_task_graph({
        "type": event_type,
        "task_id": task_id,
        **extra,
    })


def _patched_create_task(coro, *, name=None, context=None) -> asyncio.Task:
    """Patched asyncio.create_task that records task creation."""
    global _concurrent_tasks, _max_concurrent_tasks

    parent_task_id = _get_current_task_id()

    # Call original create_task
    if context is not None:
        task = _original_create_task(coro, name=name, context=context)
    elif name is not None:
        task = _original_create_task(coro, name=name)
    else:
        task = _original_create_task(coro)

    task_id = _get_task_id(task)
    coro_name = getattr(coro, "__qualname__", None) or getattr(coro, "__name__", str(coro))

    _record_task_event(
        "TASK_CREATE",
        task_id=task_id,
        coro_name=coro_name,
        parent_task_id=parent_task_id,
        task_name=name or task.get_name(),
    )

    with _task_counter_lock:
        _concurrent_tasks += 1
        if _concurrent_tasks > _max_concurrent_tasks:
            _max_concurrent_tasks = _concurrent_tasks

    def _on_done(t: asyncio.Task) -> None:
        global _concurrent_tasks
        tid = _get_task_id(t)
        exc = t.exception() if not t.cancelled() else None
        _record_task_event(
            "TASK_COMPLETE",
            task_id=tid,
            cancelled=t.cancelled(),
            has_exception=exc is not None,
        )
        with _task_counter_lock:
            _concurrent_tasks -= 1

    task.add_done_callback(_on_done)
    return task


def install_task_graph_capture() -> None:
    """Patch asyncio.create_task to capture task graph."""
    global _original_create_task, _task_counter, _task_id_map
    global _concurrent_tasks, _max_concurrent_tasks

    if _original_create_task is not None:
        return  # Already patched

    _task_counter = 0
    _task_id_map = {}
    _concurrent_tasks = 0
    _max_concurrent_tasks = 0

    _original_create_task = asyncio.create_task
    asyncio.create_task = _patched_create_task


def uninstall_task_graph_capture() -> None:
    """Restore original asyncio.create_task."""
    global _original_create_task, _task_id_map
    global _concurrent_tasks, _max_concurrent_tasks

    if _original_create_task is not None:
        asyncio.create_task = _original_create_task
        _original_create_task = None
    _task_id_map = {}
    _concurrent_tasks = 0
    _max_concurrent_tasks = 0


def get_task_graph_stats() -> Dict[str, Any]:
    """Get task graph capture statistics for manifest."""
    total_tasks = _task_counter
    return {
        "total_tasks_captured": total_tasks,
        "max_concurrent_tasks": _max_concurrent_tasks,
        "async_risk_level": (
            "safe" if _max_concurrent_tasks <= 1
            else "risky" if _max_concurrent_tasks <= 4
            else "unsupported"
        ),
    }

# Async runtimes and their typical uses
ASYNC_RUNTIMES = {
    "asyncio": "stdlib",
    "anyio": "starlette/fastapi",
    "trio": "alternative",
}


def detect_async_runtime() -> Dict[str, Any]:
    """Detect async runtime and framework.

    Returns dict with:
        - async_detected: bool
        - runtime: str or None (asyncio, anyio, trio)
        - framework: str or None (pytest-asyncio, etc.)
        - event_loop_running: bool
        - async_http: bool (httpx AsyncClient detected)
    """
    result = {
        "async_detected": False,
        "runtime": None,
        "framework": None,
        "event_loop_running": False,
        "async_http": False,
        "anyio_backend": None,
    }

    # Check for running event loop (asyncio)
    try:
        loop = asyncio.get_running_loop()
        result["async_detected"] = True
        result["runtime"] = "asyncio"
        result["event_loop_running"] = True
    except RuntimeError:
        pass

    # Check for pytest-asyncio
    if "pytest_asyncio" in sys.modules:
        result["async_detected"] = True
        result["framework"] = "pytest-asyncio"
        result["runtime"] = result["runtime"] or "asyncio"

    # Check for anyio (FastAPI/Starlette use this)
    if "anyio" in sys.modules:
        result["async_detected"] = True
        result["runtime"] = result["runtime"] or "anyio"
        result["framework"] = result["framework"] or "anyio"

        # Detect anyio backend
        try:
            import anyio
            result["anyio_backend"] = anyio.get_current_task().name if hasattr(anyio, "get_current_task") else None
        except Exception:
            pass

    # Check for trio
    if "trio" in sys.modules:
        result["async_detected"] = True
        result["runtime"] = "trio"
        result["framework"] = result["framework"] or "trio"

    # Check for httpx async usage
    if "httpx" in sys.modules:
        try:
            import httpx
            # Check if AsyncClient is in the module's namespace (imported)
            if hasattr(httpx, "AsyncClient"):
                # Mark as potentially async HTTP if httpx is loaded
                result["async_http"] = result["async_detected"]
        except Exception:
            pass

    return result


def _get_recorder() -> "Recorder":
    """Get the global recorder instance."""
    if _recorder is None:
        from autodebug_probe.recorder import get_recorder
        return get_recorder()
    return _recorder


async def _patched_asyncio_sleep(delay: float, result: Any = None) -> Any:
    """Patched asyncio.sleep that records the sleep request.

    During capture: Records the sleep duration.
    During replay: Uses minimal sleep to yield control.
    """
    recorder = _get_recorder()

    # Record the sleep request
    recorder.record_time_rng({
        "type": "ASYNC_SLEEP",
        "delay": delay,
    })

    # Actually sleep a minimal amount to yield control
    # This prevents blocking while maintaining async semantics
    if delay > 0:
        await _original_asyncio_sleep(min(delay, 0.001))
    else:
        await _original_asyncio_sleep(0)

    return result


def install_asyncio_sleep_patch() -> None:
    """Install asyncio.sleep patch for time capture."""
    global _original_asyncio_sleep

    if _original_asyncio_sleep is not None:
        return  # Already patched

    _original_asyncio_sleep = asyncio.sleep
    asyncio.sleep = _patched_asyncio_sleep


def uninstall_asyncio_sleep_patch() -> None:
    """Restore original asyncio.sleep."""
    global _original_asyncio_sleep

    if _original_asyncio_sleep is not None:
        asyncio.sleep = _original_asyncio_sleep
        _original_asyncio_sleep = None


def install_httpx_async_patch() -> None:
    """Install httpx AsyncClient patch for HTTP capture."""
    global _original_httpx_async_request, _original_httpx_async_send

    try:
        import httpx
    except ImportError:
        return

    if _original_httpx_async_request is not None:
        return  # Already patched

    _original_httpx_async_request = httpx.AsyncClient.request

    async def patched_request(self, method: str, url, **kwargs):
        """Patched AsyncClient.request that captures HTTP calls."""
        from autodebug_probe.http_patch import _record_request_response

        # Extract request body
        req_body = kwargs.get("content") or kwargs.get("data") or kwargs.get("json")
        content_type = None
        if "headers" in kwargs:
            content_type = kwargs["headers"].get("Content-Type") or kwargs["headers"].get("content-type")

        # Make the actual request
        response = await _original_httpx_async_request(self, method, url, **kwargs)

        # Record the request/response
        _record_request_response(
            method=method,
            url=str(url),
            req_body=req_body,
            content_type=content_type,
            response=response,
            is_async=True,
        )

        return response

    httpx.AsyncClient.request = patched_request

    # Also patch _send for lower-level capture
    _original_httpx_async_send = httpx.AsyncClient._send_single_request

    async def patched_send(self, request):
        """Patched AsyncClient._send_single_request for complete capture."""
        response = await _original_httpx_async_send(self, request)
        return response

    httpx.AsyncClient._send_single_request = patched_send


def uninstall_httpx_async_patch() -> None:
    """Restore original httpx AsyncClient methods."""
    global _original_httpx_async_request, _original_httpx_async_send

    try:
        import httpx
    except ImportError:
        return

    if _original_httpx_async_request is not None:
        httpx.AsyncClient.request = _original_httpx_async_request
        _original_httpx_async_request = None

    if _original_httpx_async_send is not None:
        httpx.AsyncClient._send_single_request = _original_httpx_async_send
        _original_httpx_async_send = None


class AsyncDeterminismGuard:
    """Guard to detect concurrent async I/O operations.

    For async code to be deterministic, we cannot have concurrent I/O.
    This guard tracks pending I/O operations and raises ReplayDivergence
    if concurrent I/O is detected.
    """

    def __init__(self, strict: bool = False):
        """Initialize the guard.

        Args:
            strict: If True, raise on concurrent I/O. If False, just warn.
        """
        self.pending_io: Dict[str, int] = {}
        self.total_concurrent_detected = 0
        self.strict = strict
        self._lock: Optional[asyncio.Lock] = None

    @property
    def lock(self) -> asyncio.Lock:
        """Lazy-init the lock (must be created in async context)."""
        if self._lock is None:
            self._lock = asyncio.Lock()
        return self._lock

    async def begin_io(self, io_type: str) -> None:
        """Mark the beginning of an I/O operation.

        Args:
            io_type: Type of I/O (HTTP, SQL, FILE)

        Raises:
            ReplayDivergence: If concurrent I/O is detected in strict mode.
        """
        async with self.lock:
            total_pending = sum(self.pending_io.values())
            if total_pending > 0:
                self.total_concurrent_detected += 1
                if self.strict:
                    from autodebug_replay.transcripts import ReplayDivergence
                    raise ReplayDivergence(
                        "CONCURRENT_ASYNC_IO",
                        {
                            "new_io_type": io_type,
                            "pending_io": dict(self.pending_io),
                            "total_pending": total_pending,
                        }
                    )
                else:
                    import warnings
                    warnings.warn(
                        f"Concurrent async I/O detected: starting {io_type} while {self.pending_io} pending",
                        RuntimeWarning,
                        stacklevel=3,
                    )

            self.pending_io[io_type] = self.pending_io.get(io_type, 0) + 1

    async def end_io(self, io_type: str) -> None:
        """Mark the end of an I/O operation."""
        async with self.lock:
            if io_type in self.pending_io:
                self.pending_io[io_type] -= 1
                if self.pending_io[io_type] <= 0:
                    del self.pending_io[io_type]

    def get_stats(self) -> Dict[str, Any]:
        """Get concurrent I/O detection statistics."""
        return {
            "concurrent_io_detected": self.total_concurrent_detected,
            "pending_io": dict(self.pending_io),
        }


# Global guard instance
_async_guard: Optional[AsyncDeterminismGuard] = None


def get_async_guard(strict: bool = False) -> AsyncDeterminismGuard:
    """Get or create the global async determinism guard."""
    global _async_guard
    if _async_guard is None:
        _async_guard = AsyncDeterminismGuard(strict=strict)
    return _async_guard


def install_all_async_patches() -> None:
    """Install all async-related patches."""
    install_asyncio_sleep_patch()
    install_httpx_async_patch()
    install_task_graph_capture()


def uninstall_all_async_patches() -> None:
    """Uninstall all async-related patches."""
    uninstall_asyncio_sleep_patch()
    uninstall_httpx_async_patch()
    uninstall_task_graph_capture()
