# src/fmatch/worker_pool.py
from __future__ import annotations

import logging
import multiprocessing as mp
import queue
import threading
import time
import traceback
from concurrent.futures import Future, ProcessPoolExecutor
from threading import Event, Lock
from typing import Any, Callable, List, Dict, Optional

__all__ = ["WorkerPool", "get_pool", "shutdown_pool"]

log = logging.getLogger(__name__)

# ──────────────────────────────────────────────────────────────
# Live-metrics queue (shared across all processes)
# ──────────────────────────────────────────────────────────────
_metrics_mgr: mp.Manager | None = None
_metrics_q: mp.Queue | None = None


def get_metrics_queue() -> mp.Queue:
    """Singleton multiprocessing.Queue for (timestamp, rows_processed)."""
    global _metrics_mgr, _metrics_q
    if _metrics_q is None:
        _metrics_mgr = mp.Manager()
        _metrics_q = _metrics_mgr.Queue()
    return _metrics_q


def report_progress(rows: int) -> None:
    """Cheap, non-blocking push used by workers to emit progress."""
    q = _metrics_q or get_metrics_queue()
    q.put((time.time(), rows))


# ──────────────────────────────────────────────────────────────
# Worker-pool implementation
# ──────────────────────────────────────────────────────────────
def _log_future_exc(fut: Future) -> None:
    exc = fut.exception()
    if exc is not None:
        log.critical(
            "Worker task crashed:\n%s",
            "".join(traceback.format_exception(type(exc), exc, exc.__traceback__)),
        )


class WorkerPool:
    """
    Facade over ProcessPoolExecutor **with an optional in-process queue mode**
    (used by pytest to avoid heavy forking).
    """

    # ── lifecycle ─────────────────────────────────────────────
    def __init__(self, max_workers: Optional[int] = None):
        self._executor = ProcessPoolExecutor(max_workers=max_workers)
        self.max_workers: int = self._executor._max_workers  # type: ignore[attr-defined]

        self._errors: List[Dict[str, Any]] = []
        self._errors_lock = Lock()

        self._queue: queue.Queue[
            tuple[Callable[..., Any], tuple[Any, ...], Dict[str, Any], Future]
        ] = queue.Queue()
        self._shutdown_evt = Event()
        self._thread: threading.Thread | None = None  # in-process consumer
        self._is_closed = False

        log.info("WorkerPool created (max_workers=%s)", self.max_workers)

    # ── task submission ───────────────────────────────────────
    def submit(self, fn: Callable[..., Any], *args: Any, **kwargs: Any) -> Future:
        """Submit a task; route to queue if queue-mode is active."""
        if self._is_closed:
            raise RuntimeError("WorkerPool is closed; cannot submit new tasks.")

        # in-process queue path (pytest)
        if self._thread is not None:
            fut: Future = Future()
            self._queue.put((fn, args, kwargs, fut))
            return fut

        # normal multi-process path
        def _wrapper(*_a: Any, **_kw: Any) -> Any:
            try:
                return fn(*_a, **_kw)
            except Exception as exc:  # pragma: no cover
                err = {"error": repr(exc), "traceback": traceback.format_exc()}
                with self._errors_lock:
                    self._errors.append(err)
                log.error("Worker task failed: %s", exc, exc_info=False)
                raise

        fut = self._executor.submit(_wrapper, *args, **kwargs)
        fut.add_done_callback(_log_future_exc)
        return fut

    # ── single-process queue mode (for pytest) ────────────────
    def start(self) -> None:
        """Kick off the in-process consumer thread once."""
        if self._thread is not None:
            return  # already started

        log.info("Starting WorkerPool in single-process queue mode.")

        def _run_loop() -> None:
            while not self._shutdown_evt.is_set():
                try:
                    fn, args, kwargs, fut = self._queue.get(timeout=0.1)
                except queue.Empty:
                    continue
                try:
                    fut.set_result(fn(*args, **kwargs))
                except Exception as exc:  # pragma: no cover
                    fut.set_exception(exc)
                    with self._errors_lock:
                        self._errors.append(
                            {"error": repr(exc), "traceback": traceback.format_exc()}
                        )
                finally:
                    self._queue.task_done()

        self._thread = threading.Thread(
            target=_run_loop, name="WorkerPoolInProcess", daemon=True
        )
        self._thread.start()

    # ── shutdown ──────────────────────────────────────────────
    def close(self, *, wait: bool = True) -> None:
        """Gracefully shut down executor (and thread, if any)."""
        if self._is_closed:
            return

        log.info("Closing WorkerPool (wait=%s)…", wait)
        self._is_closed = True
        self._shutdown_evt.set()

        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=2.0)

        self._executor.shutdown(wait=wait, cancel_futures=False)
        log.info("WorkerPool closed.")

    # ── diagnostics helpers ───────────────────────────────────
    def healthy(self) -> bool:
        return not self._is_closed and not self._errors

    def summary(self) -> Dict[str, Any]:
        try:
            qsize = self._executor._work_queue.qsize()  # type: ignore[attr-defined]
        except AttributeError:  # pragma: no cover
            qsize = "N/A"

        return {
            "closed": self._is_closed,
            "mode": "in_process_queue" if self._thread else "process_pool",
            "max_workers": self.max_workers,
            "in_flight_approx": qsize,
            "captured_errors": len(self._errors),
        }

    @property
    def errors(self) -> List[Dict[str, Any]]:
        with self._errors_lock:
            return list(self._errors)


# ──────────────────────────────────────────────────────────────
# Per-process singleton helpers
# ──────────────────────────────────────────────────────────────
_POOL: WorkerPool | None = None
_POOL_LOCK = Lock()


def get_pool(*, max_workers: Optional[int] = None) -> WorkerPool:
    """
    Return (or create) the WorkerPool singleton.
    Ensures the metrics queue exists **before** the first fork.
    """
    global _POOL
    if _POOL is None:
        with _POOL_LOCK:
            if _POOL is None:  # double-checked
                get_metrics_queue()  # metrics queue first
                _POOL = WorkerPool(max_workers=max_workers)
    return _POOL


def shutdown_pool(*, wait: bool = True) -> None:
    """
    Close the pool **and** tear down the metrics manager/queue.
    Avoids resource-tracker warnings at interpreter exit.
    """
    global _POOL, _metrics_mgr, _metrics_q

    if _POOL is not None:
        _POOL.close(wait=wait)
        _POOL = None

    if _metrics_q is not None:
        try:
            _metrics_q.close()
        except Exception:
            log.debug("metrics queue close failed", exc_info=True)
        _metrics_q = None

    if _metrics_mgr is not None:
        try:
            _metrics_mgr.shutdown()
        except Exception:
            log.debug("metrics manager shutdown failed", exc_info=True)
        _metrics_mgr = None
