"""
lambdagent.agentruntime.cancel — per-run cancellation registry.

Lets the host (e.g. agentpaas) signal cancellation to a running react loop +
provider subprocess from a different thread, without having to plumb a
`cancel_check_fn` through every level of the call stack.

Usage:
    # Host (HTTP handler thread): before starting agent execution
    cancel.begin_run(run_id)
    try:
        result = compiled_term.apply(input_text, ctx)
    except cancel.CancelledRun:
        # The user clicked stop; run was killed mid-step.
        ...
    finally:
        cancel.end_run(run_id)

    # Host (separate thread, e.g. cancel endpoint): request stop
    cancel.request_cancel(run_id)

    # Provider (inside the run thread): check + register subprocess
    proc = subprocess.Popen(...)
    cancel.bind_proc(proc)
    try:
        while proc.poll() is None:
            if cancel.is_cancelled():
                proc.kill()
                raise cancel.CancelledRun()
            ...
    finally:
        cancel.unbind_proc()

The registry is per-run (keyed by run_id); the current run is tracked via a
thread-local set by `begin_run()`. Providers don't need to know the run_id —
they just call `is_cancelled()` / `bind_proc()` which look it up.

Sub-agent calls inherit the parent thread's run_id (single-threaded ReAct
loop), so the same registry covers nested compiled-term invocations.
"""
from __future__ import annotations

import subprocess
import threading
from typing import Callable, Dict, Optional


_lock = threading.Lock()
_events: Dict[str, threading.Event] = {}
_procs: Dict[str, subprocess.Popen] = {}
# Per-run hook fired whenever the bound subprocess pid changes (bind/unbind).
# Lets the host persist the pid somewhere durable (DB) so a restarted backend
# can still find and kill the orphan subprocess.
_on_proc_change: Dict[str, Callable[[Optional[int]], None]] = {}
_thread_local = threading.local()


class CancelledRun(Exception):
    """Raised when the host has requested cancellation of the current run."""


def begin_run(
    run_id: str,
    on_proc_change: Optional[Callable[[Optional[int]], None]] = None,
) -> None:
    """Mark `run_id` as cancellable and bind it to the calling thread.

    Idempotent: re-binding the same run_id replaces any prior event (e.g. if
    a previous attempt for this id crashed without end_run).

    `on_proc_change`, if provided, is invoked synchronously with the bound
    Popen.pid (int) each time bind_proc runs, and with None on unbind_proc.
    Use it to persist the live pid so cancellation survives a host restart.
    """
    if not run_id:
        return
    with _lock:
        _events[run_id] = threading.Event()
        # Drop any stale proc registration carrying over from a prior attempt.
        _procs.pop(run_id, None)
        if on_proc_change is not None:
            _on_proc_change[run_id] = on_proc_change
        else:
            _on_proc_change.pop(run_id, None)
    _thread_local.run_id = run_id


def end_run(run_id: str) -> None:
    """Release per-run resources. Safe to call multiple times."""
    if not run_id:
        return
    with _lock:
        _events.pop(run_id, None)
        _procs.pop(run_id, None)
        cb = _on_proc_change.pop(run_id, None)
    # Notify outside the lock — callback may touch DB / IO.
    if cb is not None:
        try:
            cb(None)
        except Exception:
            pass
    if getattr(_thread_local, "run_id", None) == run_id:
        _thread_local.run_id = None


def request_cancel(run_id: str) -> bool:
    """Signal cancellation + kill bound subprocess (if any).

    Returns True if `run_id` was registered and we attempted to stop it,
    False if `run_id` is unknown (already finished, or never registered).
    """
    if not run_id:
        return False
    with _lock:
        ev = _events.get(run_id)
        proc = _procs.get(run_id)
    if ev is None:
        return False
    ev.set()
    if proc is not None and proc.poll() is None:
        try:
            proc.kill()
        except Exception:
            # Process may have died between poll() and kill() — fine.
            pass
    return True


def is_cancelled(run_id: Optional[str] = None) -> bool:
    """Check whether the given run (or thread-local current run) has been
    asked to cancel. Returns False if no run is bound."""
    if run_id is None:
        run_id = getattr(_thread_local, "run_id", None)
    if not run_id:
        return False
    with _lock:
        ev = _events.get(run_id)
    return ev.is_set() if ev is not None else False


def bind_proc(proc: subprocess.Popen, run_id: Optional[str] = None) -> None:
    """Associate a subprocess with the current run so request_cancel() can
    kill it. The most recently bound proc replaces any prior one — providers
    typically run one subprocess at a time per run.

    Also fires the per-run on_proc_change callback (if any) with proc.pid so
    the host can persist it durably (e.g. write to DB) for restart-resilient
    cancellation.
    """
    if run_id is None:
        run_id = getattr(_thread_local, "run_id", None)
    if not run_id:
        return
    with _lock:
        _procs[run_id] = proc
        cb = _on_proc_change.get(run_id)
    if cb is not None:
        try:
            cb(proc.pid)
        except Exception:
            pass


def unbind_proc(run_id: Optional[str] = None) -> None:
    """Release the proc binding (call in finally after subprocess.wait)."""
    if run_id is None:
        run_id = getattr(_thread_local, "run_id", None)
    if not run_id:
        return
    with _lock:
        _procs.pop(run_id, None)
        cb = _on_proc_change.get(run_id)
    if cb is not None:
        try:
            cb(None)
        except Exception:
            pass


def current_run_id() -> Optional[str]:
    """Diagnostic: what run is this thread bound to?"""
    return getattr(_thread_local, "run_id", None)
