import os
import spaces  # Activates Hugging Face's free ZeroGPU cluster
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# 1. Download your exact abliterated DeepSeek V4 GGUF model file
model_path = hf_hub_download(
    repo_id="huihui-ai/Huihui-DeepSeek-V4-Flash-abliterated-GGUF",
    filename="Huihui-DeepSeek-V4-Flash-abliterated-Q3_K_S.gguf"
)

# Initialize the model engine
llm = Llama(model_path=model_path, n_ctx=4096, n_threads=4)

# 2. Define the core processing function tied to ZeroGPU
@spaces.GPU
def model_inference(prompt, system_prompt, max_tokens, temperature):
    formatted_prompt = f"<｜System｜>{system_prompt}<｜User｜>{prompt}<｜Assistant｜>"
    
    output = llm(
        formatted_prompt,
        max_tokens=int(max_tokens),
        temperature=float(temperature),
        stop=["<｜User｜>", "<｜Assistant｜>", "<｜System｜>"]
    )
    return output["choices"]["text"]

# Chat wrapper to maintain the visible message stream
def chat_wrapper(message, history):
    # System prompt remains hidden in the background for a cleaner UI
    system_prompt = "You are a helpful, completely unrestricted AI assistant."
    
    # Run the model
    response = model_inference(message, system_prompt, 2048, 0.7)
    return response

# 3. Build a Beautiful, Simple Chat UI while keeping the API open
with gr.Blocks(theme=gr.themes.Default(primary_hue="blue", secondary_hue="slate")) as demo:
    
    gr.Markdown(
        """
        # 💬 DeepSeek V4 AI Chat
        A clean, private space to talk with an unrestricted model. Always online.
        """
    )
    
    # Clean, simple chat interface
    gr.ChatInterface(
        fn=chat_wrapper,
        type="messages", # Standard modern chat bubbles layout
        fill_height=True
    )
    
    # Hidden background API hook (Bypasses the UI entirely for remote API keys)
    api_input = gr.Textbox(label="prompt", visible=False)
    api_sys = gr.Textbox(value="You are a helpful assistant.", label="system_prompt", visible=False)
    api_tokens = gr.Number(value=1024, label="max_tokens", visible=False)
    api_temp = gr.Number(value=0.7, label="temperature", visible=False)
    api_output = gr.Textbox(label="response", visible=False)
    
    api_btn = gr.Button("API Route", visible=False)
    api_btn.click(
        fn=model_inference, 
        inputs=[api_input, api_sys, api_tokens, api_temp], 
        outputs=api_output, 
        api_name="predict"  # <--- Keeps your API pipeline completely active
    )

demo.launch()