"""Download APUS-OpenJev-v1 merged 9B weights and launch standard vLLM serving."""

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

MODEL_REPO = "apus-ailab/APUS-OpenJev-v1"
MODEL_SUBFOLDER = "9B-5949"
MODEL_ALIAS = "APUS-OpenJev-v1-9B"


def engine_command(model, host, port):
    return [
        sys.executable,
        "-m",
        "vllm.entrypoints.openai.api_server",
        "--model",
        str(model),
        "--served-model-name",
        MODEL_ALIAS,
        "--host",
        host,
        "--port",
        str(port),
        "--dtype",
        "bfloat16",
        "--tensor-parallel-size",
        "1",
        "--max-model-len",
        "8192",
        "--language-model-only",
        "--gpu-memory-utilization",
        "0.7",
        "--max-num-seqs",
        "32",
        "--max-num-batched-tokens",
        "8192",
        "--max-logprobs",
        "256",
        "--logprobs-mode",
        "raw_logprobs",
        "--generation-config",
        "vllm",
        "--no-enable-prefix-caching",
    ]


def select_catalog_model(repo, revision):
    from huggingface_hub import hf_hub_download
    from huggingface_hub.errors import EntryNotFoundError

    try:
        path = hf_hub_download(repo_id=repo, filename="meta.yaml", revision=revision)
    except EntryNotFoundError:
        # Preserve explicit older revisions published before the family catalog.
        path = hf_hub_download(
            repo_id=repo, filename="bundle-manifest.json", revision=revision
        )
        manifest = json.loads(Path(path).read_text())
        if manifest.get("repo_id") != repo:
            raise ValueError("Unexpected legacy family manifest")
        matches = [
            item for item in manifest["files"]
            if item.get("path") == MODEL_SUBFOLDER + "/config.json"
        ]
        if len(matches) != 1:
            raise ValueError("Missing or duplicate legacy config identity")
        entry = {
            "subfolder": MODEL_SUBFOLDER,
            "config_sha256": matches[0]["sha256"],
        }
    else:
        catalog = json.loads(Path(path).read_text())
        if (
            catalog.get("schema") != "apus-openjev-family.v1"
            or catalog.get("repo_id") != repo
        ):
            raise ValueError("Unexpected model-family catalog")
        entry = catalog["variants"][MODEL_SUBFOLDER]
    if entry.get("subfolder") != MODEL_SUBFOLDER or not re.fullmatch(
        r"[0-9a-f]{64}", entry.get("config_sha256", "")
    ):
        raise ValueError("Unexpected model variant identity")
    return entry


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--revision",
        default="main",
        help="Commit SHA or branch; resolved to an immutable SHA before downloading",
    )
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8000)
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print download and launch plan without network or GPU access",
    )
    args = parser.parse_args()
    if not 1 <= args.port <= 65535:
        parser.error("--port must be between 1 and 65535")
    if args.dry_run:
        print(
            json.dumps(
                {
                    "repo_id": MODEL_REPO,
                    "revision": args.revision,
                    "subfolder": MODEL_SUBFOLDER,
                    "command": engine_command(
                        f"<snapshot>/{MODEL_SUBFOLDER}", args.host, args.port
                    ),
                },
                indent=2,
            )
        )
        return
    from importlib.metadata import version
    from huggingface_hub import HfApi, snapshot_download

    if version("vllm") != "0.29.0":
        parser.error("Install vllm==0.29.0 in the active Python environment")
    revision = HfApi().model_info(MODEL_REPO, revision=args.revision).sha
    if not revision or not re.fullmatch(r"[0-9a-f]{40}", revision):
        parser.error("The Hub did not return an immutable model commit SHA")
    entry = select_catalog_model(MODEL_REPO, revision)
    snapshot = snapshot_download(
        repo_id=MODEL_REPO, revision=revision, allow_patterns=[f"{MODEL_SUBFOLDER}/*"]
    )
    model = Path(snapshot) / MODEL_SUBFOLDER
    if not (model / "config.json").is_file() or not any(model.glob("*.safetensors")):
        parser.error(
            "Downloaded subfolder is missing config.json or merged safetensors weights"
        )
    if hashlib.sha256((model / "config.json").read_bytes()).hexdigest() != entry["config_sha256"]:
        parser.error("Downloaded model config does not match family catalog")
    if (model / "adapter_config.json").exists():
        parser.error("Expected merged weights, not an adapter-only directory")
    print(
        json.dumps(
            {
                "repo_id": MODEL_REPO,
                "revision": revision,
                "model_path": str(model),
                "served_model_name": MODEL_ALIAS,
            }
        ),
        flush=True,
    )
    os.environ["VLLM_USE_FLASHINFER_SAMPLER"] = "0"
    command = engine_command(model, args.host, args.port)
    os.execv(sys.executable, command)


if __name__ == "__main__":
    main()