"""Serve a trained model on vLLM or SGLang — an OpenAI-compatible server, launched and watched.

`launch_openai_server` is the vLLM / SGLang twin of `ollama_push.push_to_ollama` and
`bankml_push.push_to_bankml`: resolve the launcher → build the exact argv → start it detached
(own session, log file and pid file under the run dir) → poll until the tag answers on
`/v1/models` (and `/health` for vLLM) → optionally swap mindX's fallback model to it.
`stop_openai_server` reads the pid file, sends SIGTERM to the server's process group and
verifies it is gone.

Two ways to serve a LoRA, chosen per call:

- **native (default — no merge):** the base model plus the adapter, as the server loads it.
  vLLM: `--enable-lora --lora-modules <tag>=<adapter> --max-lora-rank R`; the adapter is listed
  on `/v1/models` under `<tag>`. SGLang: `--enable-lora --lora-paths <tag>=<adapter>
  --max-lora-rank R`; a client selects it as `model: "<served-name>:<tag>"`.
- **merged (`merge=True`):** `merge_lora_adapter` folds the adapter into the base weights
  (`uv sync --extra ml`) and the merged directory is served under `<tag>`.

On a host without a GPU the CPU backends are used: vLLM's CPU build reads
`VLLM_CPU_KVCACHE_SPACE` (GiB; set here unless the environment already sets it) and wants
`--dtype bfloat16`; SGLang takes `--device cpu`. A config that needs a GPU (an FP8 / MXFP4 /
GPTQ checkpoint, or tensor parallel > 1) is refused on a GPU-less host, with the reason.

Every flag below is from the upstream server-argument docs (vLLM `vllm serve` CLI reference and
LoRA / CPU-installation pages; SGLang "Server Arguments" and LoRA pages). Anything else a site
needs goes through `extra_args`, verbatim — this module does not guess at flags.

Clean-room: vLLM and SGLang are reached only as a subprocess and over HTTP. Neither is imported
here (SGLang's presence is checked with `importlib.util.find_spec`, which does not import it).

Every public function returns a result; nothing here raises past `ServerLaunchResult` /
`ServerStopResult`.
"""

from __future__ import annotations

import importlib.util
import json
import os
import shutil
import signal
import subprocess
import sys
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Literal

import httpx
from pydantic import BaseModel, ConfigDict, Field

Flavour = Literal["vllm", "sglang"]
Dtype = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]

VLLM_INSTALL_HINT = (
    "install vLLM with `uv sync --extra serve` (GPU / ROCm), or CPU-only per "
    "https://docs.vllm.ai/en/latest/getting_started/installation/cpu.html "
    "(e.g. `uv pip install vllm --torch-backend cpu`); or pass --server-bin"
)
SGLANG_INSTALL_HINT = (
    "SGLang is not a mindXtrain extra; install it into the interpreter that will run it "
    "(`uv pip install sglang`, see https://docs.sglang.io) or pass --server-bin <python>"
)

# vLLM `--max-lora-rank` choices (vllm serve CLI reference); SGLang takes any int (default 16).
VLLM_LORA_RANKS: tuple[int, ...] = (1, 8, 16, 32, 64, 128, 256, 320, 512)

# `quantize.scheme` → vLLM `--quantization` (as vllm_launcher.build_vllm_command maps it).
VLLM_QUANTIZATION: dict[str, str] = {"quark_fp8": "fp8", "quark_mxfp4": "mxfp4",
                                     "gptq_rocm": "gptq"}

LaunchStatus = Literal[
    "ready", "dry_run", "refused", "missing", "no_gpu", "merge_failed", "already_running",
    "exited", "timeout", "error",
]
StopStatus = Literal["stopped", "not_running", "still_running", "error"]


class OpenAIServerSpec(BaseModel):
    """What to serve and how. `build_argv` / `build_env` turn it into the exact launch."""

    model_config = ConfigDict(extra="forbid", frozen=True)

    flavour: Flavour
    model: str = Field(min_length=1, description="Model dir or HF repo id (the base, for LoRA).")
    tag: str = Field(min_length=1, description="The name clients use (served name or LoRA name).")
    lora_path: str | None = Field(default=None, description="Adapter dir served natively.")
    lora_rank: int | None = Field(default=None, ge=1, description="Adapter rank (`r`).")
    host: str = "127.0.0.1"
    port: int = Field(default=8000, ge=1024, le=65535)
    dtype: Dtype = "auto"
    max_model_len: int = Field(default=8192, ge=512)
    tensor_parallel: int = Field(default=1, ge=1)
    quantization: str | None = None
    cpu: bool = False
    cpu_kvcache_gib: int = Field(default=4, ge=1)
    extra_args: tuple[str, ...] = ()

    @property
    def served_name(self) -> str:
        """The name the base (or merged) model is served under."""
        return f"{self.tag}-base" if self.lora_path else self.tag

    @property
    def client_model(self) -> str:
        """What an OpenAI client puts in `model` to reach the tag."""
        if self.lora_path and self.flavour == "sglang":
            return f"{self.served_name}:{self.tag}"
        return self.tag

    @property
    def base_url(self) -> str:
        return f"http://{self.host}:{self.port}"


def _vllm_rank(r: int) -> int:
    return next((c for c in VLLM_LORA_RANKS if c >= r), VLLM_LORA_RANKS[-1])


def default_launcher(flavour: Flavour) -> list[str]:
    """The launcher prefix assumed when none is resolved (used by dry runs)."""
    if flavour == "vllm":
        return ["vllm"]
    return [sys.executable, "-m", "sglang.launch_server"]


def build_argv(spec: OpenAIServerSpec, launcher: list[str] | None = None) -> list[str]:
    """The exact argv for `spec`. `launcher` is `["vllm"]` or `[python, "-m", "sglang.launch_server"]`."""
    prefix = list(launcher or default_launcher(spec.flavour))
    dtype = "bfloat16" if (spec.cpu and spec.dtype == "auto") else spec.dtype
    if spec.flavour == "vllm":
        argv = [*prefix, "serve", spec.model,
                "--served-model-name", spec.served_name,
                "--host", spec.host, "--port", str(spec.port),
                "--dtype", dtype,
                "--max-model-len", str(spec.max_model_len),
                "--tensor-parallel-size", str(spec.tensor_parallel)]
        if spec.quantization:
            argv += ["--quantization", spec.quantization]
        if spec.lora_path:
            argv += ["--enable-lora", "--lora-modules", f"{spec.tag}={spec.lora_path}",
                     "--max-lora-rank", str(_vllm_rank(spec.lora_rank or 16))]
    else:
        argv = [*prefix, "--model-path", spec.model,
                "--served-model-name", spec.served_name,
                "--host", spec.host, "--port", str(spec.port),
                "--dtype", dtype,
                "--context-length", str(spec.max_model_len),
                "--tp", str(spec.tensor_parallel)]
        if spec.cpu:
            argv += ["--device", "cpu"]
        else:
            argv += ["--mem-fraction-static", "0.85"]
        if spec.lora_path:
            argv += ["--enable-lora", "--lora-paths", f"{spec.tag}={spec.lora_path}",
                     "--max-lora-rank", str(spec.lora_rank or 16)]
    return [*argv, *spec.extra_args]


def build_env(spec: OpenAIServerSpec, base_env: dict[str, str] | None = None) -> dict[str, str]:
    """Environment variables this launch adds (the CPU backend's), without overriding the caller's."""
    env = dict(os.environ if base_env is None else base_env)
    added: dict[str, str] = {}
    if spec.cpu and spec.flavour == "vllm" and "VLLM_CPU_KVCACHE_SPACE" not in env:
        added["VLLM_CPU_KVCACHE_SPACE"] = str(spec.cpu_kvcache_gib)
    return added


def gpu_refusal(spec: OpenAIServerSpec) -> str | None:
    """None when `spec` can run without a GPU; else why it cannot."""
    if spec.quantization:
        return (f"the checkpoint is {spec.quantization}-quantized, which {spec.flavour} serves "
                "on a GPU only — serve the unquantized checkpoint (--checkpoint .../checkpoint)")
    if spec.tensor_parallel > 1:
        return f"tensor_parallel={spec.tensor_parallel} needs that many GPUs"
    return None


def detect_gpu() -> str | None:
    """A GPU this host exposes (ROCm `/dev/kfd`, NVIDIA `/dev/nvidia0`), or None."""
    for dev, name in (("/dev/kfd", "rocm"), ("/dev/nvidia0", "cuda")):
        if Path(dev).exists():
            return name
    return None


def adapter_rank(adapter_dir: Path) -> int | None:
    """The LoRA rank `r` from `adapter_config.json`, or None."""
    try:
        r = json.loads((Path(adapter_dir) / "adapter_config.json").read_text("utf-8")).get("r")
    except (OSError, ValueError):
        return None
    return int(r) if isinstance(r, int) and r > 0 else None


def resolve_launcher(flavour: Flavour, server_bin: str | None = None) -> tuple[list[str] | None, str]:
    """(launcher argv prefix, "") or (None, reason with the install hint). Never raises.

    vLLM: `server_bin` or `vllm` on PATH. SGLang: `server_bin` is a Python interpreter that has
    `sglang` (checked in a subprocess), else this interpreter when `find_spec("sglang")` finds it.
    """
    if flavour == "vllm":
        binary = server_bin or shutil.which("vllm")
        if not binary or not (shutil.which(binary) or Path(binary).is_file()):
            return None, f"`vllm` not found on PATH — {VLLM_INSTALL_HINT}"
        return [binary], ""
    if server_bin:
        probe = "import importlib.util,sys; sys.exit(importlib.util.find_spec('sglang') is None)"
        try:
            rc = subprocess.run([server_bin, "-c", probe], capture_output=True, timeout=60,
                                check=False).returncode
        except (OSError, subprocess.SubprocessError) as exc:
            return None, f"cannot run {server_bin}: {exc} — {SGLANG_INSTALL_HINT}"
        if rc != 0:
            return None, f"{server_bin} has no `sglang` package — {SGLANG_INSTALL_HINT}"
        return [server_bin, "-m", "sglang.launch_server"], ""
    try:
        found = importlib.util.find_spec("sglang") is not None
    except (ImportError, ValueError):
        found = False
    if not found:
        return None, f"`sglang` is not importable from {sys.executable} — {SGLANG_INSTALL_HINT}"
    return [sys.executable, "-m", "sglang.launch_server"], ""


def server_dir(run_dir: Path, flavour: Flavour) -> Path:
    """Where the log, pid file and launch record live: `<run_dir>/serve/<flavour>/`."""
    return Path(run_dir) / "serve" / flavour


def _alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
    except ProcessLookupError:
        return False
    except PermissionError:
        return True
    # A zombie still answers kill(0); /proc says whether it is really running.
    try:
        stat = Path(f"/proc/{pid}/stat").read_text()
        return stat.rsplit(")", 1)[1].split()[0] != "Z"
    except (OSError, IndexError):
        return True


def _read_pid(pid_file: Path) -> int | None:
    try:
        return int(pid_file.read_text().strip())
    except (OSError, ValueError):
        return None


def _tail(path: Path, n: int = 2000) -> str:
    try:
        return path.read_text(encoding="utf-8", errors="replace")[-n:]
    except OSError:
        return ""


# ---- results ------------------------------------------------------------------------------


@dataclass(frozen=True)
class ServerLaunchResult:
    """Outcome of `launch_openai_server`. `ready` (or `dry_run`) is success."""

    status: LaunchStatus
    flavour: Flavour
    tag: str
    reason: str = ""
    argv: tuple[str, ...] = ()
    env: dict[str, str] = field(default_factory=dict)
    pid: int | None = None
    base_url: str = ""
    client_model: str = ""
    models: tuple[str, ...] = ()
    log_path: Path | None = None
    pid_file: Path | None = None
    merged_dir: Path | None = None
    mindx_fallback_swap: dict[str, str] | None = None
    log: tuple[str, ...] = ()

    @property
    def ok(self) -> bool:
        return self.status in ("ready", "dry_run")


@dataclass(frozen=True)
class ServerStopResult:
    status: StopStatus
    flavour: Flavour
    pid: int | None = None
    reason: str = ""

    @property
    def ok(self) -> bool:
        return self.status in ("stopped", "not_running")


# ---- readiness ----------------------------------------------------------------------------


def wait_ready(
    spec: OpenAIServerSpec,
    *,
    timeout_s: float,
    poll_interval_s: float = 2.0,
    proc: Any = None,
    http_client: httpx.Client | None = None,
    sleep: Callable[[float], None] = time.sleep,
) -> tuple[str, tuple[str, ...]]:
    """Poll until `/v1/models` lists the expected id (and vLLM's `/health` is 200).

    Returns ("ready" | "exited" | "timeout", model ids last seen). The expected id is the tag
    for vLLM (served or LoRA name) and the served name for SGLang.
    """
    expect = spec.tag if spec.flavour == "vllm" else spec.served_name
    client = http_client or httpx.Client(timeout=5.0)
    seen: tuple[str, ...] = ()
    deadline = time.monotonic() + timeout_s
    try:
        while True:
            if proc is not None and proc.poll() is not None:
                return "exited", seen
            try:
                healthy = True
                if spec.flavour == "vllm":
                    healthy = client.get(f"{spec.base_url}/health").status_code == 200
                if healthy:
                    resp = client.get(f"{spec.base_url}/v1/models")
                    if resp.status_code == 200:
                        data = resp.json().get("data") or []
                        seen = tuple(str(m.get("id")) for m in data if isinstance(m, dict))
                        if expect in seen:
                            return "ready", seen
            except (httpx.HTTPError, ValueError, AttributeError):
                pass
            if time.monotonic() >= deadline:
                return "timeout", seen
            sleep(poll_interval_s)
    finally:
        if http_client is None:
            client.close()


# ---- launch / stop ------------------------------------------------------------------------


def launch_openai_server(
    flavour: Flavour,
    *,
    base_model: str,
    tag: str,
    run_dir: Path,
    adapter_dir: Path | None = None,
    model_dir: Path | None = None,
    merge: bool = False,
    host: str = "127.0.0.1",
    port: int = 8000,
    dtype: Dtype = "auto",
    max_model_len: int = 8192,
    tensor_parallel: int = 1,
    quantization: str | None = None,
    cpu: bool | None = None,
    cpu_kvcache_gib: int = 4,
    extra_args: tuple[str, ...] | list[str] = (),
    server_bin: str | None = None,
    dry_run: bool = False,
    ready_timeout_s: float = 600.0,
    poll_interval_s: float = 2.0,
    register_with_mindx: bool = False,
    mindx_base_url: str | None = None,
    http_client: httpx.Client | None = None,
    sleep: Callable[[float], None] = time.sleep,
    sink: Callable[[str], None] | None = None,
) -> ServerLaunchResult:
    """Launch vLLM / SGLang on a trained checkpoint, detached, and wait for it to answer.

    Give exactly one of `adapter_dir` (a PEFT LoRA over `base_model`; served natively unless
    `merge`) or `model_dir` (a full model directory — merged or quantized). `cpu=None` decides
    from `detect_gpu()`. `dry_run` builds and returns the argv and touches nothing.

    The server runs in its own session with stdout+stderr in `<run_dir>/serve/<flavour>/
    server.log` and its pid in `server.pid`; it outlives this call (stop it with
    `stop_openai_server`). On `timeout` it is left running — big models load slowly.
    """
    lines: list[str] = []

    def emit(line: str) -> None:
        lines.append(line)
        if sink:
            sink(line)

    def done(status: LaunchStatus, **kw: Any) -> ServerLaunchResult:
        return ServerLaunchResult(status=status, flavour=flavour, tag=tag, log=tuple(lines), **kw)

    try:
        if (adapter_dir is None) == (model_dir is None):
            return done("error", reason="give exactly one of adapter_dir or model_dir")
        gpu = detect_gpu() if cpu is None else (None if cpu else "given")
        on_cpu = gpu is None
        sdir = server_dir(run_dir, flavour)
        lora = None if (adapter_dir is None or merge) else Path(adapter_dir)
        spec = OpenAIServerSpec(
            flavour=flavour,
            model=base_model if lora else str(model_dir or sdir / "merged"),
            tag=tag, lora_path=str(lora.resolve()) if lora else None,
            lora_rank=adapter_rank(lora) if lora else None,
            host=host, port=port, dtype=dtype, max_model_len=max_model_len,
            tensor_parallel=tensor_parallel,
            quantization=quantization if flavour == "vllm" else None,
            cpu=on_cpu, cpu_kvcache_gib=cpu_kvcache_gib, extra_args=tuple(extra_args),
        )
        needs_gpu = gpu_refusal(spec) if on_cpu else None
        if flavour == "sglang" and quantization and on_cpu:
            needs_gpu = (f"the checkpoint is {quantization}-quantized, which sglang serves on a "
                         "GPU only — serve the unquantized checkpoint (--checkpoint .../checkpoint)")

        if dry_run:
            launcher, why = resolve_launcher(flavour, server_bin)
            argv = build_argv(spec, launcher)
            env = build_env(spec)
            if merge and adapter_dir is not None:
                emit(f"[serve-{flavour}] would merge {adapter_dir} into {sdir / 'merged'} first")
            for note in (why, needs_gpu and f"needs a GPU: {needs_gpu}"):
                if note:
                    emit(f"[serve-{flavour}] note: {note}")
            if env:
                emit(f"[serve-{flavour}] env: " + " ".join(f"{k}={v}" for k, v in env.items()))
            return done("dry_run", argv=tuple(argv), env=env, base_url=spec.base_url,
                        client_model=spec.client_model, reason=why or (needs_gpu or ""))

        if needs_gpu:
            return done("no_gpu", reason=f"no GPU on this host and {needs_gpu}")
        launcher, why = resolve_launcher(flavour, server_bin)
        if launcher is None:
            return done("missing", reason=why)

        pid_file, log_path = sdir / "server.pid", sdir / "server.log"
        old = _read_pid(pid_file)
        if old is not None and _alive(old):
            return done("already_running", pid=old, pid_file=pid_file, log_path=log_path,
                        reason=(f"{flavour} already running as pid {old} ({pid_file}); "
                                f"stop it first (`serve --to {flavour} --stop`)"))

        merged: Path | None = None
        if merge and adapter_dir is not None:
            try:
                from mindxtrain.deploy.ollama_push import merge_lora_adapter

                merged = merge_lora_adapter(base_model, Path(adapter_dir), sdir / "merged",
                                            sink=emit)
            except ImportError as exc:
                return done("merge_failed", reason=f"{exc} — merging needs `uv sync --extra ml`")
            except Exception as exc:  # a merge failure is a result, not a crash
                return done("merge_failed", reason=f"{type(exc).__name__}: {exc}")
            spec = spec.model_copy(update={"model": str(merged)})
        elif model_dir is not None and not Path(model_dir).is_dir():
            return done("error", reason=f"model_dir {model_dir} is not a directory")

        argv = build_argv(spec, launcher)
        added = build_env(spec)
        sdir.mkdir(parents=True, exist_ok=True)
        emit(f"[serve-{flavour}] $ {' '.join(argv)}")
        with log_path.open("ab") as logf:
            proc = subprocess.Popen(argv, stdout=logf, stderr=subprocess.STDOUT,
                                    stdin=subprocess.DEVNULL, start_new_session=True,
                                    env={**os.environ, **added})
        pid_file.write_text(f"{proc.pid}\n", encoding="utf-8")
        (sdir / "launch.json").write_text(json.dumps({
            "flavour": flavour, "tag": tag, "pid": proc.pid, "argv": argv, "env": added,
            "base_url": spec.base_url, "client_model": spec.client_model,
            "started_at": datetime.now(tz=UTC).isoformat(),
        }, indent=1), encoding="utf-8")
        emit(f"[serve-{flavour}] pid {proc.pid}, log {log_path}; waiting for {spec.base_url}")

        state, models = wait_ready(spec, timeout_s=ready_timeout_s,
                                   poll_interval_s=poll_interval_s, proc=proc,
                                   http_client=http_client, sleep=sleep)
        common: dict[str, Any] = dict(argv=tuple(argv), env=added, pid=proc.pid,
                                      base_url=spec.base_url, client_model=spec.client_model,
                                      models=models, log_path=log_path, pid_file=pid_file,
                                      merged_dir=merged)
        if state == "exited":
            pid_file.unlink(missing_ok=True)
            return done("exited", reason=(f"{flavour} exited with code {proc.poll()} before "
                                          f"answering; log tail:\n{_tail(log_path, 1200)}"),
                        **common)
        if state == "timeout":
            return done("timeout", reason=(
                f"{spec.base_url} did not list {spec.client_model!r} within {ready_timeout_s:g}s "
                f"(seen: {', '.join(models) or 'nothing'}); still running as pid {proc.pid} — "
                f"watch {log_path}, or `serve --to {flavour} --stop`"), **common)

        swap: dict[str, str] | None = None
        if register_with_mindx:
            try:
                from mindxtrain.deploy.api_client import swap_mindx_fallback_model

                swap = swap_mindx_fallback_model(provider=flavour, model=spec.client_model,
                                                 api_url=mindx_base_url)
                emit(f"[serve-{flavour}] mindX swap: {swap.get('previous', '?')} -> "
                     f"{swap.get('current', '?')}")
            except Exception as exc:  # best-effort, as push_to_ollama
                emit(f"[serve-{flavour}] mindX registration failed (server still up): {exc}")
        emit(f"[serve-{flavour}] ready: {spec.client_model} at {spec.base_url}/v1")
        common["mindx_fallback_swap"] = swap
        return done("ready", **common)
    except Exception as exc:  # never past the result
        return done("error", reason=f"{type(exc).__name__}: {exc}")


def stop_openai_server(
    flavour: Flavour,
    run_dir: Path,
    *,
    timeout_s: float = 30.0,
    poll_interval_s: float = 0.5,
    sleep: Callable[[float], None] = time.sleep,
) -> ServerStopResult:
    """SIGTERM the server recorded in `<run_dir>/serve/<flavour>/server.pid` (its whole process
    group — it was started in its own session) and verify it is gone. Never raises."""
    pid_file = server_dir(run_dir, flavour) / "server.pid"
    try:
        pid = _read_pid(pid_file)
        if pid is None:
            return ServerStopResult("not_running", flavour, reason=f"no pid file at {pid_file}")
        if not _alive(pid):
            pid_file.unlink(missing_ok=True)
            return ServerStopResult("not_running", flavour, pid,
                                    reason=f"pid {pid} was not running (stale pid file removed)")
        try:
            os.killpg(pid, signal.SIGTERM)
        except (ProcessLookupError, PermissionError):
            os.kill(pid, signal.SIGTERM)
        deadline = time.monotonic() + timeout_s
        while _alive(pid):
            if time.monotonic() >= deadline:
                return ServerStopResult("still_running", flavour, pid, reason=(
                    f"pid {pid} still alive {timeout_s:g}s after SIGTERM; "
                    f"`kill -KILL -{pid}` ends its group"))
            sleep(poll_interval_s)
        pid_file.unlink(missing_ok=True)
        return ServerStopResult("stopped", flavour, pid)
    except Exception as exc:
        return ServerStopResult("error", flavour, reason=f"{type(exc).__name__}: {exc}")


__all__ = [
    "VLLM_LORA_RANKS",
    "VLLM_QUANTIZATION",
    "OpenAIServerSpec",
    "ServerLaunchResult",
    "ServerStopResult",
    "adapter_rank",
    "build_argv",
    "build_env",
    "default_launcher",
    "detect_gpu",
    "gpu_refusal",
    "launch_openai_server",
    "resolve_launcher",
    "server_dir",
    "stop_openai_server",
    "wait_ready",
]