import os
import sys
import subprocess

# 1. FORCE EXTRACT PRE-BUILT CUDA WHEELS INSTANTLY BEFORE ANY IMPORTS
try:
    import llama_cpp
except ModuleNotFoundError:
    print("Package missing! Forcing immediate 5-second pre-compiled wheel installation...")
    subprocess.check_call([
        sys.executable, "-m", "pip", "install", 
        "llama-cpp-python", 
        "--extra-index-url", "https://github.io"
    ])

# Enforce cache directories to live on persistent storage
os.environ["HF_HOME"] = "/data/.huggingface"
os.environ["XDG_CACHE_HOME"] = "/data/.cache"

# CRITICAL FIX: Explicitly import hf_hub_download to fix the NameError
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import spaces
import gradio as gr

# 2. Pull down the standard 1-bit 3.9 GB file variant to your persistent bucket mount
print("Checking persistent storage for Bonsai-27B weights...")
model_path = hf_hub_download(
    repo_id="prism-ml/Bonsai-27B-gguf", 
    filename="Bonsai-27B-Q1_0.gguf", 
    local_dir="/data"
)

# 3. Load the model into memory globally
print("Initializing model architecture into global memory...")
llm = Llama(
    model_path=model_path, 
    n_ctx=32768,        # Your 32K context window sweet spot
    n_gpu_layers=99,    # Dumps 100% of the layers directly to ZeroGPU VRAM
    ctk_quant=2,        # Compresses attention Keys to 4-bit format
    ctv_quant=2,        # Compresses attention Values to 4-bit format
    flash_attn=True     
)

# 4. ZeroGPU Endpoint execution loop
@spaces.GPU(duration=60)
def api_generate(prompt):
    clean_prompt = str(prompt).strip()
    if not clean_prompt:
        return "Error: Empty prompt received."
        
    formatted_prompt = f"<|user|>\n{clean_prompt}<|endoftext|>\n<|assistant|>"
    
    response = llm(
        formatted_prompt, 
        max_tokens=1024, 
        temperature=0.7, 
        top_p=0.95
    )
    
    try:
        return response["choices"][0]["text"] # Fixed nested array dictionary index format
    except (KeyError, IndexError, TypeError):
        return str(response)

# 5. Bind the public endpoint mapping
demo = gr.Interface(
    fn=api_generate, 
    inputs=gr.Textbox(label="Input Prompt / Document Dump", lines=8), 
    outputs=gr.Textbox(label="Model Output"),
    api_name="chat" # Exposes your endpoint hook as /predict
)

if __name__ == "__main__":
    demo.launch()