from huggingface_hub import hf_hub_download
import os
import shutil
import subprocess
import sys

repo = os.environ.get("MODEL_REPO", "mradermacher/NVIDIA-Nemotron-3-Nano-4B-BF16-GGUF")
filename = os.environ.get("MODEL_FILE", "NVIDIA-Nemotron-3-Nano-4B-BF16.Q4_K_M.gguf")
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
cache_dir = os.environ.get("HF_HOME", "/data/.cache/huggingface")
local_dir = os.environ.get("MODEL_DIR", "/data/models")
os.makedirs(local_dir, exist_ok=True)

model_path = os.path.join(local_dir, filename)
if os.path.exists(model_path) and os.path.getsize(model_path) > 0:
    print(f"Using existing model at {model_path}", flush=True)
else:
    print(f"Downloading {repo}/{filename}...", flush=True)
    model_path = hf_hub_download(
        repo_id=repo,
        filename=filename,
        token=token,
        cache_dir=cache_dir,
        local_dir=local_dir,
    )

server_bin = shutil.which("llama-server")
if server_bin is None:
    for candidate in ("/app/llama-server", "/usr/local/bin/llama-server", "/usr/bin/llama-server", "/llama-server"):
        if os.path.exists(candidate):
            server_bin = candidate
            break
if server_bin is None:
    print("Could not find llama-server", flush=True)
    raise FileNotFoundError("llama-server")

cmd = [
    server_bin,
    "-m", model_path,
    "--host", "0.0.0.0",
    "--port", os.environ.get("PORT", "7860"),
    "-c", os.environ.get("CTX_SIZE", "2048"),
    "-ngl", os.environ.get("N_GPU_LAYERS", "0"),
]

if os.environ.get("NO_WARMUP", "1").lower() in {"1", "true", "yes", "on"}:
    cmd.append("--no-warmup")

print("Starting:", " ".join(cmd), flush=True)
sys.stdout.flush()
subprocess.run(cmd, check=True)