from __future__ import annotations

import argparse
import json
import os
import time
import urllib.request
from pathlib import Path
from typing import Any


def _pick_qwen_snapshot() -> str | None:
    base = Path("/private/hf/hub/models--Qwen--Qwen2.5-0.5B-Instruct/snapshots")
    if not base.exists():
        return None
    snaps = sorted([p for p in base.iterdir() if p.is_dir()], reverse=True)
    return str(snaps[0]) if snaps else None


def _http_get_json(url: str) -> Any:
    with urllib.request.urlopen(url, timeout=10) as resp:
        raw = resp.read().decode("utf-8")
        return json.loads(raw)


def _wait_http_json(url: str, *, timeout_s: int) -> Any:
    deadline = time.time() + float(timeout_s)
    last_err: Exception | None = None
    while time.time() < deadline:
        try:
            return _http_get_json(url)
        except Exception as e:
            last_err = e
            time.sleep(2)
    raise RuntimeError(f"timeout waiting for {url}: {last_err!r}")


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description="Ray Serve LLM smoke test (deploy + /v1/models probe).")
    ap.add_argument("--ray-address", default="auto")
    ap.add_argument("--http-port", type=int, default=8000)
    ap.add_argument("--app-name", default="argus_llm_smoke")
    ap.add_argument("--route-prefix", default="/")
    ap.add_argument("--accelerator-type", default=os.environ.get("ARGUS_ACCELERATOR_TYPE") or "H20")
    ap.add_argument("--model-id", default="smoke-qwen-0.5b")
    ap.add_argument("--model-source", default=None, help="Local path or HF id. Default: cached Qwen snapshot under /private/hf.")
    ap.add_argument("--tensor-parallel-size", type=int, default=1)
    ap.add_argument("--num-replicas", type=int, default=1)
    ap.add_argument("--wait-s", type=int, default=600)
    args = ap.parse_args(argv)

    model_source = str(args.model_source or _pick_qwen_snapshot() or "")
    if not model_source:
        raise SystemExit("missing --model-source and no cached Qwen snapshot found under /private/hf")

    # Force offline HF behavior for the smoke test.
    os.environ.setdefault("HF_HOME", "/private/hf")
    os.environ.setdefault("HUGGINGFACE_HUB_CACHE", "/private/hf/hub")
    os.environ.setdefault("TRANSFORMERS_CACHE", "/private/hf/transformers")
    os.environ["HF_HUB_OFFLINE"] = "1"

    import ray

    ray.init(address=str(args.ray_address), ignore_reinit_error=True, log_to_driver=False)

    from ray import serve

    try:
        serve.start(proxy_location="HeadOnly", http_options={"host": "0.0.0.0", "port": int(args.http_port)})
    except Exception:
        # Best-effort: Serve may already be running in the container (e.g., started by the MVP API scheduler).
        pass

    from ray.serve.llm import LLMConfig, build_openai_app

    # Build a config dict and filter by the current Ray's LLMConfig schema, since fields
    # may differ between Ray versions.
    cfg_dict: dict[str, Any] = {
        "model_loading_config": {"model_id": str(args.model_id), "model_source": model_source},
        "accelerator_type": str(args.accelerator_type),
        "deployment_config": {"num_replicas": int(args.num_replicas)},
        "engine_kwargs": {"tensor_parallel_size": int(args.tensor_parallel_size)},
        "runtime_env": {"env_vars": {"HF_HUB_OFFLINE": "1", "HF_HOME": "/private/hf"}},
    }
    allowed = set(getattr(LLMConfig, "model_fields", {}).keys())
    if allowed:
        cfg_dict = {k: v for k, v in cfg_dict.items() if k in allowed}

    llm_cfg = LLMConfig(**cfg_dict)
    app = build_openai_app({"llm_configs": [llm_cfg]})

    serve.run(app, name=str(args.app_name), route_prefix=str(args.route_prefix))

    models_url = f"http://127.0.0.1:{int(args.http_port)}/v1/models"
    payload = _wait_http_json(models_url, timeout_s=int(args.wait_s))
    print(json.dumps(payload, indent=2, sort_keys=True))

    return 0


if __name__ == "__main__":
    raise SystemExit(main())
