#!/usr/bin/env python3
"""Hugging Face Space — MiniMax H3 Prompt Enhancer 2.6B (ZeroGPU + llama.cpp).

Loads Q4_K_M GGUF via llama-cpp-python (CUDA wheel). Generation runs inside
``@spaces.GPU`` so ``n_gpu_layers=-1`` can offload onto the ZeroGPU slice.

Companion (fast CPU 350M):
https://huggingface.co/spaces/geocine/MiniMax-H3-Prompt-Enhancer
"""

from __future__ import annotations

import ctypes
import os
import sys
from pathlib import Path

# Import spaces first so ZeroGPU can patch the runtime before Gradio starts.
import spaces  # noqa: F401

# Torch ships libcudart; preload so the CUDA llama-cpp wheel can dlopen later.
# Do NOT import llama_cpp at module scope — CUDA init belongs inside @spaces.GPU.
import torch

_torch_lib = Path(torch.__file__).resolve().parent / "lib"
for _name in ("libcudart.so.12", "libcudart.so"):
    _cand = _torch_lib / _name
    if _cand.is_file():
        try:
            ctypes.CDLL(str(_cand), mode=ctypes.RTLD_GLOBAL)
            print(f"[boot] preloaded {_cand.name} from torch", flush=True)
        except OSError as exc:
            print(f"[boot] libcudart preload skipped: {exc}", flush=True)
        break
_ld = os.environ.get("LD_LIBRARY_PATH", "")
if str(_torch_lib) not in _ld.split(":"):
    os.environ["LD_LIBRARY_PATH"] = (
        f"{_torch_lib}:{_ld}" if _ld else str(_torch_lib)
    )

ROOT = Path(__file__).resolve().parent
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from huggingface_hub import hf_hub_download

from minimax.apps.gradio import build_app
from minimax.modeling.gguf_runtime import effective_cpus, resolve_gguf_file
from minimax.paths import DEFAULT_GGUF_FILE, DEFAULT_GGUF_REPO

GGUF_REPO = os.environ.get("GGUF_REPO", DEFAULT_GGUF_REPO)
GGUF_FILE = os.environ.get("GGUF_FILE", DEFAULT_GGUF_FILE)
LOCAL_MODELS = ROOT / "models"
N_THREADS = int(os.environ.get("N_THREADS", "0")) or None
N_CTX = int(os.environ.get("N_CTX", "4096"))
N_GPU_LAYERS = int(os.environ.get("N_GPU_LAYERS", "-1"))
GPU_DURATION = int(os.environ.get("GPU_DURATION", "120"))


def resolve_model_path() -> str:
    env = os.environ.get("GGUF_PATH") or os.environ.get("MODEL_ID")
    if env and (Path(env).exists() or str(env).endswith(".gguf")):
        p = Path(env)
        if p.exists():
            return str(resolve_gguf_file(p) if p.is_dir() else p)
        if str(env).endswith(".gguf") and (LOCAL_MODELS / env).is_file():
            return str(LOCAL_MODELS / env)

    if LOCAL_MODELS.is_dir() and list(LOCAL_MODELS.glob("*.gguf")):
        path = str(resolve_gguf_file(LOCAL_MODELS))
        print(f"[boot] using local GGUF {path}", flush=True)
        return path

    print(f"[boot] resolving {GGUF_REPO}/{GGUF_FILE} …", flush=True)
    return hf_hub_download(repo_id=GGUF_REPO, filename=GGUF_FILE)


print(
    f"[boot] torch={torch.__version__}  cpus={effective_cpus()}  "
    f"n_gpu_layers={N_GPU_LAYERS}",
    flush=True,
)
MODEL_PATH = resolve_model_path()
print(
    f"[boot] GGUF={MODEL_PATH}  n_threads={N_THREADS or 'auto'}  n_ctx={N_CTX}  "
    f"gpu_duration={GPU_DURATION}s",
    flush=True,
)
print("[boot] building Gradio UI (GGUF loads on first Generate inside ZeroGPU)", flush=True)

demo = build_app(
    model_path=MODEL_PATH,
    backend="gguf",
    n_ctx=N_CTX,
    n_threads=N_THREADS,
    n_gpu_layers=N_GPU_LAYERS,
    prefer_cpu=False,
    gpu_duration=GPU_DURATION,
)
print("[boot] Gradio ready", flush=True)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1)
    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        # SSR is experimental and makes Space page loads very slow.
        ssr_mode=False,
    )