import gradio as gr
from transformers import pipeline
import os

# Cache pipelines so we don’t reload every call
pipeline_cache = {}

def get_generator(model_name):
    if model_name not in pipeline_cache:
        pipeline_cache[model_name] = pipeline(
            "text-generation",
            model=model_name,
            device=0 if os.environ.get("CUDA_VISIBLE_DEVICES") else -1
        )
    return pipeline_cache[model_name]

def build_prompt(system, examples, query, style):
    # Prefix with system role
    prompt = f"SYSTEM: {system.strip()}\n\n"
    if style == "Zero-Shot":
        prompt += f"USER: {query.strip()}\nASSISTANT:"
    else:  # Few-Shot
        # Insert up to three examples
        for ex in examples:
            ex = ex.strip()
            if ex:
                prompt += f"EXAMPLE: {ex}\n"
        prompt += f"USER: {query.strip()}\nASSISTANT:"
    return prompt

def generate(system, ex1, ex2, ex3, query, model_name, style):
    # Build prompt text
    examples = [ex1, ex2, ex3]
    prompt = build_prompt(system, examples, query, style)
    gen = get_generator(model_name)
    # Generate up to 100 tokens beyond prompt
    out = gen(prompt, max_length=len(prompt.split())+100, do_sample=True)[0]["generated_text"]
    # Strip prompt prefix
    return out[len(prompt):].strip()

# Available small/medium models on HF Hub
AVAILABLE_MODELS = [
    "gpt2",
    "EleutherAI/gpt-neo-125M",
    "distilgpt2"
]

with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.Markdown("## 🎮 Prompt Engineering Playground")
    gr.Markdown(
        "Experiment with **Zero-Shot** vs **Few-Shot** prompting across different open models."
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            system_box = gr.Textbox(
                label="🖥️ System Prompt",
                placeholder="You are a helpful assistant."
            )
            ex1 = gr.Textbox(label="📖 Example 1 (few-shot only)", placeholder="…")
            ex2 = gr.Textbox(label="📖 Example 2 (few-shot only)")
            ex3 = gr.Textbox(label="📖 Example 3 (few-shot only)")
            query_box = gr.Textbox(label="🗣️ User Prompt", placeholder="Type your query here…")
            
            with gr.Row():
                model_choice = gr.Dropdown(
                    choices=AVAILABLE_MODELS,
                    value=AVAILABLE_MODELS[0],
                    label="🤖 Model"
                )
                style_choice = gr.Radio(
                    choices=["Zero-Shot", "Few-Shot"],
                    value="Zero-Shot",
                    label="🎨 Prompt Style"
                )
            
            run_btn = gr.Button("🚀 Generate")
        
        with gr.Column(scale=3):
            output_box = gr.Textbox(
                label="💬 Model Output",
                lines=15
            )
    
    run_btn.click(
        fn=generate,
        inputs=[system_box, ex1, ex2, ex3, query_box, model_choice, style_choice],
        outputs=output_box
    )

demo.launch()