"""Launch the vLLM OpenAI server with the project's energy/token efficiency settings applied.

    python scripts/serve_vllm.py --config config/config.yaml            # start the server
    python scripts/serve_vllm.py --config config/config.yaml --dry-run  # just print the command

Reads config and assembles the vLLM command with:
  - the int4 (w4a16) model + capped context length
  - fp8 KV cache            (efficiency.kv_cache_dtype)   -> less memory/energy per token
  - prefix caching          (efficiency.enable_prefix_caching) -> static system prompt not re-prefilled
  - speculative decoding     (efficiency.speculative)       -> lower latency/energy per request
      * method=ngram : prompt-lookup, no extra model, ideal for RAG's verbatim copying
      * method=eagle : uses a DeepSpec-trained Eagle3 draft model (efficiency.speculative.draft_model)

Speculative decoding does NOT change how many tokens are generated — it makes generating them
cheaper/faster. Token COUNT is minimized elsewhere (router, cache, compact prompt, context budget).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from agent.settings import load_config  # noqa: E402


def build_command(cfg, port: int) -> list[str]:
    model_dir = str(cfg.resolve("model.quantized_dir"))
    max_len = int(cfg.get("model.max_model_len", 4096))

    cmd = [
        sys.executable, "-m", "vllm.entrypoints.openai.api_server",
        "--model", model_dir,
        "--served-model-name", cfg.get("serving.served_model", model_dir),
        "--max-model-len", str(max_len),
        "--port", str(port),
    ]

    kv = cfg.get("efficiency.kv_cache_dtype")
    if kv:
        cmd += ["--kv-cache-dtype", kv]
    if bool(cfg.get("efficiency.enable_prefix_caching", False)):
        cmd += ["--enable-prefix-caching"]

    # Speculative decoding (vLLM v1 uses a single --speculative-config JSON blob).
    method = str(cfg.get("efficiency.speculative.method", "none")).lower()
    if method == "ngram":
        spec = {
            "method": "ngram",
            "num_speculative_tokens": int(cfg.get("efficiency.speculative.num_speculative_tokens", 5)),
            "prompt_lookup_max": int(cfg.get("efficiency.speculative.prompt_lookup_max", 4)),
            "prompt_lookup_min": int(cfg.get("efficiency.speculative.prompt_lookup_min", 2)),
        }
        cmd += ["--speculative-config", json.dumps(spec)]
    elif method == "eagle":
        draft = cfg.get("efficiency.speculative.draft_model", "")
        if not draft:
            raise SystemExit("efficiency.speculative.draft_model is required when method=eagle")
        spec = {
            "method": "eagle",
            "model": draft,
            "num_speculative_tokens": int(cfg.get("efficiency.speculative.num_speculative_tokens", 5)),
        }
        cmd += ["--speculative-config", json.dumps(spec)]
    # method == "none" -> no speculative flags

    return cmd


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--config", default="config/config.yaml")
    ap.add_argument("--port", type=int, default=8001)
    ap.add_argument("--dry-run", action="store_true", help="print the command instead of running it")
    args = ap.parse_args()
    cfg = load_config(args.config)

    cmd = build_command(cfg, args.port)
    printable = " ".join(f"'{c}'" if " " in c else c for c in cmd)
    print(printable + "\n")

    if args.dry_run:
        return
    # Replace this process with vLLM.
    os.execvp(cmd[0], cmd)


if __name__ == "__main__":
    main()