"""
stylo — fiction-oriented single-string text metrics → feature vector

Libraries: textdescriptives (+ pos_proportions), lexicalrichness,
           textstat, NRCLex

Performance (long ~50k-char answers / live 256-tok windows):
  - spaCy: tagger + sentencizer only (no parser/NER/lemmatizer)
  - quality n-gram metrics: pure-Python (spaCy quality pipe was ~1s+/doc)
  - **Linux default: process pool + fork** — one spaCy per process; fork COW
    shares model pages so RAM stays near ~1× model, not N× thread copies
  - **batched nlp.pipe** inside workers (and a short live batch window) so
    free CPU is used without 32× en_core_web_sm RSS
  - thread pool fallback (Windows / METRIC_POOL=thread): thread-local spaCy,
    fewer workers, still batched when possible
"""

from __future__ import annotations

import os
import re
import signal
import sys
import threading
import warnings
from collections import Counter
from concurrent.futures import Future, ProcessPoolExecutor, ThreadPoolExecutor
from pathlib import Path
from typing import Any, Iterable

# Pin BLAS/OpenMP *before* numpy/spacy import — concurrent spaCy + multi-thread
# BLAS is a common SIGSEGV source under metric pools.
for _k, _v in (
    ("OMP_NUM_THREADS", "1"),
    ("OPENBLAS_NUM_THREADS", "1"),
    ("MKL_NUM_THREADS", "1"),
    ("NUMEXPR_NUM_THREADS", "1"),
    ("VECLIB_MAXIMUM_THREADS", "1"),
    ("BLIS_NUM_THREADS", "1"),
):
    os.environ.setdefault(_k, _v)

warnings.filterwarnings("ignore")

import numpy as np
import spacy
import textdescriptives as td
import textstat
from lexicalrichness import LexicalRichness
from nrclex import NRCLex


def _safe_count_chars(text: str, ignore_spaces: bool = True) -> int:
    """Pure-Python char count (textstat uses re.sub(r'\\s') which has SEGV'd
    under fork workers after the parent has embed/HTTP threads)."""
    if text is None:
        return 0
    if not isinstance(text, str):
        text = str(text)
    if "\x00" in text:
        text = text.replace("\x00", "")
    if ignore_spaces:
        return sum(1 for c in text if not c.isspace())
    return len(text)


def _safe_count_letters(text: str) -> int:
    """Pure-Python letter count; same rationale as ``_safe_count_chars``."""
    if text is None:
        return 0
    if not isinstance(text, str):
        text = str(text)
    if "\x00" in text:
        text = text.replace("\x00", "")
    try:
        from textstat.backend.transformations._remove_punctuation import (
            remove_punctuation,
        )

        no_ws = "".join(c for c in text if not c.isspace())
        return len(remove_punctuation(no_ws, rm_apostrophe=True))
    except Exception:
        return sum(1 for c in text if c.isalpha())


def _install_textstat_safe_counts() -> None:
    """Monkeypatch textstat count_chars/count_letters everywhere they are bound.

    Crash site (faulthandler): ``textstat.reading_time`` → ``count_chars`` →
    ``re.sub(r"\\s", "", text)`` in a ProcessPool worker mid-trial.
    """
    targets = (
        "textstat.backend.counts._count_chars",
        "textstat.backend.counts._count_letters",
        "textstat.backend.counts",
        "textstat.backend.metrics._reading_time",
        "textstat.backend.metrics._chars_per_word",
        "textstat.backend.metrics._letters_per_word",
    )
    for name in targets:
        try:
            mod = __import__(name, fromlist=["count_chars", "count_letters"])
        except Exception:
            continue
        if hasattr(mod, "count_chars"):
            setattr(mod, "count_chars", _safe_count_chars)
        if hasattr(mod, "count_letters"):
            setattr(mod, "count_letters", _safe_count_letters)


_install_textstat_safe_counts()


def _sanitize_metric_text(text: Any) -> str:
    """Normalize text before spaCy/textstat (NULs / non-str can corrupt C layers)."""
    if not isinstance(text, str):
        text = str(text) if text is not None else ""
    if "\x00" in text:
        text = text.replace("\x00", "")
    return text


# ---------------------------------------------------------------------------
# spaCy / NRC — thread-local (thread pool) or process-global (process pool)
# ---------------------------------------------------------------------------
_thread_local = threading.local()
# spacy.load / pipe add is not safe under concurrent first-use; serialize builds.
_NLP_BUILD_LOCK = threading.Lock()
_NRC_BUILD_LOCK = threading.Lock()
# Process-pool global (set in parent before fork for COW, and/or in initializer).
_PROCESS_NLP = None
# Persistent pool so workers keep their loaded models across trials.
_METRIC_POOL: Any = None
_METRIC_POOL_SIZE: int = 0
_METRIC_POOL_KIND: str = ""  # "process" | "thread"
_METRIC_WARNED_SUBMIT = False
_METRIC_BATCHER: "_MetricBatcher | None" = None
# Live-window batching: trade a few ms latency for pipe throughput / less IPC.
_BATCH_MAX = int(os.getenv("METRIC_BATCH_MAX") or "24")
_BATCH_WAIT_S = float(os.getenv("METRIC_BATCH_WAIT_MS") or "40") / 1000.0
# Process-pool worker PIDs (Ctrl+C / os._exit otherwise leaves them as orphans).
_KNOWN_WORKER_PIDS: set[int] = set()
# Per-process file so two concurrent optimizers never reap each other.
_WORKER_PID_FILE = Path(
    os.getenv("STYLO_PID_FILE")
    or f"/tmp/optuna-2phase-stylo-workers-{os.getuid()}-{os.getpid()}.pids"
)

# Minimal English stoplist for quality n_stop_words (avoid loading spaCy stops).
_STOP = frozenset(
    """
    a an the and or but if in on at to for of as is was are were be been being
    it its this that these those i you he she we they me him her us them my your
    his their our not no nor so than then too very can could will would should
    may might must do does did have has had with from by into about over after
    before under again further once here there when where why how all each few
    more most other some such only own same both through during above below
    """.split()
)

_WORD_RE = re.compile(r"[A-Za-z0-9']+")
_ELLIPSIS_RE = re.compile(r"\.\.\.|…")


def _build_nlp():
    """Tagger + sentencizer only — no parser/NER/lemmatizer/quality pipe."""
    nlp = spacy.load(
        "en_core_web_sm",
        disable=["parser", "ner", "lemmatizer"],
    )
    if "senter" not in nlp.pipe_names and "sentencizer" not in nlp.pipe_names:
        nlp.add_pipe("sentencizer")
    for pipe in [
        "textdescriptives/descriptive_stats",
        "textdescriptives/readability",
        "textdescriptives/pos_proportions",
        "textdescriptives/information_theory",
    ]:
        if pipe not in nlp.pipe_names:
            nlp.add_pipe(pipe)
    nlp.max_length = max(nlp.max_length, 2_000_000)
    return nlp


def _in_process_pool_worker() -> bool:
    try:
        import multiprocessing as mp

        return mp.current_process().name not in ("MainProcess",)
    except Exception:
        return False


def _get_nlp():
    """spaCy pipeline: process-global in process workers, else thread-local.

    Process pool (Linux fork): parent preloads ``_PROCESS_NLP``; children
    inherit COW pages and reuse it — one ~0.6 GiB model shared, not N copies.
    Thread pool: one model per thread (Language is not thread-safe).
    """
    global _PROCESS_NLP
    if _PROCESS_NLP is not None and (
        _METRIC_POOL_KIND == "process" or _in_process_pool_worker()
    ):
        return _PROCESS_NLP
    nlp = getattr(_thread_local, "nlp", None)
    if nlp is not None:
        return nlp
    with _NLP_BUILD_LOCK:
        nlp = getattr(_thread_local, "nlp", None)
        if nlp is None:
            nlp = _build_nlp()
            _thread_local.nlp = nlp
    return nlp


def _get_nrc() -> NRCLex:
    """Per-thread NRCLex (mutates instance state on load_token_list)."""
    nrc = getattr(_thread_local, "nrc", None)
    if nrc is not None:
        return nrc
    with _NRC_BUILD_LOCK:
        nrc = getattr(_thread_local, "nrc", None)
        if nrc is None:
            nrc = NRCLex()
            _thread_local.nrc = nrc
    return nrc


def _ensure_process_nlp():
    """Load process-global spaCy once (parent before fork, or child init)."""
    global _PROCESS_NLP
    if _PROCESS_NLP is not None:
        return _PROCESS_NLP
    with _NLP_BUILD_LOCK:
        if _PROCESS_NLP is None:
            _PROCESS_NLP = _build_nlp()
    return _PROCESS_NLP


def _process_pool_initializer() -> None:
    """Child process entry: BLAS pins already set at import; ensure spaCy ready."""
    for _k, _v in (
        ("OMP_NUM_THREADS", "1"),
        ("OPENBLAS_NUM_THREADS", "1"),
        ("MKL_NUM_THREADS", "1"),
        ("NUMEXPR_NUM_THREADS", "1"),
        ("VECLIB_MAXIMUM_THREADS", "1"),
        ("BLIS_NUM_THREADS", "1"),
    ):
        os.environ.setdefault(_k, _v)
    # Re-apply in case child import order re-bound textstat names.
    _install_textstat_safe_counts()
    # Fork COW: parent already loaded _PROCESS_NLP — reuse, do not rebuild.
    _ensure_process_nlp()
    _get_nrc()


def _total_ram_gb() -> float | None:
    """Best-effort host RAM (GiB). Used to avoid multi-spaCy OOM on small e2."""
    # Linux (typical e2)
    try:
        with open("/proc/meminfo", encoding="utf-8") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    kb = float(line.split()[1])
                    return kb / (1024.0 * 1024.0)
    except Exception:
        pass
    # Windows
    try:
        import ctypes

        class _MemStatus(ctypes.Structure):
            _fields_ = [
                ("dwLength", ctypes.c_ulong),
                ("dwMemoryLoad", ctypes.c_ulong),
                ("ullTotalPhys", ctypes.c_ulonglong),
                ("ullAvailPhys", ctypes.c_ulonglong),
                ("ullTotalPageFile", ctypes.c_ulonglong),
                ("ullAvailPageFile", ctypes.c_ulonglong),
                ("ullTotalVirtual", ctypes.c_ulonglong),
                ("ullAvailVirtual", ctypes.c_ulonglong),
                ("sullAvailExtendedVirtual", ctypes.c_ulonglong),
            ]

        m = _MemStatus()
        m.dwLength = ctypes.sizeof(_MemStatus)
        if ctypes.windll.kernel32.GlobalMemoryStatusEx(ctypes.byref(m)):
            return float(m.ullTotalPhys) / (1024.0 ** 3)
    except Exception:
        pass
    return None


# Measured on this stack after warm extracts (process children dirty spaCy pages
# so fork COW does *not* keep children at a few MB — each is ~0.65 GiB RSS).
_STYLO_PER_PROCESS_GIB = 0.65
_STYLO_PARENT_GIB = 0.75  # parent holds one spaCy for fork preload
_STYLO_PER_THREAD_GIB = 0.60
# Whole optimizer process tree target (GiB). We only *restrict stylo workers*
# so estimated stylo + non-stylo reserve ≤ this. Override: SCRIPT_RAM_GIB=8
_SCRIPT_RAM_GIB_DEFAULT = 8.0
# Reserved for everything that is not stylo workers: Python/numpy, NovelAI
# streams, embed HTTP buffers, Optuna/SQLite, human cache arrays, etc.
# Override: NON_STYLO_RESERVE_GIB=3
_NON_STYLO_RESERVE_GIB_DEFAULT = 3.0


def _script_ram_budget_gib() -> float:
    """Ceiling for the whole script (parent + stylo children), GiB."""
    raw = (os.getenv("SCRIPT_RAM_GIB") or "").strip()
    if not raw:
        return _SCRIPT_RAM_GIB_DEFAULT
    try:
        return max(2.0, float(raw))
    except ValueError:
        return _SCRIPT_RAM_GIB_DEFAULT


def _non_stylo_reserve_gib() -> float:
    """RAM left for non-stylo parts of the optimizer (not capped elsewhere)."""
    raw = (os.getenv("NON_STYLO_RESERVE_GIB") or "").strip()
    if not raw:
        return _NON_STYLO_RESERVE_GIB_DEFAULT
    try:
        return max(0.5, float(raw))
    except ValueError:
        return _NON_STYLO_RESERVE_GIB_DEFAULT


def _stylo_ram_budget_gib() -> float:
    """How much RAM stylo is allowed so script total stays ≤ SCRIPT_RAM_GIB.

    Default: SCRIPT_RAM_GIB − NON_STYLO_RESERVE_GIB (8 − 3 = 5 GiB).
    Explicit STYLO_RAM_GIB=N overrides the remainder math.
    """
    raw = (os.getenv("STYLO_RAM_GIB") or "").strip()
    if raw:
        try:
            return max(1.0, float(raw))
        except ValueError:
            pass
    return max(1.0, _script_ram_budget_gib() - _non_stylo_reserve_gib())


def _stylo_worker_cap_for_budget(want: int) -> int:
    """Max workers so estimated stylo RSS fits the stylo slice of SCRIPT_RAM_GIB."""
    want = max(1, int(want))
    budget = _stylo_ram_budget_gib()
    if _use_process_pool():
        # parent model + N children × measured dirty RSS
        room = budget - _STYLO_PARENT_GIB
        if room < _STYLO_PER_PROCESS_GIB:
            return 1
        return max(1, min(want, int(room / _STYLO_PER_PROCESS_GIB)))
    return max(1, min(want, int(budget / _STYLO_PER_THREAD_GIB)))


def _ram_worker_cap(want: int, *, per_worker_gib: float) -> int:
    """Cap workers by *host* free RAM (whole machine), separate from stylo budget."""
    want = max(1, int(want))
    ram = _total_ram_gb()
    if ram is None:
        return min(want, 4)
    headroom = 4.0
    by_ram = max(1, int((float(ram) - headroom) / max(per_worker_gib, 0.05)))
    return max(1, min(want, by_ram))


def _use_process_pool() -> bool:
    """Linux fork process pool by default; override with METRIC_POOL=thread|process."""
    raw = (os.getenv("METRIC_POOL") or "").strip().lower()
    if raw in ("thread", "threads", "t"):
        return False
    if raw in ("process", "proc", "p"):
        return True
    return sys.platform.startswith("linux")


def _default_worker_want() -> int:
    """Aim at CPU count (capped), not a huge fixed N of model copies."""
    cpu = os.cpu_count() or 4
    return max(2, min(int(cpu), 8))


def _metric_workers(n_texts: int = 0) -> int:
    """
    Concurrent stylo workers (processes on Linux, threads elsewhere).

    Default: ``min(cpu_count, 8)``, then capped so:

      stylo_est + NON_STYLO_RESERVE ≤ SCRIPT_RAM_GIB   (default 8 GiB)

    Only **stylo worker count** is restricted — streams/embeds/Optuna are not
    capped in code; they get the reserve. Process children ~0.65 GiB each after
    warm (spaCy dirties COW pages).

    Override:
      SCRIPT_RAM_GIB=8           whole-script ceiling (default)
      NON_STYLO_RESERVE_GIB=3    RAM reserved for non-stylo (default)
      STYLO_RAM_GIB=N            explicit stylo slice (skips remainder math)
      METRIC_WORKERS=N | 0|unlimited
      METRIC_WORKERS_FORCE=1     skip *host* RAM cap (still respects script budget)
      METRIC_POOL=process|thread
    """
    del n_texts
    env = (os.getenv("METRIC_WORKERS") or "").strip().lower()
    force = (os.getenv("METRIC_WORKERS_FORCE") or "").strip().lower() in (
        "1",
        "true",
        "yes",
    )

    if env in ("0", "unlimited", "inf", "max"):
        want = 64
    elif env:
        try:
            want = max(1, int(env))
        except ValueError:
            want = _default_worker_want()
    else:
        want = _default_worker_want()

    # Keep whole script ≤ SCRIPT_RAM_GIB by only shrinking stylo workers.
    want = _stylo_worker_cap_for_budget(want)

    if force:
        return max(1, int(want))
    per = _STYLO_PER_PROCESS_GIB if _use_process_pool() else _STYLO_PER_THREAD_GIB
    return _ram_worker_cap(want, per_worker_gib=per)


class _DaemonThreadPoolExecutor(ThreadPoolExecutor):
    """Like ThreadPoolExecutor but worker threads are daemons."""

    def _adjust_thread_count(self) -> None:  # type: ignore[override]
        import weakref

        from concurrent.futures.thread import _threads_queues, _worker

        if self._idle_semaphore.acquire(timeout=0):
            return

        def weakref_cb(_: Any, q: Any = self._work_queue) -> None:
            q.put(None)

        num_threads = len(self._threads)
        if num_threads < self._max_workers:
            thread_name = f"{self._thread_name_prefix or self}_{num_threads}"
            t = threading.Thread(
                name=thread_name,
                target=_worker,
                args=(
                    weakref.ref(self, weakref_cb),
                    self._work_queue,
                    self._initializer,
                    self._initargs,
                ),
                daemon=True,
            )
            t.start()
            self._threads.add(t)
            _threads_queues[t] = self._work_queue


def _snapshot_pool_worker_pids(pool: Any) -> set[int]:
    """PIDs of ProcessPoolExecutor children (empty for thread pools)."""
    out: set[int] = set()
    procs = getattr(pool, "_processes", None)
    if not procs:
        return out
    try:
        values = procs.values() if hasattr(procs, "values") else []
    except Exception:
        return out
    for p in list(values):
        try:
            pid = getattr(p, "pid", None)
            if pid:
                out.add(int(pid))
        except Exception:
            pass
    return out


def _write_worker_pidfile(pids: set[int]) -> None:
    try:
        _WORKER_PID_FILE.write_text(
            "\n".join(str(p) for p in sorted(pids) if p > 1) + ("\n" if pids else ""),
            encoding="utf-8",
        )
    except Exception:
        pass


def _read_worker_pidfile() -> set[int]:
    out: set[int] = set()
    try:
        if not _WORKER_PID_FILE.is_file():
            return out
        for line in _WORKER_PID_FILE.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if line.isdigit():
                out.add(int(line))
    except Exception:
        pass
    return out


def _clear_worker_pidfile() -> None:
    try:
        _WORKER_PID_FILE.unlink(missing_ok=True)
    except Exception:
        pass


def _kill_pids(pids: set[int], *, sig: int = signal.SIGKILL) -> int:
    """Best-effort kill; returns how many signals were sent."""
    n = 0
    me = os.getpid()
    for pid in list(pids):
        if pid <= 1 or pid == me:
            continue
        try:
            os.kill(pid, sig)
            n += 1
        except ProcessLookupError:
            pass
        except PermissionError:
            pass
        except Exception:
            pass
    return n


def register_metric_pool_workers(pool: Any | None = None) -> None:
    """Record process-pool PIDs so Ctrl+C can kill them (and next run reaps).

    Writes **only** currently known live pool PIDs (does not merge with a stale
    pidfile — that could SIGKILL a recycled PID after a long run).
    """
    global _KNOWN_WORKER_PIDS
    pool = pool if pool is not None else _METRIC_POOL
    if pool is None:
        return
    pids = _snapshot_pool_worker_pids(pool)
    if not pids:
        return
    _KNOWN_WORKER_PIDS.update(pids)
    _write_worker_pidfile(set(_KNOWN_WORKER_PIDS))


def _pid_cmdline(pid: int) -> str:
    try:
        with open(f"/proc/{pid}/cmdline", "rb") as f:
            return f.read().replace(b"\0", b" ").decode("utf-8", "replace")
    except Exception:
        return ""


def _looks_like_our_optimizer(pid: int) -> bool:
    """True if PID is still this project’s novelai_preset_optimize process."""
    cmd = _pid_cmdline(pid)
    return "novelai_preset_optimize.py" in cmd


def reap_stale_metric_workers() -> int:
    """Kill stylo process workers left behind by a previous Ctrl+C / os._exit.

    Scans this user’s ``stylo-workers-*-*.pids`` files. Only kills PIDs that:
      - still exist, and
      - cmdline still contains novelai_preset_optimize.py, and
      - are **not** children of a *live* optimizer main (another concurrent run).
    """
    me = os.getpid()
    total = 0
    # Own file + any leftover files from dead parent PIDs of this user.
    files: list[Path] = []
    try:
        for p in Path("/tmp").glob(f"optuna-2phase-stylo-workers-{os.getuid()}-*.pids"):
            files.append(p)
    except Exception:
        pass
    if _WORKER_PID_FILE not in files and _WORKER_PID_FILE.is_file():
        files.append(_WORKER_PID_FILE)

    for path in files:
        try:
            text = path.read_text(encoding="utf-8")
        except Exception:
            continue
        # If file is for another still-running main (…-<otherpid>.pids), skip
        # reaping that main’s workers while it lives.
        try:
            owner = int(path.stem.rsplit("-", 1)[-1])
        except ValueError:
            owner = 0
        if owner and owner != me:
            try:
                os.kill(owner, 0)
                if _looks_like_our_optimizer(owner):
                    continue  # peer optimizer still running — do not touch
            except ProcessLookupError:
                pass  # owner dead → reap its workers
            except Exception:
                continue

        stale: set[int] = set()
        for line in text.splitlines():
            line = line.strip()
            if line.isdigit():
                stale.add(int(line))
        alive: set[int] = set()
        for pid in stale:
            if pid == me:
                continue
            try:
                os.kill(pid, 0)
            except ProcessLookupError:
                continue
            except PermissionError:
                continue
            except Exception:
                continue
            if not _looks_like_our_optimizer(pid):
                continue  # PID recycled to something else
            alive.add(pid)
        if alive:
            total += _kill_pids(alive, sig=signal.SIGKILL)
        try:
            path.unlink(missing_ok=True)
        except Exception:
            pass

    if total:
        print(
            f"  reaped {total} leftover stylo worker(s) from previous run "
            f"(Ctrl+C orphans)",
            flush=True,
        )
    return total


_POOL_CREATE_LOCK = threading.Lock()
_POOL_BROKEN = False


def _mark_metric_pool_broken() -> None:
    """Next ``_get_metric_pool`` will recreate after BrokenProcessPool / kill."""
    global _POOL_BROKEN
    _POOL_BROKEN = True


def _get_metric_pool(n_workers: int) -> Any:
    """Process pool (Linux fork + COW spaCy) or daemon thread pool."""
    global _METRIC_POOL, _METRIC_POOL_SIZE, _METRIC_POOL_KIND, _PROCESS_NLP
    global _POOL_BROKEN
    kind = "process" if _use_process_pool() else "thread"
    with _POOL_CREATE_LOCK:
        if (
            _METRIC_POOL is not None
            and _METRIC_POOL_SIZE == n_workers
            and _METRIC_POOL_KIND == kind
            and not _POOL_BROKEN
        ):
            return _METRIC_POOL

        was_broken = _POOL_BROKEN
        prev_kind = _METRIC_POOL_KIND
        if _METRIC_POOL is not None or _POOL_BROKEN:
            # Tear down previous / broken pool + kill process workers.
            shutdown_metric_pool(wait=False)
            _POOL_BROKEN = False

        # Re-forking after the parent has embed/HTTP/batcher threads is a common
        # SIGSEGV source (crash often surfaces later in re.sub / C extensions).
        # If a process pool already died mid-run, stay on threads.
        if (
            kind == "process"
            and was_broken
            and prev_kind == "process"
            and threading.active_count() > 2
        ):
            print(
                "warning: process stylo pool broke mid-run; "
                "falling back to threads (fork-after-threads is unsafe)",
                flush=True,
            )
            kind = "thread"
            _PROCESS_NLP = None

        if kind == "process":
            import multiprocessing as mp

            # Parent loads once; fork children share pages (COW) until they write.
            _ensure_process_nlp()
            try:
                ctx = mp.get_context("fork")
                _METRIC_POOL = ProcessPoolExecutor(
                    max_workers=n_workers,
                    mp_context=ctx,
                    initializer=_process_pool_initializer,
                )
                _METRIC_POOL_KIND = "process"
                _METRIC_POOL_SIZE = n_workers
                return _METRIC_POOL
            except Exception as e:
                print(
                    f"warning: process stylo pool failed ({type(e).__name__}: {e}); "
                    f"falling back to threads",
                    flush=True,
                )
                kind = "thread"
                _PROCESS_NLP = None  # threads use thread-local copies

        try:
            _METRIC_POOL = _DaemonThreadPoolExecutor(
                max_workers=n_workers, thread_name_prefix="obj1-metrics"
            )
        except Exception:
            _METRIC_POOL = ThreadPoolExecutor(
                max_workers=n_workers, thread_name_prefix="obj1-metrics"
            )
        _METRIC_POOL_KIND = "thread"
        _METRIC_POOL_SIZE = n_workers
        return _METRIC_POOL


def shutdown_metric_pool(*, wait: bool = False) -> None:
    """Drop the stylo pool and **kill process workers** (no Ctrl+C orphans).

    ``ProcessPoolExecutor.shutdown(wait=False)`` alone leaves children running
    (especially mid-spaCy). We SIGKILL registered worker PIDs so RSS returns.
    """
    global _METRIC_POOL, _METRIC_POOL_SIZE, _METRIC_POOL_KIND, _METRIC_BATCHER
    global _KNOWN_WORKER_PIDS
    batcher = _METRIC_BATCHER
    _METRIC_BATCHER = None
    if batcher is not None:
        try:
            batcher.close()
        except Exception:
            pass
    pool = _METRIC_POOL
    _METRIC_POOL = None
    kind = _METRIC_POOL_KIND
    _METRIC_POOL_SIZE = 0
    _METRIC_POOL_KIND = ""

    pids: set[int] = set(_KNOWN_WORKER_PIDS)
    if pool is not None:
        pids |= _snapshot_pool_worker_pids(pool)
    pids |= _read_worker_pidfile()

    # Signal workers before shutdown so they do not outlive os._exit.
    if kind == "process" or pids:
        _kill_pids(pids, sig=signal.SIGKILL)

    if pool is not None:
        try:
            pool.shutdown(wait=wait, cancel_futures=True)
        except TypeError:
            try:
                pool.shutdown(wait=wait)
            except Exception:
                pass
        except Exception:
            pass
        # Second pass: anything still listed on the executor.
        pids |= _snapshot_pool_worker_pids(pool)
        _kill_pids(pids, sig=signal.SIGKILL)

    # multiprocessing may still track children
    try:
        import multiprocessing as mp

        for child in mp.active_children():
            try:
                child.kill()
            except Exception:
                try:
                    child.terminate()
                except Exception:
                    pass
    except Exception:
        pass

    _KNOWN_WORKER_PIDS.clear()
    _clear_worker_pidfile()


def _doc_to_td_dict(doc) -> dict:
    out = td.extract_dict(doc)
    if isinstance(out, list):
        if not out:
            return {}
        return out[0] if isinstance(out[0], dict) else {}
    return out if isinstance(out, dict) else {}


def _td(STRING: str) -> dict:
    return _doc_to_td_dict(_get_nlp()(STRING))


def _vector_from_doc(doc, STRING: str) -> np.ndarray:
    """Feature vector from an already-parsed spaCy Doc + original text."""
    d = _doc_to_td_dict(doc)
    vals = _assemble(d, STRING)
    out = np.asarray(vals, dtype=np.float64)
    if out.shape != (_N_TOTAL,):
        raise RuntimeError(
            f"extract_vector length {out.shape[0]} != expected {_N_TOTAL}"
        )
    return out


def _safe(v, default=0.0) -> float:
    try:
        return float(v) if v is not None else default
    except Exception:
        return default


def _log10_perplexity(v, default=0.0) -> float:
    """log10 of textdescriptives perplexity; raw exp-scale values are unusable."""
    x = _safe(v, default=float("nan"))
    if not np.isfinite(x) or x <= 0.0:
        return default
    y = float(np.log10(x))
    if not np.isfinite(y):
        return default
    return float(np.clip(y, 0.0, 300.0))


# ---------------------------------------------------------------------------
# Fast pure-Python quality metrics (replaces textdescriptives/quality pipe)
# ---------------------------------------------------------------------------

def _chr_fraction_of_dup_spans(text: str, spans: list[str]) -> float:
    """Fraction of characters belonging to items that appear more than once."""
    if not text or not spans:
        return 0.0
    counts = Counter(spans)
    dup_chars = sum(len(s) * c for s, c in counts.items() if c > 1 and s)
    return float(min(1.0, dup_chars / max(len(text), 1)))


def _ngram_chr_len(gram: tuple[str, ...]) -> int:
    """Character length of words joined by single spaces."""
    if not gram:
        return 0
    return sum(len(w) for w in gram) + (len(gram) - 1)


def _word_ngram_chr_fraction(words: list[str], n: int) -> float:
    """
    Character-mass of word n-grams that repeat at least once, over total
    character mass of all word n-grams (textdescriptives-style duplicate_*gram).
    Uses tuple keys (no join) for speed on long docs.
    """
    if len(words) < n:
        return 0.0
    counts: Counter[tuple[str, ...]] = Counter()
    total = 0
    for i in range(len(words) - n + 1):
        g = tuple(words[i : i + n])
        counts[g] += 1
        total += _ngram_chr_len(g)
    if total <= 0:
        return 0.0
    dup = sum(_ngram_chr_len(g) * c for g, c in counts.items() if c > 1)
    return float(min(1.0, dup / total))


def _top_ngram_chr_fraction(words: list[str], n: int) -> float:
    """Character mass of the single most frequent word n-gram / all n-gram mass."""
    if len(words) < n:
        return 0.0
    counts: Counter[tuple[str, ...]] = Counter()
    total = 0
    for i in range(len(words) - n + 1):
        g = tuple(words[i : i + n])
        counts[g] += 1
        total += _ngram_chr_len(g)
    if total <= 0 or not counts:
        return 0.0
    top_g, top_c = counts.most_common(1)[0]
    return float(min(1.0, (_ngram_chr_len(top_g) * top_c) / total))


def _quality_vals(STRING: str) -> list[float]:
    """17 quality features — same names/order as former textdescriptives quality."""
    text = STRING or ""
    n_chars = max(len(text), 1)
    words = _WORD_RE.findall(text)
    n_words = max(len(words), 1)
    lower_words = [w.lower() for w in words]

    alpha_chars = sum(1 for c in text if c.isalpha())
    alpha_ratio = alpha_chars / n_chars
    mean_word_length = float(np.mean([len(w) for w in words])) if words else 0.0
    n_stop_words = float(sum(1 for w in lower_words if w in _STOP))
    proportion_ellipsis = float(len(_ELLIPSIS_RE.findall(text))) / n_words
    # lines starting with bullet-like markers
    lines = text.splitlines()
    if lines:
        bullets = sum(
            1
            for ln in lines
            if ln.lstrip().startswith(("-", "*", "•", "·"))
            or (ln.lstrip()[:2].isdigit() and ln.lstrip()[1:2] in ".)")
        )
        proportion_bullet_points = bullets / len(lines)
    else:
        proportion_bullet_points = 0.0

    # non-empty lines / paragraphs
    nonempty_lines = [ln for ln in lines if ln.strip()]
    paragraphs = [p for p in re.split(r"\n\s*\n", text) if p.strip()]
    duplicate_line_chr_fraction = _chr_fraction_of_dup_spans(text, nonempty_lines)
    duplicate_paragraph_chr_fraction = _chr_fraction_of_dup_spans(text, paragraphs)

    dup_ngrams = [_word_ngram_chr_fraction(lower_words, n) for n in range(5, 11)]
    top_ngrams = [_top_ngram_chr_fraction(lower_words, n) for n in range(2, 5)]

    # symbol '#' to word ratio (textdescriptives symbol_to_word_ratio_#)
    symbol_hash_to_word_ratio = text.count("#") / n_words

    return [
        _safe(alpha_ratio),
        _safe(mean_word_length),
        _safe(n_stop_words),
        _safe(proportion_ellipsis),
        _safe(proportion_bullet_points),
        _safe(duplicate_line_chr_fraction),
        _safe(duplicate_paragraph_chr_fraction),
        *[_safe(x) for x in dup_ngrams],
        *[_safe(x) for x in top_ngrams],
        _safe(symbol_hash_to_word_ratio),
    ]


# ---------------------------------------------------------------------------
# Fixed feature names (order must match the vector)
# ---------------------------------------------------------------------------
FEATURE_NAMES: list[str] = [
    # TextDescriptives – descriptive
    "token_length_mean",
    "token_length_median",
    "token_length_std",
    "sentence_length_mean",
    "sentence_length_median",
    "sentence_length_std",
    "syllables_per_token_mean",
    "syllables_per_token_median",
    "syllables_per_token_std",
    "proportion_unique_tokens",
    "n_tokens",
    "n_unique_tokens",
    "n_characters",
    "n_sentences",
    # TextDescriptives – readability
    "flesch_reading_ease",
    "flesch_kincaid_grade",
    "smog",
    "gunning_fog",
    "automated_readability_index",
    "coleman_liau_index",
    "lix",
    "rix",
    # Quality (fast pure-Python; same names as former textdescriptives quality)
    "alpha_ratio",
    "mean_word_length",
    "n_stop_words",
    "proportion_ellipsis",
    "proportion_bullet_points",
    "duplicate_line_chr_fraction",
    "duplicate_paragraph_chr_fraction",
    "duplicate_5gram_chr_fraction",
    "duplicate_6gram_chr_fraction",
    "duplicate_7gram_chr_fraction",
    "duplicate_8gram_chr_fraction",
    "duplicate_9gram_chr_fraction",
    "duplicate_10gram_chr_fraction",
    "top_2gram_chr_fraction",
    "top_3gram_chr_fraction",
    "top_4gram_chr_fraction",
    "symbol_hash_to_word_ratio",
    # TextDescriptives – information theory (perplexity → log10)
    "entropy",
    "log10_perplexity",
    # TextDescriptives – POS proportions
    "pos_prop_NOUN",
    "pos_prop_VERB",
    "pos_prop_ADJ",
    "pos_prop_ADV",
    "pos_prop_PRON",
    "pos_prop_DET",
    "pos_prop_ADP",
    "pos_prop_AUX",
    "pos_prop_CCONJ",
    "pos_prop_SCONJ",
    "pos_prop_PART",
    "pos_prop_INTJ",
    "pos_prop_PROPN",
    "pos_prop_NUM",
    "pos_prop_PUNCT",
    # lexicalrichness
    "mattr",
    "mtld",
    "hdd",
    "maas",
    "herdan",
    "summer",
    "dugast",
    "yulek",
    "yulei",
    "herdanvm",
    "simpsond",
    "vocd",
    # textstat
    "dale_chall_readability_score",
    "linsear_write_formula",
    "spache_readability",
    "mcalpine_eflaw",
    "reading_time",
    "syllable_count",
    "lexicon_count",
    "sentence_count",
    "char_count",
    "letter_count",
    "polysyllabcount",
    "monosyllabcount",
    "difficult_words",
    "text_standard_float",
    # NRCLex
    "nrc_fear",
    "nrc_anger",
    "nrc_anticipation",
    "nrc_trust",
    "nrc_surprise",
    "nrc_sadness",
    "nrc_disgust",
    "nrc_joy",
    "nrc_positive",
    "nrc_negative",
]

# spaCy/TD dict slice: descriptive(14) + readability(8) — quality is pure Python
_N_DESC_READ = 22
_N_QUALITY = 17
_N_INFO_POS = 17  # entropy+perp(2) + POS(15)
_N_SPACY_CORE = _N_DESC_READ + _N_INFO_POS  # values taken from spaCy dict
_N_TOTAL = len(FEATURE_NAMES)


def _desc_read_vals(d: dict) -> list[float]:
    return [
        _safe(d.get("token_length_mean")),
        _safe(d.get("token_length_median")),
        _safe(d.get("token_length_std")),
        _safe(d.get("sentence_length_mean")),
        _safe(d.get("sentence_length_median")),
        _safe(d.get("sentence_length_std")),
        _safe(d.get("syllables_per_token_mean")),
        _safe(d.get("syllables_per_token_median")),
        _safe(d.get("syllables_per_token_std")),
        _safe(d.get("proportion_unique_tokens")),
        _safe(d.get("n_tokens")),
        _safe(d.get("n_unique_tokens")),
        _safe(d.get("n_characters")),
        _safe(d.get("n_sentences")),
        _safe(d.get("flesch_reading_ease")),
        _safe(d.get("flesch_kincaid_grade")),
        _safe(d.get("smog")),
        _safe(d.get("gunning_fog")),
        _safe(d.get("automated_readability_index")),
        _safe(d.get("coleman_liau_index")),
        _safe(d.get("lix")),
        _safe(d.get("rix")),
    ]


def _info_pos_vals(d: dict) -> list[float]:
    return [
        _safe(d.get("entropy")),
        _log10_perplexity(d.get("perplexity")),
        _safe(d.get("pos_prop_NOUN")),
        _safe(d.get("pos_prop_VERB")),
        _safe(d.get("pos_prop_ADJ")),
        _safe(d.get("pos_prop_ADV")),
        _safe(d.get("pos_prop_PRON")),
        _safe(d.get("pos_prop_DET")),
        _safe(d.get("pos_prop_ADP")),
        _safe(d.get("pos_prop_AUX")),
        _safe(d.get("pos_prop_CCONJ")),
        _safe(d.get("pos_prop_SCONJ")),
        _safe(d.get("pos_prop_PART")),
        _safe(d.get("pos_prop_INTJ")),
        _safe(d.get("pos_prop_PROPN")),
        _safe(d.get("pos_prop_NUM")),
        _safe(d.get("pos_prop_PUNCT")),
    ]


def _spacy_core_vals(d: dict) -> list[float]:
    """desc+read + info+POS (quality filled separately)."""
    return _desc_read_vals(d) + _info_pos_vals(d)


def _lex_textstat_nrc_vals(STRING: str) -> list[float]:
    """lexicalrichness + textstat + NRCLex."""
    vals: list[float] = []

    try:
        lex = LexicalRichness(STRING)
        vals.extend(
            [
                _safe(lex.mattr(window_size=50)),
                _safe(lex.mtld(threshold=0.72)),
                _safe(lex.hdd(draws=42)),
                _safe(lex.Maas),
                _safe(lex.Herdan),
                _safe(lex.Summer),
                _safe(lex.Dugast),
                _safe(lex.yulek),
                _safe(lex.yulei),
                _safe(lex.herdanvm),
                _safe(lex.simpsond),
                _safe(lex.vocd()),
            ]
        )
    except Exception:
        vals.extend([0.0] * 12)

    try:
        vals.extend(
            [
                _safe(textstat.dale_chall_readability_score(STRING)),
                _safe(textstat.linsear_write_formula(STRING)),
                _safe(textstat.spache_readability(STRING)),
                _safe(textstat.mcalpine_eflaw(STRING)),
                _safe(textstat.reading_time(STRING)),
                _safe(textstat.syllable_count(STRING)),
                _safe(textstat.lexicon_count(STRING)),
                _safe(textstat.sentence_count(STRING)),
                _safe(textstat.char_count(STRING)),
                _safe(textstat.letter_count(STRING)),
                _safe(textstat.polysyllabcount(STRING)),
                _safe(textstat.monosyllabcount(STRING)),
                _safe(textstat.difficult_words(STRING)),
                _safe(textstat.text_standard(STRING, float_output=True)),
            ]
        )
    except Exception:
        vals.extend([0.0] * 14)

    try:
        tokens = [
            w.strip(".,;:!?\"'()[]{}").lower() for w in STRING.split() if w.strip()
        ]
        emotion = _get_nrc()
        if tokens:
            emotion.load_token_list(tokens)
        else:
            emotion.load_raw_text(STRING)
        raw = getattr(emotion, "raw_emotion_scores", None) or {}
        word_count = max(len(tokens) if tokens else len(STRING.split()), 1)
        vals.extend(
            [
                _safe(raw.get("fear", 0) / word_count),
                _safe(raw.get("anger", 0) / word_count),
                _safe(raw.get("anticipation", 0) / word_count),
                _safe(raw.get("trust", 0) / word_count),
                _safe(raw.get("surprise", 0) / word_count),
                _safe(raw.get("sadness", 0) / word_count),
                _safe(raw.get("disgust", 0) / word_count),
                _safe(raw.get("joy", 0) / word_count),
                _safe(raw.get("positive", 0) / word_count),
                _safe(raw.get("negative", 0) / word_count),
            ]
        )
    except Exception:
        vals.extend([0.0] * 10)

    return vals


def _non_spacy_block(STRING: str) -> list[float]:
    """Quality + lexicalrichness + textstat + NRC (everything not needing spaCy)."""
    return _quality_vals(STRING) + _lex_textstat_nrc_vals(STRING)


def _assemble(d: dict, STRING: str, non_spacy: list[float] | None = None) -> list[float]:
    """desc/read from spaCy, quality..nrc from non_spacy block (or compute)."""
    rest = non_spacy if non_spacy is not None else _non_spacy_block(STRING)
    # rest = quality(17) + lex(12) + textstat(14) + nrc(10) = 53
    q = rest[:_N_QUALITY]
    tail = rest[_N_QUALITY:]
    vals = _desc_read_vals(d) + q + _info_pos_vals(d) + tail
    return vals


def extract_vector(STRING: str) -> np.ndarray:
    """Return a fixed-order float vector of all metrics (single text)."""
    text = _sanitize_metric_text(STRING)
    doc = _get_nlp()(text)
    return _vector_from_doc(doc, text)


def _extract_batch(texts: list[str]) -> list[np.ndarray]:
    """Run many windows on one worker via nlp.pipe (better CPU per GiB)."""
    cleaned = [_sanitize_metric_text(t) for t in texts]
    if not cleaned:
        return []
    if len(cleaned) == 1:
        return [extract_vector(cleaned[0])]
    nlp = _get_nlp()
    # nrc is thread-local; one worker thread/process at a time uses it.
    out: list[np.ndarray] = []
    for doc, text in zip(
        nlp.pipe(cleaned, batch_size=min(64, len(cleaned))), cleaned
    ):
        out.append(_vector_from_doc(doc, text))
    return out


# Top-level alias for ProcessPool pickling.
def _pool_extract_batch(texts: list[str]) -> list[list[float]]:
    """Process-pool entry: return plain lists (smaller/safer than ndarray pickle)."""
    vecs = _extract_batch(texts)
    return [v.tolist() for v in vecs]


def _pool_extract_one(text: str) -> list[float]:
    return extract_vector(text).tolist()


class _MetricBatcher:
    """Coalesce live single-window submits into nlp.pipe batches.

    Small wait (default 40ms) / max batch (24) so streaming stylo uses pipe
    throughput instead of one IPC + one nlp() per window.
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        self._cv = threading.Condition(self._lock)
        self._pending: list[tuple[str, Future]] = []
        self._closed = False
        self._thread = threading.Thread(
            target=self._loop, name="stylo-batcher", daemon=True
        )
        self._thread.start()

    def submit(self, text: str) -> Future:
        fut: Future = Future()
        with self._cv:
            if self._closed:
                fut.set_exception(RuntimeError("stylo batcher closed"))
                return fut
            self._pending.append((text, fut))
            if len(self._pending) >= max(1, _BATCH_MAX):
                self._cv.notify()
        return fut

    def close(self) -> None:
        with self._cv:
            self._closed = True
            self._cv.notify_all()

    def _loop(self) -> None:
        while True:
            with self._cv:
                if self._closed and not self._pending:
                    return
                if not self._pending:
                    self._cv.wait(timeout=max(0.01, _BATCH_WAIT_S))
                if not self._pending:
                    if self._closed:
                        return
                    continue
                take = min(len(self._pending), max(1, _BATCH_MAX))
                batch = self._pending[:take]
                del self._pending[:take]
            texts = [t for t, _f in batch]
            futs = [f for _t, f in batch]
            try:
                workers = _metric_workers()
                try:
                    pool = _get_metric_pool(workers)
                except Exception:
                    _mark_metric_pool_broken()
                    raise

                def _apply_rows(
                    batch_futs: list[Future], rows: list
                ) -> None:
                    if len(rows) != len(batch_futs):
                        raise RuntimeError(
                            f"stylo batch length mismatch: "
                            f"got {len(rows)} rows for {len(batch_futs)} futures"
                        )
                    for fut, row in zip(batch_futs, rows):
                        if not fut.cancelled():
                            fut.set_result(np.asarray(row, dtype=np.float64))

                # One pool task runs the whole batch on one worker (pipe).
                # Fan out large flushes across workers for multi-core.
                if len(texts) <= 4 or workers <= 1:
                    rows = pool.submit(_pool_extract_batch, texts).result()
                    _apply_rows(futs, rows)
                else:
                    n_chunks = min(workers, len(texts))
                    chunk_size = (len(texts) + n_chunks - 1) // n_chunks
                    chunks: list[list[str]] = []
                    fut_chunks: list[list[Future]] = []
                    for i in range(0, len(texts), chunk_size):
                        chunks.append(texts[i : i + chunk_size])
                        fut_chunks.append(futs[i : i + chunk_size])
                    pfuts = [pool.submit(_pool_extract_batch, c) for c in chunks]
                    for pf, fc in zip(pfuts, fut_chunks):
                        rows = pf.result()
                        _apply_rows(fc, rows)
            except Exception as e:
                # BrokenProcessPool / killed workers → recreate next call.
                ename = type(e).__name__
                if "Broken" in ename or "Terminated" in ename:
                    _mark_metric_pool_broken()
                for fut in futs:
                    if not fut.cancelled() and not fut.done():
                        fut.set_exception(e)


def _get_batcher() -> _MetricBatcher:
    global _METRIC_BATCHER
    with _POOL_CREATE_LOCK:
        if _METRIC_BATCHER is None:
            _METRIC_BATCHER = _MetricBatcher()
        return _METRIC_BATCHER


def submit_extract_vector(text: str, *, n_texts_hint: int = 5):
    """
    Queue one window metric extract (batched + process/thread pool).

    Used live as each 256-tok window completes. Returns
    concurrent.futures.Future[np.ndarray].
    """
    global _METRIC_WARNED_SUBMIT
    del n_texts_hint
    # Ensure pool exists (and process model preloaded) before batcher runs.
    try:
        _get_metric_pool(_metric_workers())
        return _get_batcher().submit(
            text if isinstance(text, str) else str(text)
        )
    except Exception as e:
        if not _METRIC_WARNED_SUBMIT:
            print(
                f"warning: stylo submit failed ({type(e).__name__}: {e})",
                flush=True,
            )
            _METRIC_WARNED_SUBMIT = True
        raise


def _preload_worker_slot(i: int) -> int:
    """Warm one pool slot (process or thread)."""
    _ = i
    extract_vector("The night was cold and the river moved slowly under the moon.")
    return 0


def warmup_metric_workers(n_texts: int = 5) -> int:
    """
    Create the metric pool and warm each worker.

    Process mode: parent loads spaCy once (fork COW), then each child is
    touched with a smoke extract. Thread mode: each thread loads its own copy.
    """
    del n_texts
    workers = _metric_workers()
    pool = _get_metric_pool(workers)
    kind = _METRIC_POOL_KIND or "thread"
    info = metric_pool_info()
    print(
        f"  stylo pool: kind={kind} workers={workers} "
        f"stylo≈{info['stylo_ram_est_gib']}g "
        f"script≈{info['script_ram_est_gib']}g/"
        f"≤{info['script_ram_budget_gib']}g "
        f"(reserve {info['non_stylo_reserve_gib']}g non-stylo) "
        f"batch_max={_BATCH_MAX} batch_wait_ms={_BATCH_WAIT_S * 1000:.0f}",
        flush=True,
    )
    # Touch batcher so the coalescing thread is alive before the first trial.
    _get_batcher()
    futs = [pool.submit(_preload_worker_slot, i) for i in range(workers)]
    # Process workers exist after first submit — track PIDs for Ctrl+C cleanup.
    register_metric_pool_workers(pool)
    for i, fut in enumerate(futs):
        try:
            fut.result(timeout=180.0)
            print(f"  stylo worker {i + 1}/{workers} ready", flush=True)
        except Exception as e:
            print(
                f"warning: stylo worker warm {i + 1}/{workers} failed "
                f"({type(e).__name__}: {e})",
                flush=True,
            )
        register_metric_pool_workers(pool)
    return workers


def metric_pool_info(n_texts: int = 5) -> dict[str, int | str | float | None]:
    """Diagnostics for logs: workers vs host CPUs / RAM / pool kind."""
    del n_texts
    cpu = os.cpu_count() or 2
    ram = _total_ram_gb()
    workers = _metric_workers()
    kind = "process" if _use_process_pool() else "thread"
    if _METRIC_POOL_KIND:
        kind = _METRIC_POOL_KIND
    stylo_budget = _stylo_ram_budget_gib()
    script_budget = _script_ram_budget_gib()
    reserve = _non_stylo_reserve_gib()
    if kind == "process" or _use_process_pool():
        stylo_est = _STYLO_PARENT_GIB + workers * _STYLO_PER_PROCESS_GIB
    else:
        stylo_est = workers * _STYLO_PER_THREAD_GIB
    return {
        "cpu_count": cpu,
        "ram_gb": None if ram is None else round(ram, 1),
        "workers": workers,
        "pool_kind": kind,
        "script_ram_budget_gib": script_budget,
        "non_stylo_reserve_gib": reserve,
        "stylo_ram_budget_gib": stylo_budget,
        "stylo_ram_est_gib": round(stylo_est, 2),
        "script_ram_est_gib": round(stylo_est + reserve, 2),
        "override": (os.getenv("METRIC_WORKERS") or "").strip() or "auto",
    }


def extract_vectors(texts: Iterable[str]) -> np.ndarray:
    """
    Batch extract for multiple texts (shape: n_texts × n_features).

    Splits work across pool workers; each worker uses ``nlp.pipe`` on its chunk.
    """
    texts_list = [t if isinstance(t, str) else str(t) for t in texts]
    n = len(texts_list)
    if n == 0:
        return np.zeros((0, _N_TOTAL), dtype=np.float64)
    if n == 1:
        return extract_vector(texts_list[0]).reshape(1, -1)

    workers = _metric_workers()
    pool = _get_metric_pool(workers)
    n_chunks = min(workers, n)
    chunk_size = (n + n_chunks - 1) // n_chunks
    chunks = [
        texts_list[i : i + chunk_size] for i in range(0, n, chunk_size)
    ]
    parts = list(pool.map(_pool_extract_batch, chunks))
    vectors = [np.asarray(row, dtype=np.float64) for part in parts for row in part]
    return np.stack(vectors, axis=0)


if __name__ == "__main__":
    text = (
        "The old house stood silent under the moon. "
        "She whispered, 'I am afraid,' and the shadows seemed to listen. "
        "Fear and hope fought inside her heart."
    )
    vec = extract_vector(text)
    print(f"Vector shape: {vec.shape}")
    print(f"Features: {len(FEATURE_NAMES)}")
    for name, value in zip(FEATURE_NAMES, vec):
        print(f"{name:45s} = {value:10.4f}")
