import os
import gradio as gr
import torch
from transformers import pipeline, BitsAndBytesConfig

# 1. Setup Auth
token = os.environ.get("HF_TOKEN")

# 2. Configuration for 4-bit Quantization (Saves RAM)
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16
)

print("Loading 3 models into RAM... please wait about 5 minutes.")

# Load Gemma 3 1B (Uses ~1GB RAM)
pipe_1b = pipeline("text-generation", model="google/gemma-3-1b-it", 
                   device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config})

# Load Gemma 3 4B (Uses ~3.5GB RAM in 4-bit)
pipe_4b = pipeline("text-generation", model="google/gemma-3-4b-it", 
                   device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config})

# Load Phi-4 Mini (Uses ~2.5GB RAM in 4-bit)
pipe_phi = pipeline("text-generation", model="microsoft/Phi-4-mini-instruct", 
                    device_map="cpu", token=token, model_kwargs={"quantization_config": bnb_config})

def battle(message):
    results = []
    # Format message for Instruct models
    msgs = [{"role": "user", "content": message}]
    
    # Generate from all three
    out_1b = pipe_1b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
    out_4b = pipe_4b(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
    out_phi = pipe_phi(msgs, max_new_tokens=128)[0]['generated_text'][-1]['content']
    
    return out_1b, out_4b, out_phi

# 3. UI Layout
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ⚔️ Triple Model Local Arena")
    gr.Markdown("Running locally on 16GB CPU RAM via 4-bit Quantization.")
    
    with gr.Row():
        txt = gr.Textbox(label="Prompt", placeholder="Compare their logic...", scale=4)
        btn = gr.Button("Battle!", variant="primary", scale=1)
        
    with gr.Row():
        with gr.Column():
            gr.Markdown("### ⚡ Gemma 3 1B")
            res_1 = gr.Markdown()
        with gr.Column():
            gr.Markdown("### 🧠 Gemma 3 4B")
            res_2 = gr.Markdown()
        with gr.Column():
            gr.Markdown("### 🛠️ Phi-4 Mini")
            res_3 = gr.Markdown()

    btn.click(battle, inputs=txt, outputs=[res_1, res_2, res_3])

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