import os
import sys
import logging
import traceback
import spaces
import gradio as gr
from huggingface_hub import hf_hub_download

logging.basicConfig(level=logging.INFO)
log = logging.getLogger(__name__)

MODEL_REPO = "sakamakismile/gemma-4-12B-coder-fable5-composer2.5-GGUF"
MODEL_FILE = "gemma-4-12B-coder-fable5-composer2.5-Q4_K_M.gguf"
MODEL_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), MODEL_FILE)

if not os.path.exists(MODEL_PATH):
    log.info("Downloading model (7.38 GB)...")
    hf_hub_download(
        repo_id=MODEL_REPO,
        filename=MODEL_FILE,
        local_dir=os.path.dirname(os.path.abspath(__file__)),
    )
    log.info("Download complete")

def _setup_cuda_paths():
    nvidia_pkgs = ["nvidia", "nvidia.cuda_runtime", "nvidia.cublas"]
    candidates = set()
    for mod_name in nvidia_pkgs:
        try:
            mod = __import__(mod_name, fromlist=["__path__"])
            pkg_path = mod.__path__[0]
            candidates.add(os.path.join(os.path.dirname(pkg_path), mod_name.split(".")[-1], "lib"))
        except Exception:
            pass
    for p in sys.path:
        candidates.add(os.path.join(p, "nvidia", "cuda_runtime", "lib"))
        candidates.add(os.path.join(p, "nvidia", "cublas", "lib"))
    candidates.add("/usr/local/lib/python3.12/site-packages/nvidia/cuda_runtime/lib")
    candidates.add("/usr/local/lib/python3.12/site-packages/nvidia/cublas/lib")
    current = os.environ.get("LD_LIBRARY_PATH", "")
    merged = []
    for p in candidates:
        if os.path.exists(p) and p not in merged:
            merged.append(p)
    for p in merged:
        os.environ["LD_LIBRARY_PATH"] = p + ":" + os.environ.get("LD_LIBRARY_PATH", "")
    log.info(f"CUDA lib paths: {merged}")

@spaces.GPU
def cuda_test():
    import torch
    return {
        "cuda_available": torch.cuda.is_available(),
        "device_count": torch.cuda.device_count(),
        "device_name": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "N/A",
    }

@spaces.GPU
def generate(messages, max_tokens=1024, temperature=0.7, top_p=0.95):
    try:
        _setup_cuda_paths()
        log.info(f"LD_LIBRARY_PATH={os.environ.get('LD_LIBRARY_PATH', '')[:200]}")
        log.info("Importing llama_cpp...")
        from llama_cpp import Llama as _Llama
        log.info("Import OK, loading model...")
        llm = _Llama(
            model_path=MODEL_PATH,
            n_gpu_layers=-1,
            n_ctx=8192,
            verbose=False
        )
        log.info("Model loaded, generating...")
        output = llm.create_chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        result = output["choices"][0]["message"]["content"].strip()
        log.info("Generation complete")
        return result
    except Exception as e:
        err = f"GPU Error: {type(e).__name__}: {e}\n{traceback.format_exc()}"
        log.error(err)
        return err

def predict(message, history):
    if not os.path.exists(MODEL_PATH):
        return "Model is still downloading... Please wait ~5 minutes and try again."
    messages = []
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})
    return generate(messages)

def check_status():
    exists = os.path.exists(MODEL_PATH)
    size = os.path.getsize(MODEL_PATH) if exists else 0
    cuda = cuda_test()
    return f"Model file: {'exists' if exists else 'missing'} ({size/1e9:.1f} GB)\nGPU: {cuda}"

with gr.Blocks(title="Gemma Coder Zero", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Gemma 4 12B Coder Zero")
    gr.Markdown("Powered by llama.cpp on ZeroGPU (RTX Pro 6000 Blackwell)")

    with gr.Tabs():
        with gr.TabItem("Chat"):
            gr.ChatInterface(
                fn=predict,
                title="Gemma Coder",
                description="Ask any coding question!"
            )
        with gr.TabItem("Status"):
            status_btn = gr.Button("Check GPU & Model Status")
            status_out = gr.Textbox(label="Status")
            status_btn.click(fn=check_status, outputs=status_out)

    gr.Markdown("---\nFirst request is slow (~5 min) while the 7.38 GB model downloads. Subsequent requests are fast.")

demo.queue(default_concurrency_limit=1)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)