import os
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# 1. Download the optimized 4-bit quantized Gemma 4 model directly to local cache
print("Downloading Gemma 4 E4B-it Q4_K_M...")
model_path = hf_hub_download(
    repo_id="unsloth/gemma-4-E4B-it-GGUF",
    filename="gemma-4-E4B-it-Q4_K_M.gguf"
)
print(f"Download complete! Saved to: {model_path}")
print("Loading model architecture into system RAM...")

# 2. Configure the inference engine strictly for a 2-vCPU / 16GB RAM container
llm = Llama(
    model_path=model_path,
    n_threads=2,          # Pin execution strictly to the 2 available virtual CPU cores
    n_ctx=2048,           # Cap the context window slightly to maintain lightning-fast CPU speeds
    n_batch=512,          # Balanced token chunk evaluation batching
    verbose=False
)
print("Gemma 4 is fully loaded and ready!")

def extract_text(content):
    """
    Helper function to parse Gradio 6 structured content blocks.
    Converts list-wrapped blocks into a clean plain text string for llama.cpp.
    """
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        text_parts = []
        for block in content:
            if isinstance(block, dict) and block.get("type") == "text":
                text_parts.append(block.get("text", ""))
            elif hasattr(block, "type") and getattr(block, "type") == "text":
                text_parts.append(getattr(block, "text", ""))
        return "".join(text_parts)
    if isinstance(content, dict):
        return content.get("text", "")
    return str(content)

def predict(message, history, system_prompt, temperature, top_p, enable_thinking):
    messages = []
    
    # Inject native reasoning/thinking rules into the system space if enabled
    sys_content = system_prompt if system_prompt else "You are a helpful AI assistant."
    if enable_thinking:
        sys_content = "<|think|>\n" + sys_content
        
    messages.append({"role": "system", "content": sys_content})
    
    # Process and sanitize Gradio 6 structured chat history lists into standard strings
    for msg in history:
        if isinstance(msg, dict):
            role = msg.get("role", "user")
            raw_content = msg.get("content", "")
        else:
            role = getattr(msg, "role", "user")
            raw_content = getattr(msg, "content", "")
        
        clean_content = extract_text(raw_content)
        messages.append({"role": role, "content": clean_content})
        
    # Extract the clean user prompt from the input message object
    user_content = extract_text(message)
    messages.append({"role": "user", "content": user_content})
    
    try:
        response_stream = llm.create_chat_completion(
            messages=messages,
            temperature=temperature,
            top_p=top_p,
            stream=True
        )
        
        partial_text = ""

        for chunk in response_stream:
            choices = chunk.get("choices", [])
        
            if not choices:
                continue
        
            delta = choices[0].get("delta", {})
        
            content = delta.get("content")
        
            if content:
                partial_text += content
                yield partial_text
    except Exception as e:
        yield f"An execution error occurred: {str(e)}"

# 3. Render Gradio UI layout
with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Optimized Gemma 4 E4B-it Workspace")
    gr.Markdown("Running locally on a single 2-vCPU instance using `llama.cpp` + 4-bit quantization.")
    
    with gr.Accordion("Inference Configurations", open=False):
        system_prompt = gr.Textbox(
            value="You are an expert software engineer and helpful AI assistant.", 
            label="System Prompt"
        )
        enable_thinking = gr.Checkbox(
            value=False, 
            label="Enable Deep Reasoning (Thinking Mode)", 
            info="Turn off for high-speed conversational text. Turn on for complex programming/logic."
        )
        temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature")
        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top P")

    gr.ChatInterface(
        fn=predict,
        additional_inputs=[system_prompt, temperature, top_p, enable_thinking]
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())