"""On-demand llama-server lead for cross-machine shard plans.

When the allocator shards a model across the GPU pool, the LEAD must load it
with llama.cpp's RPC backend (--rpc host:port,... + --tensor-split). The
0.3.x python binding can't do that (Llama() lost its rpc_servers param), so
the lead runs a managed ``llama-server`` subprocess instead and the existing
HTTP runner (LlamaCppRunner) talks to it.

One process per (model_key, rpc_servers, tensor_split), cached for the life
of this process — same lifecycle semantics as the in-process Llama singleton:
a changed shard plan takes effect on the next fresh load, not retroactively.

``ensure_shard_server`` returns the base_url on success or None on any
failure, so callers can always fall back to local loading.
"""
from __future__ import annotations

import os
import socket
import subprocess
import threading
import time
from typing import Optional

import httpx

from .imports import *

_SERVERS: dict = {}            # (model_key, rpc, ts) -> {"proc": Popen, "base_url": str}
_LOCK = threading.Lock()

_PORT_BASE = int(os.environ.get("HUGPY_SHARD_PORT_BASE", "8790"))
_HEALTH_TIMEOUT_S = float(os.environ.get("HUGPY_SHARD_HEALTH_TIMEOUT", "180"))


def _server_bin() -> Optional[str]:
    """llama-server binary, resolved per-OS (env -> `hugpy install-engine` -> PATH)."""
    from .....engine.resolve import server_bin

    return server_bin()


def _free_port(start: int) -> int:
    for port in range(start, start + 200):
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            if s.connect_ex(("127.0.0.1", port)) != 0:
                return port
    raise OSError(f"no free port in [{start}, {start + 200})")


def _wait_healthy(base_url: str, proc: subprocess.Popen) -> bool:
    """Poll /health until ok; loading a big model over RPC can take a while."""
    deadline = time.monotonic() + _HEALTH_TIMEOUT_S
    while time.monotonic() < deadline:
        if proc.poll() is not None:
            return False                      # server died during load
        try:
            with httpx.Client(timeout=2.0) as client:
                if client.get(f"{base_url}/health").status_code == 200:
                    return True
        except Exception:
            pass
        time.sleep(1.5)
    return False


def _prune_dead() -> None:
    for key in [k for k, v in _SERVERS.items() if v["proc"].poll() is not None]:
        logger.warning("shard server for %s exited; dropping from cache", key[0])
        _SERVERS.pop(key, None)


def ensure_shard_server(
    model_key: str,
    rpc_servers: str,
    tensor_split=None,
) -> Optional[str]:
    """Get-or-spawn the llama-server lead for this shard plan.

    Returns its base_url ("http://127.0.0.1:<port>") once /health passes,
    or None (binary missing / no GGUF / failed to come up) so the caller can
    fall back to a local load.
    """
    ts_key = tuple(tensor_split) if tensor_split else ()
    cache_key = (model_key, rpc_servers, ts_key)

    with _LOCK:
        _prune_dead()
        hit = _SERVERS.get(cache_key)
        if hit:
            return hit["base_url"]

        binary = _server_bin()
        if not binary:
            logger.warning("ensure_shard_server: no llama-server binary "
                           "(set LLAMA_SERVER_BIN); cannot lead a shard")
            return None

        try:
            cfg = get_model_config(model_key)
            model_dir = ensure_model(model_key)
            model_path = get_gguf_file(model_dir, cfg)
        except Exception as exc:
            logger.warning("ensure_shard_server: model resolve failed for %s: %s",
                           model_key, exc)
            return None
        if not model_path:
            logger.warning("ensure_shard_server: no GGUF for %s — only llama_cpp "
                           "models can shard via RPC", model_key)
            return None

        port = _free_port(_PORT_BASE)
        base_url = f"http://127.0.0.1:{port}"
        cmd = [
            binary,
            "-m", os.fspath(model_path),
            "--host", "127.0.0.1",
            "--port", str(port),
            "--rpc", rpc_servers,
            "-ngl", "999",                    # all layers onto the pooled backends
            "-c", str(DEFAULT_N_CTX),
        ]
        if ts_key:
            cmd += ["--tensor-split", ",".join(str(x) for x in ts_key)]

        # Vision GGUF: load the multimodal projector beside the model. No-op for
        # text models (no projector present).
        from .....imports.src.utils import find_mmproj
        mmproj = find_mmproj(model_path)
        if mmproj:
            cmd += ["--mmproj", os.fspath(mmproj)]

        logger.info("shard lead: %s", " ".join(cmd))
        try:
            from ....._platform.procutil import popen_detached
            proc = popen_detached(                 # detaches so it survives reloads (cross-OS)
                cmd,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except OSError as exc:
            logger.error("shard lead failed to start: %s", exc)
            return None

    # Health-wait outside the lock so a slow load doesn't serialize other models.
    if not _wait_healthy(base_url, proc):
        logger.error("shard lead for %s never became healthy; killing pid %s",
                     model_key, proc.pid)
        proc.kill()
        return None

    with _LOCK:
        _SERVERS[cache_key] = {"proc": proc, "base_url": base_url}
    logger.info("shard lead healthy for %s at %s (rpc=%s split=%s)",
                model_key, base_url, rpc_servers, ts_key or "auto")
    return base_url
