"""Single-flight (request coalescing): dedupe concurrent calls for one key.

Only one caller (the leader) actually runs ``fn`` for a given key; every other
caller that arrives while it is in flight waits and receives the same result or
exception. Once the in-flight call completes, the key is freed so the next call
runs ``fn`` again.

Correctness notes:
- A global lock protects the in-flight registry; each in-flight call has its
  own condition variable for waiters.
- The leader runs ``fn`` outside any lock so other keys (and registry
  operations) are not blocked.
- Result or exception is broadcast to all waiters; waiters re-raise the same
  exception object the leader raised.
- Registry entry is removed in a ``finally`` so a crashing leader cannot leak
  a stuck key.
- ``fn`` is called exactly once per in-flight window.
"""
from __future__ import annotations

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


class _InFlight:
    __slots__ = ("key", "fn", "cond", "result", "exc", "done", "waiters")

    def __init__(self, key: str, fn: Callable[[], Any]) -> None:
        self.key = key
        self.fn = fn
        self.cond = threading.Condition(threading.Lock())
        self.result: Any = None
        self.exc: Optional[BaseException] = None
        self.done: bool = False
        self.waiters: int = 0


class SingleFlight:
    def __init__(self) -> None:
        self._mu = threading.Lock()
        self._inflight: Dict[Any, _InFlight] = {}

    @property
    def inflight_count(self) -> int:
        with self._mu:
            return len(self._inflight)

    def do(self, key: Any, fn: Callable[[], Any]) -> Any:
        # Register (or join) an in-flight call.
        with self._mu:
            c = self._inflight.get(key)
            leader = c is None
            if leader:
                c = _InFlight(key, fn)
                self._inflight[key] = c
            c.waiters += 1

        if leader:
            try:
                result = fn()
            except BaseException as e:  # noqa: BLE001 - must propagate anything
                with c.cond:
                    c.exc = e
                    c.done = True
                    c.cond.notify_all()
                raise
            else:
                with c.cond:
                    c.result = result
                    c.done = True
                    c.cond.notify_all()
                return result
            finally:
                with self._mu:
                    # Only remove if still ours (defensive against re-registration).
                    if self._inflight.get(key) is c:
                        del self._inflight[key]
        # Follower path.
        with c.cond:
            while not c.done:
                c.cond.wait()
            if c.exc is not None:
                raise c.exc
            return c.result