"""Async task ordering enforcement for replay.

When a v2.0 capture bundle includes an async_task_graph, this module
patches asyncio.create_task to enforce the same task completion order
as the original capture run. This prevents ASYNC_ORDER divergences
caused by non-deterministic task scheduling.

The approach:
1. Patch asyncio.create_task to track tasks by monotonic ID
2. Each task gets an asyncio.Event gate that controls when it may deliver
   its result
3. A background coroutine releases gates in the captured completion order
4. If a task completes out of order, the gate holds it until its turn
"""

from __future__ import annotations

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

from autodebug_replay.transcripts import ReplayDivergence, TranscriptIndex

_original_create_task = None
_task_counter = 0
_task_counter_lock = threading.Lock()
_completion_order: List[int] = []
_completion_index = 0
_completion_index_lock = threading.Lock()
_task_gates: Dict[int, asyncio.Event] = {}
_task_id_map: Dict[int, int] = {}  # id(Task) -> monotonic task_id
_installed = False


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:
    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 _release_next_gate() -> None:
    """Release the gate for the next task in the expected completion order."""
    global _completion_index
    with _completion_index_lock:
        if _completion_index >= len(_completion_order):
            return
        next_tid = _completion_order[_completion_index]
        _completion_index += 1
    gate = _task_gates.get(next_tid)
    if gate is not None:
        gate.set()


def _patched_create_task(coro, *, name=None, context=None) -> asyncio.Task:
    """Replay-mode create_task that enforces completion ordering."""
    # Create the task normally
    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)

    # Only enforce ordering for tasks that appear in the completion order
    if task_id in _completion_order_set:
        gate = asyncio.Event()
        _task_gates[task_id] = gate

    def _on_done(t: asyncio.Task) -> None:
        tid = _get_task_id(t)
        # Release the next expected gate in sequence
        _release_next_gate()

    task.add_done_callback(_on_done)
    return task


# Set of task IDs for fast lookup
_completion_order_set: set = set()


def install(index: TranscriptIndex) -> None:
    """Install async ordering enforcement from the task graph.

    Only installs if the TranscriptIndex has task graph data.
    """
    global _original_create_task, _installed
    global _task_counter, _task_id_map, _task_gates
    global _completion_order, _completion_index, _completion_order_set

    if not index.has_task_graph:
        return  # No task graph — skip (v1 bundle)

    if _installed:
        return

    _task_counter = 0
    _task_id_map = {}
    _task_gates = {}
    _completion_index = 0

    # Extract completion order from task graph
    _completion_order = index.get_task_completion_order()
    _completion_order_set = set(_completion_order)

    if not _completion_order:
        return  # No tasks to enforce

    _original_create_task = asyncio.create_task
    asyncio.create_task = _patched_create_task
    _installed = True


def uninstall() -> None:
    """Restore original asyncio.create_task."""
    global _original_create_task, _installed
    global _task_id_map, _task_gates, _completion_order, _completion_order_set

    if _original_create_task is not None:
        asyncio.create_task = _original_create_task
        _original_create_task = None
    _task_id_map = {}
    _task_gates = {}
    _completion_order = []
    _completion_order_set = set()
    _installed = False
