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

@spaces.GPU
def generate_code(prompt):
    generator = pipeline(
        'text-generation',
        model='unsloth/Qwen2.5-Coder-14B',
        dtype=torch.bfloat16,
        device_map="auto"
    )
    
    full_prompt = f"""You are an expert coding assistant. Write clean, efficient code based on the user's request.
    
User Request: {prompt}

Code:
"""
    result = generator(
        full_prompt,
        max_new_tokens=512,
        do_sample=True,
        temperature=0.2,
        top_p=0.95,
    )
    
    generated = result[0]['generated_text']
    code_start = generated.find("Code:")
    if code_start != -1:
        return generated[code_start + len("Code:"):].strip()
    return generated

# --- Improved UI ---
with gr.Blocks(title="AI Coder", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧠 My Free AI Coder")
    gr.Markdown("Powered by Qwen2.5-Coder-14B on Hugging Face ZeroGPU")

    with gr.Row():
        with gr.Column(scale=4):
            prompt = gr.Textbox(
                lines=5,
                placeholder="Ask me to write some code, e.g., 'Write a Python function to reverse a linked list'",
                label="Your Request"
            )
            submit = gr.Button("🚀 Generate", variant="primary")
            clear = gr.Button("🗑️ Clear")
        with gr.Column(scale=6):
            # Use gr.Code for syntax highlighting + copy button
            output = gr.Code(
                label="Generated Code",
                language="python",  # You can make this dynamic later
                interactive=False,
                lines=20
            )

    # Wire up the buttons
    submit.click(fn=generate_code, inputs=prompt, outputs=output)
    clear.click(fn=lambda: "", inputs=[], outputs=prompt)
    # Also clear output when clear is pressed (optional)
    clear.click(fn=lambda: "", inputs=[], outputs=output)

demo.launch()