import os, sys
import asyncio
import subprocess
from pathlib import Path

from vllm.utils.argparse_utils import FlexibleArgumentParser
from vllm.entrypoints.openai.api_server import run_server
from vllm.entrypoints.openai.cli_args import make_arg_parser


def _pick_free_gpu():
    if "CUDA_VISIBLE_DEVICES" in os.environ:
        return os.environ["CUDA_VISIBLE_DEVICES"]
    try:

        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=index,memory.free", "--format=csv,noheader"],
            text=True,
        )
        lines = [l.strip() for l in out.strip().split("\n") if l.strip()]
        best = max(lines, key=lambda x: int(x.split(",")[1].strip().split()[0]))
        gpu = best.split(",")[0].strip()
    except Exception:
        gpu = "0"
        
    return gpu


if __name__ == "__main__":
    gpu = _pick_free_gpu()
    if len(sys.argv) > 1:
        gpu = sys.argv[1]

    os.environ["CUDA_VISIBLE_DEVICES"] = gpu
    os.environ["FLASHINFER_DISABLE_VERSION_CHECK"] = "1"

    model = str(Path(__file__).resolve().parent / "models/Qwen2.5-7B-Instruct")

    parser = make_arg_parser(FlexibleArgumentParser())
    args = parser.parse_args([
        "--model", model,
        "--host", "0.0.0.0",
        "--port", "8000",
        "--served-model-name", "Qwen2.5-7B-Instruct",
        "--trust-remote-code",
        "--max-model-len", "32768",
        "--dtype", "bfloat16",
        "--gpu-memory-utilization", "0.85",
        "--enforce-eager",
        "--enable-auto-tool-choice",
        "--tool-call-parser", "hermes",
    ])

    asyncio.run(run_server(args))