"""Bounded process pool for the nightly PDF parse stages (certificates and audit reports).

Workers only fetch and parse; the parent records every outcome, so the SQLite work DB keeps a
single writer. Design: docs/superpowers/specs/2026-10-05-nightly-parallel-pdf-parsing-design.md.
"""

from __future__ import annotations

import time
from collections import deque
from concurrent.futures import FIRST_COMPLETED, Future, ProcessPoolExecutor, wait
from concurrent.futures.process import BrokenProcessPool
from dataclasses import asdict, dataclass, field
from typing import Callable, Generic, Mapping, Sequence, TypeVar

T = TypeVar("T")
R = TypeVar("R")

# A pool broken by a dying worker fails every in-flight future, so the culprit is unknown.
# Each affected item gets this many fresh attempts before it is recorded as a crash.
BROKEN_POOL_RETRIES = 1


@dataclass
class PoolStats:
    workers: int = 0
    submitted: int = 0
    completed: int = 0
    timeouts: int = 0
    crashes: int = 0
    recycles: int = 0
    stopped_no_time: int = 0
    wall_s: float = 0.0
    peak_child_rss_mb: float | None = None
    peak_in_flight: int = 0
    peak_by_capacity: dict[str, int] = field(default_factory=dict)
    submitted_by_capacity: dict[str, int] = field(default_factory=dict)
    completed_by_capacity: dict[str, int] = field(default_factory=dict)
    stopped_by_capacity: dict[str, int] = field(default_factory=dict)
    first_admission_s: dict[str, float] = field(default_factory=dict)
    last_completion_s: dict[str, float] = field(default_factory=dict)

    def as_dict(self) -> dict:
        return asdict(self)


@dataclass
class _Pending(Generic[T]):
    item: T
    attempts: int = 0
    deadline: float | None = None


def peak_child_rss_mb() -> float | None:
    """Largest resident set of any reaped child process (Linux/macOS); None on Windows."""
    try:
        import resource
    except ImportError:
        return None
    peak = resource.getrusage(resource.RUSAGE_CHILDREN).ru_maxrss
    # Linux reports KiB, macOS bytes.
    import sys

    return round(peak / (1024 * 1024 if sys.platform == "darwin" else 1024), 1)


def _terminate(pool: ProcessPoolExecutor) -> None:
    """Kill the pool's worker processes; shutdown(wait=False) alone leaves a wedged one running."""
    # _processes is private but has been the executor's pid -> Process map since Python 3.2.
    for process in list((getattr(pool, "_processes", None) or {}).values()):
        try:
            process.terminate()
        except Exception:  # noqa: BLE001 -- already gone
            pass
    pool.shutdown(wait=False, cancel_futures=True)


def run_bounded_pool(
    items: Sequence[T],
    work: Callable[[T], R],
    record: Callable[[T, R], None],
    *,
    workers: int,
    on_timeout: Callable[[T], R],
    on_crash: Callable[[T, BaseException], R],
    should_continue: Callable[[], bool] | None = None,
    lane: Callable[[T], str] | None = None,
    lane_limits: Mapping[str, int] | None = None,
    capacity_keys: Callable[[T], Sequence[str]] | None = None,
    capacity_limits: Mapping[str, int] | None = None,
    item_timeout_s: float = 180.0,
    task_timeout_s: Callable[[T], float] | None = None,
    max_tasks_per_child: int = 25,
    initializer: Callable[..., None] | None = None,
    initargs: tuple = (),
) -> PoolStats:
    """Run ``work`` over ``items`` in worker processes, recording each outcome in the parent.

    ``work`` must be a picklable top-level function. Submission follows ``items`` order, except
    that an item whose lane is at its limit waits while later items from other lanes start.
    Alternatively ``capacity_keys`` supplies multiple independent capacity names per item;
    every matching capacity must have space. The legacy lane arguments remain unchanged.
    ``should_continue`` is asked before every submission; once it says no, nothing new starts,
    in-flight items finish, and the never-started count lands in ``stopped_no_time``. An item
    running longer than ``item_timeout_s`` is recorded through ``on_timeout`` and the worker
    processes are killed; innocent in-flight items are re-queued at the front. ``initializer``
    runs once in every worker process (also after a recycle) to set up per-process state.
    ``task_timeout_s``, when supplied, returns the fixed timeout for each item in the parent;
    omission retains the legacy shared ``item_timeout_s`` deadline.
    """
    workers = max(1, int(workers))
    if capacity_keys is not None and (lane is not None or lane_limits is not None):
        raise ValueError("use capacity_keys/capacity_limits or lane/lane_limits, not both")
    limits = {name: max(1, int(limit)) for name, limit in (lane_limits or {}).items()}
    lane_of = lane or (lambda _item: "")
    if capacity_keys is not None:
        limits = dict(capacity_limits or {})
        if any(type(limit) is not int or limit < 1 for limit in limits.values()):
            raise ValueError("capacity limits must be positive integers")

    def keys_of(item: T) -> tuple[str, ...]:
        return tuple(dict.fromkeys(capacity_keys(item))) if capacity_keys else (lane_of(item),)

    def expires_at(entry: _Pending[T], submitted_at: float) -> float:
        # Mixed deadlines belong to the logical document, including innocent preemptions.
        # Legacy shared-timeout mode deliberately retains its original per-submission clock.
        return entry.deadline if task_timeout_s is not None else submitted_at + item_timeout_s

    stats = PoolStats(workers=workers)
    started = time.monotonic()
    pending: deque[_Pending[T]] = deque(_Pending(item) for item in items)
    stopping = False

    def _next_submittable(running_lanes: dict[str, int]) -> _Pending[T] | None:
        for index, entry in enumerate(pending):
            if all(running_lanes.get(name, 0) < limits.get(name, workers)
                   for name in keys_of(entry.item)):
                del pending[index]
                return entry
        return None

    def _finish(item: T, result: R) -> None:
        record(item, result)
        stats.completed += 1
        for name in keys_of(item):
            stats.completed_by_capacity[name] = stats.completed_by_capacity.get(name, 0) + 1
            stats.last_completion_s[name] = time.monotonic() - started

    while pending and not stopping:
        pool = ProcessPoolExecutor(
            max_workers=workers,
            max_tasks_per_child=max_tasks_per_child,
            initializer=initializer,
            initargs=initargs,
        )
        in_flight: dict[Future, tuple[_Pending[T], float]] = {}
        running_lanes: dict[str, int] = {}
        recycle = False
        try:
            while (pending and not stopping) or in_flight:
                while not stopping and pending and len(in_flight) < workers:
                    if should_continue is not None and not should_continue():
                        stopping = True
                        break
                    entry = _next_submittable(running_lanes)
                    if entry is None:
                        break
                    if task_timeout_s is not None:
                        now = time.monotonic()
                        if entry.deadline is None:
                            entry.deadline = now + task_timeout_s(entry.item)
                        elif entry.deadline <= now:
                            stats.timeouts += 1
                            _finish(entry.item, on_timeout(entry.item))
                            continue
                    in_flight[pool.submit(work, entry.item)] = (entry, time.monotonic())
                    for name in keys_of(entry.item):
                        running_lanes[name] = running_lanes.get(name, 0) + 1
                        stats.peak_by_capacity[name] = max(stats.peak_by_capacity.get(name, 0), running_lanes[name])
                        stats.submitted_by_capacity[name] = stats.submitted_by_capacity.get(name, 0) + 1
                        stats.first_admission_s.setdefault(name, time.monotonic() - started)
                    stats.submitted += 1
                    stats.peak_in_flight = max(stats.peak_in_flight, len(in_flight))
                if not in_flight:
                    break
                next_deadline = min(expires_at(entry, submitted_at)
                                    for entry, submitted_at in in_flight.values())
                wait_s = max(0.0, next_deadline - time.monotonic())
                finished, _ = wait(set(in_flight), timeout=wait_s, return_when=FIRST_COMPLETED)
                broken: list[_Pending[T]] = []
                for future in finished:
                    entry, _submitted_at = in_flight.pop(future)
                    for name in keys_of(entry.item):
                        running_lanes[name] -= 1
                    try:
                        result = future.result()
                    except BrokenProcessPool:
                        broken.append(entry)
                        continue
                    except Exception as exc:  # noqa: BLE001 -- work() raised inside the worker
                        stats.crashes += 1
                        _finish(entry.item, on_crash(entry.item, exc))
                        continue
                    _finish(entry.item, result)
                if broken:
                    # Everything still in flight is lost with the pool too.
                    broken.extend(entry for entry, _submitted_at in in_flight.values())
                    in_flight.clear()
                    for entry in reversed(broken):
                        if entry.attempts < BROKEN_POOL_RETRIES:
                            entry.attempts += 1
                            pending.appendleft(entry)
                        else:
                            stats.crashes += 1
                            _finish(entry.item, on_crash(entry.item, BrokenProcessPool("worker process died")))
                    recycle = True
                    break
                now = time.monotonic()
                overdue = [
                    future for future, (entry, submitted_at) in in_flight.items()
                    if now >= expires_at(entry, submitted_at)
                ]
                if overdue:
                    for future in overdue:
                        entry, _submitted_at = in_flight.pop(future)
                        stats.timeouts += 1
                        _finish(entry.item, on_timeout(entry.item))
                    for entry, _submitted_at in reversed(list(in_flight.values())):
                        pending.appendleft(entry)
                    in_flight.clear()
                    recycle = True
                    break
        except BaseException:
            _terminate(pool)
            raise
        if recycle:
            stats.recycles += 1
            _terminate(pool)
        else:
            pool.shutdown(wait=True)
    stats.stopped_no_time = len(pending)
    for entry in pending:
        for name in keys_of(entry.item):
            stats.stopped_by_capacity[name] = stats.stopped_by_capacity.get(name, 0) + 1
    stats.wall_s = round(time.monotonic() - started, 1)
    stats.peak_child_rss_mb = peak_child_rss_mb()
    return stats