"""
ClinIQ — Modal deployment.
Runs Qwen2.5-3B-Instruct GGUF via llama.cpp on A10G GPU.

Earns: 🦙 Llama Champion (llama.cpp) + 🐜 Tiny Titan (≤4B params)

Deploy:   modal deploy modal_inference.py
Test:     modal run modal_inference.py
Endpoint: printed after deploy — set as MODAL_ENDPOINT Space secret
"""

import modal

# ── Image: llama.cpp + CUDA + model baked in ───────────────────────────────────

REPO_ID  = "bartowski/Qwen2.5-3B-Instruct-GGUF"
FILENAME = "Qwen2.5-3B-Instruct-Q4_K_M.gguf"
MODEL_PATH = f"/model/{FILENAME}"

image = (
    # Use CUDA 12.4 devel image — includes nvcc and toolkit needed to build llama.cpp with GPU
    modal.Image.from_registry(
        "nvidia/cuda:12.4.0-devel-ubuntu22.04",
        add_python="3.11",
    )
    .apt_install("git", "cmake", "build-essential", "libcurl4-openssl-dev", "wget")
    .run_commands(
        "git clone --depth 1 https://github.com/ggerganov/llama.cpp /llama.cpp",
        # Disable tests (they fail to link libcuda stubs in build containers)
        # Link against CUDA stubs so the shared libs resolve during build
        "cd /llama.cpp && cmake -B build "
        "  -DGGML_CUDA=ON "
        "  -DCMAKE_CUDA_ARCHITECTURES=86 "
        "  -DLLAMA_BUILD_TESTS=OFF "
        "  -DLLAMA_BUILD_EXAMPLES=OFF "
        "  -DCMAKE_EXE_LINKER_FLAGS='-L/usr/local/cuda/lib64/stubs -lcuda' "
        "  -DCMAKE_SHARED_LINKER_FLAGS='-L/usr/local/cuda/lib64/stubs -lcuda' "
        "&& cmake --build build --config Release --target llama-server -j$(nproc)",
        "cp /llama.cpp/build/bin/llama-server /usr/local/bin/llama-server",
    )
    .pip_install("huggingface_hub>=0.24.0", "httpx", "fastapi[standard]")
    .env({"HF_XET_HIGH_PERFORMANCE": "1"})
    .run_commands(
        f"hf download {REPO_ID} {FILENAME} --local-dir /model"
    )
)

app = modal.App("cliniq-inference", image=image)


# ── Inference class ────────────────────────────────────────────────────────────

@app.cls(
    gpu="A10G",
    scaledown_window=300,
)
@modal.concurrent(max_inputs=8)
class LlamaCppServer:

    @modal.enter()
    def start(self):
        import subprocess, time, httpx

        self._proc = subprocess.Popen(
            [
                "llama-server",
                "--model",        MODEL_PATH,
                "--ctx-size",     "4096",
                "--n-gpu-layers", "99",
                "--port",         "8080",
                "--host",         "127.0.0.1",
                "--threads",      "4",
            ],
        )
        # Poll until healthy
        for _ in range(90):
            try:
                if httpx.get("http://127.0.0.1:8080/health", timeout=2).status_code == 200:
                    print("✅ llama-server ready")
                    return
            except Exception:
                pass
            time.sleep(2)
        raise RuntimeError("llama-server did not start in 3 minutes")

    @modal.exit()
    def stop(self):
        self._proc.terminate()

    @modal.method()
    def generate(self, prompt: str, max_tokens: int = 600, json_mode: bool = False) -> str:
        import httpx

        payload = {
            "prompt":      prompt,
            "n_predict":   max_tokens,
            "temperature": 0.0,
            "stop":        ["<|im_end|>", "<|endoftext|>", "<|im_start|>"],
        }
        r = httpx.post("http://127.0.0.1:8080/completion", json=payload, timeout=120)
        r.raise_for_status()
        return r.json()["content"].strip()


# ── Web endpoint (called from Gradio Space) ────────────────────────────────────

@app.function()
@modal.fastapi_endpoint(method="POST", label="cliniq-infer")
def infer(item: dict) -> dict:
    """
    POST body: {"prompt": str, "max_tokens": int, "json_mode": bool}
    Returns:   {"text": str}
    """
    server = LlamaCppServer()
    text = server.generate.remote(
        item["prompt"],
        item.get("max_tokens", 600),
        item.get("json_mode", False),
    )
    return {"text": text}


# ── Health check endpoint ──────────────────────────────────────────────────────

@app.function()
@modal.fastapi_endpoint(method="GET", label="cliniq-health")
def health() -> dict:
    return {"status": "ok", "model": FILENAME}


# ── Local test ─────────────────────────────────────────────────────────────────

@app.local_entrypoint()
def test():
    server = LlamaCppServer()
    prompt = (
        "<|im_start|>system\nYou are a helpful clinical assistant.<|im_end|>\n"
        "<|im_start|>user\nWhat are the first-line treatments for community-acquired pneumonia?<|im_end|>\n"
        "<|im_start|>assistant\n"
    )
    print("\n=== Test Output ===")
    print(server.generate.remote(prompt, max_tokens=300))
    print("\n=== Structured Test ===")
    struct_prompt = (
        "<|im_start|>system\nExtract as JSON only.<|im_end|>\n"
        "<|im_start|>user\n"
        "Document: Patient takes Metformin 1000mg BID and Lisinopril 10mg daily. "
        "Allergic to Penicillin (rash).\n"
        'List medications as JSON: [{"name":"...","dose":"...","frequency":"..."}]<|im_end|>\n'
        "<|im_start|>assistant\n"
    )
    print(server.generate.remote(struct_prompt, max_tokens=200, json_mode=True))