"""sshj.probe — reachability probing (bounded pool).

B2 step 3d: moved verbatim from sshj_cli.py (single-file era) — no
behavior change. Owns PROBE_TIMEOUT (main's direct probe imports
it back). Edge: tailnet.ssh_host (probe args resolve via it).
"""

import os
import subprocess
import threading
import time
from concurrent.futures import ThreadPoolExecutor

from .tailnet import ssh_host


PROBE_TIMEOUT = 5          # seconds, max per reachability probe

PROBE_POLL = 0.2           # probe wait granularity: a stopped pool must be able to KILL its in-flight ssh (below)

PROBE_POOL_DEFAULT = 3


def _probe_args(h):
    if h.get("src"):
        tgt = ssh_host(h)
        args = ["ssh"]
        if h.get("port") not in ("", "22"):
            args += ["-p", h["port"]]
        args += ["-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=no",
                 "-o", "LogLevel=ERROR", "--", tgt, "true"]
        return args
    args = ["ssh", "-o", "BatchMode=yes", "-o", "StrictHostKeyChecking=no",
            "-o", "LogLevel=ERROR", "--", h["alias"], "true"]
    return args


def probe(h, timeout_s, stop_event=None):
    """Returns (up: bool, latency_ms: float|None) via a timed ssh round-trip.

    `stop_event` (optional, a threading.Event): when set, the in-flight ssh
    is KILLED and the probe returns down — a quit must be instant, so the
    pool's workers can be interrupted (2026-09-21: Esc blocked on
    concurrent.futures' atexit hook waiting for in-flight probes)."""
    try:
        t = max(1, min(int(float(timeout_s)), PROBE_TIMEOUT)) if str(timeout_s) not in ("", "none") else PROBE_TIMEOUT
    except ValueError:
        t = PROBE_TIMEOUT
    args = _probe_args(h)
    args[1:1] = ["-o", f"ConnectTimeout={t}"]
    try:
        t0 = time.time()
        p = subprocess.Popen(args, stdout=subprocess.DEVNULL,
                             stderr=subprocess.DEVNULL)
        end = time.time() + t + 2
        while p.poll() is None:
            if stop_event is not None and stop_event.is_set():
                p.kill()
                p.wait()
                return False, None
            if time.time() >= end:
                p.kill()
                p.wait()
                return False, None
            time.sleep(PROBE_POLL)
        up = p.returncode == 0
        ms = (time.time() - t0) * 1000.0
        return up, (ms if up else None)
    except Exception:
        return False, None


class Checker(threading.Thread):
    """Probes hosts in a bounded worker pool (B14), in priority order,
    until the screen dies. One probe at a time (the old behavior) is
    SSHJ_PROBE_CONCURRENCY=1; default is 3 — a pool of N ssh round-trips
    that each cost <5s used to serialise to N×5s, so the pool probes
    visible rows first and the rest fill in as workers free up."""

    def __init__(self, hosts):
        super().__init__(daemon=True)
        self.by_alias = {h["alias"]: h for h in hosts}
        self.status = {}        # alias -> "check" | "up" | "down"
        self.latency = {}       # alias -> ms (measured ssh round-trip)
        self.lock = threading.Lock()
        self.order = [h["alias"] for h in hosts]
        self.kick = threading.Event()
        self.done = False
        # set on stop: in-flight probes KILL their ssh and return, so the
        # pool's (non-daemon) workers finish within one poll tick. That is
        # what makes a quit instant — concurrent.futures' atexit hook
        # otherwise waits for EVERY in-flight worker future (2026-09-21).
        self.stop_ev = threading.Event()
        self.pool_size = self._pool_size()
        self.pool = ThreadPoolExecutor(
            max_workers=self.pool_size, thread_name_prefix="probe")

    @staticmethod
    def _pool_size():
        try:
            n = int(os.environ.get("SSHJ_PROBE_CONCURRENCY", ""))
            return max(1, min(16, n))
        except ValueError:
            return PROBE_POOL_DEFAULT

    def add_host(self, h):
        with self.lock:
            if h["alias"] in self.by_alias:
                return
            self.by_alias[h["alias"]] = h
            self.order.append(h["alias"])
        self.kick.set()

    def reprioritize(self, ordered_aliases):
        with self.lock:
            self.order = list(dict.fromkeys(ordered_aliases))
        self.kick.set()

    def stop(self):
        self.done = True
        self.stop_ev.set()
        self.kick.set()

    def run_once(self):
        """Probe every host not yet measured, concurrently under the pool.
        Returns seconds the pass took (0.0 when there was nothing to do).
        The thread loop calls this; tests call it directly (deterministic)."""
        with self.lock:
            todo = [a for a in self.order if a not in self.status]
            for a in todo:
                self.status[a] = "check"
            hosts = [self.by_alias[a] for a in todo]
        if not hosts:
            return 0.0
        t0 = time.time()
        for h in hosts:
            self.pool.submit(self._probe_one, h)
        time.sleep(0)  # yield so workers start before the loop re-checks
        # Wait for THIS pass to settle: every todo host has a status.
        while not self.done:
            with self.lock:
                missing = [a for a in todo if self.status.get(a) not in
                           ("up", "down")]
            if not missing:
                break
            time.sleep(0.05)
        return time.time() - t0

    def _probe_one(self, h):
        alias = h["alias"]
        up, ms = probe(h, h.get("timeout", ""), stop_event=self.stop_ev)
        with self.lock:
            if self.done and alias not in self.by_alias:
                return
            self.status[alias] = "up" if up else "down"
            if ms is not None:
                self.latency[alias] = ms
        self.kick.set()

    def run(self):
        while not self.done:
            self.run_once()
            if self.done:
                break
            self.kick.wait(1.0)
            self.kick.clear()


