#!/usr/bin/env python3
"""Detect a Llama 4 variant from its config.json.

The unified image can serve both:

  - Llama-4-Scout-17B-16E-Instruct  (16 local experts, BF16) -> TP8, PP1, 8 GPUs
  - Llama-4-Maverick-17B-128E-Instruct-FP8 (128 local experts, FP8) -> TP8, PP2, 16 GPUs

The detected values are printed as shell assignments by default, so run.sh and
entrypoint.sh can `eval` them.  `--json` is used by verify.sh.

Detection is based on `text_config.num_local_experts`, with a model-directory
name fallback if the field is missing.  Unsupported configs fail loudly rather
than silently starting with the wrong TP/PP.
"""

from __future__ import annotations

import argparse
import json
import pathlib
import shlex
import sys

PROFILES = {
    "scout": {
        "tp": 8,
        "pp": 1,
        "port": 8091,
        "served_name": "llama4-scout",
        "gpu_count": 8,
        "description": "Llama 4 Scout 17B-16E Instruct (BF16, TP8, 8 GPUs)",
    },
    "maverick": {
        "tp": 8,
        "pp": 2,
        "port": 8092,
        "served_name": "llama4-maverick",
        "gpu_count": 16,
        "description": "Llama 4 Maverick 17B-128E Instruct FP8 (TP8+PP2, 16 GPUs)",
    },
}


def load_config(model_dir: str) -> dict:
    path = pathlib.Path(model_dir) / "config.json"
    if not path.is_file():
        raise ValueError(f"config.json not found: {path}")
    try:
        return json.loads(path.read_text())
    except json.JSONDecodeError as exc:
        raise ValueError(f"invalid config.json: {exc}") from exc


def detect(model_dir: str) -> dict:
    cfg = load_config(model_dir)
    model_type = str(cfg.get("model_type", ""))
    if model_type != "llama4":
        raise ValueError(
            f"unsupported model_type={model_type!r}; expected 'llama4'"
        )

    text = cfg.get("text_config") or {}
    experts = text.get("num_local_experts")
    if experts is None:
        experts = text.get("num_experts")

    if experts is not None:
        experts = int(experts)
        if experts >= 64:
            kind = "maverick"
        elif experts <= 32:
            kind = "scout"
        else:
            raise ValueError(
                f"cannot determine Llama 4 variant: num_local_experts={experts}; "
                "expected 16 (Scout) or 128 (Maverick)"
            )
        detection = f"num_local_experts={experts}"
    else:
        name = pathlib.Path(model_dir).name.lower()
        if "maverick" in name:
            kind = "maverick"
        elif "scout" in name:
            kind = "scout"
        else:
            raise ValueError(
                "cannot determine Llama 4 variant: text_config.num_local_experts "
                "is missing and the directory name does not contain scout/maverick"
            )
        detection = "directory name"

    profile = dict(PROFILES[kind])
    profile.update(
        {
            "kind": kind,
            "model_type": model_type,
            "architecture": (cfg.get("architectures") or [""])[0],
            "num_local_experts": experts,
            "quantized": bool(cfg.get("quantization_config")),
            "detection": detection,
            "model_dir": str(pathlib.Path(model_dir).resolve()),
        }
    )
    return profile


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Detect Llama 4 Scout/Maverick from config.json"
    )
    parser.add_argument("model_dir", help="directory containing config.json")
    parser.add_argument(
        "--json",
        action="store_true",
        help="print JSON instead of shell assignments",
    )
    args = parser.parse_args()

    try:
        profile = detect(args.model_dir)
    except ValueError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 2

    if args.json:
        print(json.dumps(profile, ensure_ascii=False, indent=2))
        return 0

    values = {
        "DETECTED_MODEL_KIND": profile["kind"],
        "DETECTED_TP": profile["tp"],
        "DETECTED_PP": profile["pp"],
        "DETECTED_PORT": profile["port"],
        "DETECTED_SERVED_NAME": profile["served_name"],
        "DETECTED_GPU_COUNT": profile["gpu_count"],
        "DETECTED_DESCRIPTION": profile["description"],
    }
    for key, value in values.items():
        print(f"{key}={shlex.quote(str(value))}")
    return 0


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