from __future__ import annotations

from threading import Event, Lock


class RunRequestTracker:
    """
    Request-level tracking state extracted from FlowRunRef.

    This tracker owns:
    - in-flight accounting
    - per-request pending counters
    - per-request source sequence
    - request completion events

    FlowRunRef remains responsible for run lifecycle transitions
    (ingress state + run status).
    """

    def __init__(self) -> None:
        self._inflight_requests: int = 0
        self._request_pending: dict[str, int] = {}
        self._request_producer_done: dict[str, bool] = {}
        self._request_source_seq: dict[str, int] = {}
        self._request_events: dict[str, Event] = {}
        self._lock = Lock()

    def inflight_requests(self) -> int:
        with self._lock:
            return self._inflight_requests

    def ack_inflight_delta(self, delta: int) -> int:
        with self._lock:
            next_value = self._inflight_requests + delta
            if next_value < 0:
                raise ValueError("FlowRun in-flight request counter cannot be negative.")
            self._inflight_requests = next_value
            return self._inflight_requests

    def start_request(self, request_id: str, initial_delta: int = 1) -> Event:
        if not request_id:
            raise ValueError("request_id must be non-empty.")
        if initial_delta < 0:
            raise ValueError("initial_delta must be non-negative.")

        with self._lock:
            event = self._request_events.get(request_id)
            if event is None:
                event = Event()
                self._request_events[request_id] = event

            pending = self._request_pending.get(request_id, 0) + initial_delta
            self._request_pending[request_id] = pending
            self._request_producer_done.setdefault(request_id, False)
            self._request_source_seq.setdefault(request_id, 0)

            if initial_delta > 0:
                self._inflight_requests += initial_delta
            if pending <= 0 and self._request_producer_done.get(request_id, False):
                self._complete_request_locked(request_id, event=event)
            return event

    def request_pending(self, request_id: str) -> int | None:
        with self._lock:
            return self._request_pending.get(request_id)

    def request_producer_done(self, request_id: str) -> bool:
        with self._lock:
            if request_id not in self._request_pending:
                return True
            return bool(self._request_producer_done.get(request_id, False))

    def request_done(self, request_id: str) -> bool:
        with self._lock:
            if request_id not in self._request_pending:
                return True
            event = self._request_events.get(request_id)
            if event is None:
                return False
            return event.is_set()

    def mark_producer_done(self, request_id: str) -> bool:
        if not request_id:
            return False
        with self._lock:
            if request_id not in self._request_pending:
                return True
            self._request_producer_done[request_id] = True
            pending = int(self._request_pending.get(request_id, 0))
            if pending > 0:
                return False
            self._complete_request_locked(request_id)
            return True

    def mark_all_producer_done(self) -> tuple[str, ...]:
        with self._lock:
            request_ids = tuple(self._request_pending.keys())
            for request_id in request_ids:
                self._request_producer_done[request_id] = True
            for request_id in tuple(self._request_pending.keys()):
                pending = int(self._request_pending.get(request_id, 0))
                if pending <= 0:
                    self._complete_request_locked(request_id)
            return request_ids

    def pending_request_ids(self) -> tuple[str, ...]:
        with self._lock:
            return tuple(self._request_pending.keys())

    def request_source_final_seq(self, request_id: str) -> int | None:
        with self._lock:
            seq = self._request_source_seq.get(request_id)
            if seq is None or seq <= 0:
                return None
            return seq

    def next_request_meta(self, request_id: str) -> dict[str, int]:
        with self._lock:
            next_source_seq = self._request_source_seq.get(request_id, 0) + 1
            self._request_source_seq[request_id] = next_source_seq
        return {"source_seq": next_source_seq}

    def ack_request_delta(self, request_id: str, delta: int) -> int:
        if not request_id:
            return 0

        with self._lock:
            pending = self._request_pending.get(request_id)
            if pending is None:
                return 0

            next_pending = pending + delta
            if next_pending < 0:
                raise ValueError("FlowRun request pending counter cannot be negative.")

            next_inflight = self._inflight_requests + delta
            if next_inflight < 0:
                raise ValueError("FlowRun in-flight request counter cannot be negative.")

            self._inflight_requests = next_inflight
            if next_pending == 0:
                if self._request_producer_done.get(request_id, False):
                    self._complete_request_locked(request_id)
                else:
                    # Request event-chain is drained but ingress side has not
                    # marked producer_done yet; keep request open.
                    self._request_pending[request_id] = 0
            else:
                self._request_pending[request_id] = next_pending
            return next_pending

    def _complete_request_locked(self, request_id: str, *, event: Event | None = None) -> None:
        self._request_pending.pop(request_id, None)
        self._request_source_seq.pop(request_id, None)
        self._request_producer_done.pop(request_id, None)
        stored_event = self._request_events.pop(request_id, None)
        request_event = event or stored_event
        if request_event is not None:
            request_event.set()


__all__ = ["RunRequestTracker"]
