from __future__ import annotations

import threading
from typing import Any


_lock = threading.Lock()
_request_runs: dict[str, Any] = {}


def register_run_request(run: Any, request_id: str, *, initial_delta: int = 1):
    start_request = getattr(run, "start_request", None)
    if not callable(start_request):
        raise RuntimeError(
            "FlowDef requires runtime request tracker support. "
            "Run must implement start_request(request_id, initial_delta)."
        )
    event = start_request(request_id, initial_delta=initial_delta)
    with _lock:
        _request_runs[request_id] = run
    return event


def ack_run_request(request_id: str, delta: int) -> int:
    with _lock:
        run = _request_runs.get(request_id)
    if run is None:
        return 0

    ack_request_delta = getattr(run, "ack_request_delta", None)
    if not callable(ack_request_delta):
        with _lock:
            _request_runs.pop(request_id, None)
        raise RuntimeError(
            "FlowDef requires runtime request tracker support. "
            "Run must implement ack_request_delta(request_id, delta)."
        )

    pending = ack_request_delta(request_id, delta)
    if pending <= 0:
        with _lock:
            _request_runs.pop(request_id, None)
    return pending


def reset_request_registry() -> None:
    with _lock:
        _request_runs.clear()


__all__ = [
    "register_run_request",
    "ack_run_request",
    "reset_request_registry",
]
