import gradio as gr
import spaces
import torch
from transformers import pipeline

# 1. Load the smarter 1.5B model
model_id = "Qwen/Qwen2.5-1.5B-Instruct"
pipe = pipeline(
    "text-generation", 
    model=model_id, 
    dtype=torch.bfloat16, 
    device_map="auto"
)

# 2. Local generation on GPU
@spaces.GPU
def generate_text(prompt, temperature, max_tokens):
    try:
        messages = [{"role": "user", "content": prompt}]
        
        outputs = pipe(
            messages, 
            max_new_tokens=int(max_tokens), 
            temperature=float(temperature), 
            do_sample=True if temperature > 0 else False
        )
        
        return outputs[0]["generated_text"][-1]["content"]
               
    except Exception as e:
        return f"An error occurred: {str(e)}"

# 3. Custom CSS to inject modern rounded corners and gradient accent borders
custom_css = """
.header-box {
    text-align: center;
    padding: 20px;
    margin-bottom: 20px;
    border-radius: 12px;
    background: linear-gradient(135deg, #4f46e5, #06b6d4);
    color: white;
}
.header-box h1 {
    color: white !important;
    margin: 0;
    font-weight: 800;
}
.header-box p {
    color: #e0f2fe !important;
    margin-top: 5px;
}
.generate-btn {
    background: linear-gradient(135deg, #4f46e5, #06b6d4) !important;
    color: white !important;
    border: none !important;
    font-weight: bold !important;
}
"""

# 4. Applying a soft, highly polished UI Theme
theme = gr.themes.Soft(
    primary_hue="indigo",
    secondary_hue="cyan",
    neutral_hue="slate"
)

# 5. Build the styled UI Blocks
with gr.Blocks(theme=theme, css=custom_css) as demo:
    
    # Custom Gradient Header (Created safely using a custom div structure inside gr.HTML)
    gr.HTML(
        """
        <div class="header-box">
            <h1>🚀 My Generative AI Playground</h1>
            <p>Powered by Qwen-2.5-1.5B and running locally on ZeroGPU</p>
        </div>
        """
    )

    with gr.Row():
        # Left Side: Inputs & Interactive Settings
        with gr.Column(scale=4):
            input_text = gr.Textbox(
                placeholder="Type your creative ideas, math problems, or questions here...", 
                label="💡 What would you like to ask?",
                lines=5
            )
            
            # Clickable prompt helper buttons
            gr.Markdown("### 🌟 Try these quick prompts:")
            with gr.Row():
                btn_story = gr.Button("📝 Write a Sci-Fi Story", size="sm")
                btn_code = gr.Button("🐍 Python Palindrome Code", size="sm")
                btn_news = gr.Button("📰 Trending News Now", size="sm")
            
            # Interactive parameters under a collapsible menu
            with gr.Accordion("⚙️ Advanced AI Settings", open=False):
                temp_slider = gr.Slider(
                    minimum=0.0, 
                    maximum=1.2, 
                    value=0.7, 
                    step=0.1, 
                    label="Creativity (Temperature)", 
                    info="Higher values mean more creative, lower means focused."
                )
                token_slider = gr.Slider(
                    minimum=50, 
                    maximum=1000, 
                    value=300, 
                    step=50, 
                    label="Response Length (Max Tokens)"
                )
            
            submit_btn = gr.Button("⚡ Generate Response", elem_classes="generate-btn")
            
        # Right Side: Clean output box
        with gr.Column(scale=5):
            output_text = gr.Textbox(
                label="✨ AI Response", 
                interactive=False, 
                lines=15, 
                placeholder="The AI's masterpiece will generate here..."
            )
            
    # Connect UI Click Elements
    submit_btn.click(
        fn=generate_text, 
        inputs=[input_text, temp_slider, token_slider], 
        outputs=output_text
    )
    
    # Hook up the quick prompt buttons to auto-fill the input box!
    btn_story.click(lambda: "Write a short, engaging story about a lost astronaut who finds a cosmic garden floating in space.", outputs=input_text)
    btn_code.click(lambda: "Write a complete Python function that checks if a string is a palindrome. Explain how it works step-by-step.", outputs=input_text)
    btn_news.click(lambda: "In clear, simple terms, explain the trending news happening now.", outputs=input_text)

# Launch the interactive application
demo.launch()