from __future__ import annotations

from dataclasses import dataclass
from threading import Event, Lock
import time
from typing import Any

from flownet.core.request_id import new_stream_request_id


@dataclass
class _StreamRunState:
    source_id: str | None = None
    pending: int = 0
    done_marked: bool = False
    completed: bool = False
    updated_at: float = 0.0


class StreamRunTracker:
    def __init__(self):
        self._states: dict[str, _StreamRunState] = {}
        self._events: dict[str, Event] = {}
        self._request_sources: dict[str, str] = {}
        self._stats = {
            "tracked_requests": 0,
            "completed_requests": 0,
            "ack_events": 0,
            "done_events": 0,
            "underflow_events": 0,
        }
        self._lock = Lock()

    def new_request_id(self, source_id: str) -> str:
        request_id = new_stream_request_id(source_id)
        with self._lock:
            self._request_sources[request_id] = source_id
            self._ensure_state_locked(request_id)
        return request_id

    def start_request(self, request_id: str, initial_delta: int = 1) -> Event:
        self.event_chain_done(request_id, initial_delta)
        return self._event_for(request_id)

    def event_chain_done(self, request_id: str, delta: int) -> int:
        if not request_id:
            return 0
        with self._lock:
            self._stats["ack_events"] += 1
            state = self._ensure_state_locked(request_id)
            next_pending = state.pending + int(delta)
            if next_pending < 0:
                # Defensive clamp: a missing +1 should not crash runtime packet path.
                self._stats["underflow_events"] += 1
                next_pending = 0
            state.pending = next_pending
            state.updated_at = time.time()
            self._maybe_complete_locked(request_id, state)
            return state.pending

    def mark_producer_done(self, request_id: str) -> bool:
        if not request_id:
            return False
        with self._lock:
            self._stats["done_events"] += 1
            state = self._ensure_state_locked(request_id)
            state.done_marked = True
            state.updated_at = time.time()
            return self._maybe_complete_locked(request_id, state)

    def mark_done(self, request_id: str) -> bool:
        # Backward-compatible alias during producer-done naming migration.
        return self.mark_producer_done(request_id)

    def request_status(self, request_id: str) -> dict[str, Any]:
        with self._lock:
            state = self._ensure_state_locked(request_id)
            event_chain_pending = int(state.pending)
            producer_done = bool(state.done_marked)
            request_done = bool(state.completed)
            if request_done:
                status = "completed"
            elif producer_done and event_chain_pending > 0:
                status = "quiescing"
            else:
                status = "running"
            return {
                "request_id": request_id,
                "request_ref_id": request_id,
                "source_id": state.source_id,
                "pending": event_chain_pending,
                "done_marked": producer_done,
                "completed": request_done,
                "event_chain_pending": event_chain_pending,
                "event_chain_done": event_chain_pending == 0,
                "producer_done": producer_done,
                "request_done": request_done,
                "status": status,
                "final_seq": None,
                "expected_total_events": None,
            }

    def wait_done(self, request_id: str, timeout: float | None = None) -> bool:
        event = self._event_for(request_id)
        return event.wait(timeout=timeout)

    def stats(self) -> dict[str, int]:
        with self._lock:
            stats = dict(self._stats)
            stats["active_requests"] = sum(
                1
                for state in self._states.values()
                if state.pending > 0 or not state.done_marked
            )
            return stats

    def request_status_snapshot(self, request_id: str) -> dict[str, Any]:
        return self.request_status(request_id)

    def _event_for(self, request_id: str) -> Event:
        with self._lock:
            event = self._events.get(request_id)
            if event is None:
                event = Event()
                self._events[request_id] = event
            return event

    def _ensure_state_locked(self, request_id: str) -> _StreamRunState:
        state = self._states.get(request_id)
        if state is not None:
            if state.source_id is None:
                state.source_id = self._request_sources.get(request_id)
            return state
        state = _StreamRunState(
            source_id=self._request_sources.get(request_id),
            updated_at=time.time(),
        )
        self._states[request_id] = state
        self._stats["tracked_requests"] += 1
        return state

    def _maybe_complete_locked(self, request_id: str, state: _StreamRunState) -> bool:
        if state.completed:
            return True
        if not state.done_marked:
            return False
        if state.pending != 0:
            return False
        state.completed = True
        self._stats["completed_requests"] += 1
        event = self._events.get(request_id)
        if event is None:
            event = Event()
            self._events[request_id] = event
        event.set()
        return True


__all__ = ["StreamRunTracker"]
