import gradio as gr
from transformers import pipeline

# Initialize the text generation pipeline
pipe = pipeline("text-generation", model="akhaliq/MyGemmaGradioCoder")

def generate_code(user_input):
    """
    Generate code based on user input using the Gemma model
    """
    messages = [
        {"role": "user", "content": user_input},
    ]
    
    # Generate response from the model
    response = pipe(messages, max_new_tokens=512, temperature=0.7, do_sample=True)
    
    # Extract the generated text from the response
    generated_text = response[0]['generated_text']
    
    # If the response contains the full conversation, extract just the assistant's response
    if isinstance(generated_text, list):
        # Handle conversation format
        for msg in generated_text:
            if msg.get('role') == 'assistant':
                return msg.get('content', '')
        # If no assistant message found, return the last message content
        return generated_text[-1].get('content', '') if generated_text else ""
    else:
        # Handle string format - try to extract the code after the user input
        if user_input in generated_text:
            return generated_text.split(user_input)[-1].strip()
        return generated_text

# Create the Gradio interface
with gr.Blocks(title="Text to Code Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🚀 Text to Code Generator
        
        Generate code from natural language descriptions using the Gemma Gradio Coder model.
        Simply describe what you want to build, and the AI will generate the corresponding code!
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input section
            input_text = gr.Textbox(
                label="Describe what you want to code",
                placeholder="e.g., Create a Python function that calculates the factorial of a number",
                lines=5,
                max_lines=10
            )
            
            with gr.Row():
                generate_btn = gr.Button("Generate Code", variant="primary", scale=2)
                clear_btn = gr.ClearButton([input_text], value="Clear", scale=1)
            
            # Examples section
            gr.Examples(
                examples=[
                    ["Create a Python function to check if a number is prime"],
                    ["Write a JavaScript function to reverse a string"],
                    ["Create a React component for a todo list item"],
                    ["Write a SQL query to find the top 5 customers by total purchase amount"],
                    ["Create a Python class for a bank account with deposit and withdraw methods"],
                ],
                inputs=input_text,
                label="Example Prompts"
            )
        
        with gr.Column(scale=1):
            # Output section
            output_code = gr.Code(
                label="Generated Code",
                language="python",
                lines=20,
                interactive=True,
                show_line_numbers=True,
                wrap_lines=True,
                autocomplete=True
            )
            
            with gr.Row():
                copy_btn = gr.Button("📋 Copy Code", scale=1)
            
    # Add event handlers
    generate_btn.click(
        fn=generate_code,
        inputs=input_text,
        outputs=output_code,
        api_name="generate"
    )
    
    input_text.submit(
        fn=generate_code,
        inputs=input_text,
        outputs=output_code
    )
    
    # Copy functionality (using JavaScript)
    copy_btn.click(
        fn=None,
        inputs=output_code,
        outputs=None,
        js="""
        (code) => {
            navigator.clipboard.writeText(code);
            alert('Code copied to clipboard!');
            return null;
        }
        """
    )
    
    # Footer
    gr.Markdown(
        """
        ---
        💡 **Tips:**
        - Be specific about the programming language you want
        - Include details about inputs, outputs, and edge cases
        - You can edit the generated code directly in the output box
        
        Powered by [akhaliq/MyGemmaGradioCoder](https://huggingface.co/akhaliq/MyGemmaGradioCoder)
        """
    )

# Launch the app
if __name__ == "__main__":
    demo.launch()