import gradio as gr
from transformers import pipeline
import os

HF_TOKEN = os.environ["HF_TOKEN"]


# Initialize the text generation pipelines
pipe = pipeline("text-generation", model="akhaliq/MyGemmaGradioCoder")
pipe2 = pipeline("text-generation", model="google/gemma-3-270m-it", token=HF_TOKEN)

def generate_code(user_input, model_choice="Model 1"):
    """
    Generate code based on user input using the selected Gemma model
    """
    messages = [
        {"role": "user", "content": user_input},
    ]
    
    # Select pipeline based on model choice
    selected_pipe = pipe if model_choice == "Model 1 (MyGemmaGradioCoder)" else pipe2
    
    # Generate response from the model
    response = selected_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

def generate_both(user_input):
    """
    Generate code from both models for comparison
    """
    output1 = generate_code(user_input, "Model 1 (MyGemmaGradioCoder)")
    output2 = generate_code(user_input, "Model 2 (gemma-3-270m-it)")
    return output1, output2

# Create the Gradio interface
with gr.Blocks(title="Text to Code Generator - Model Comparison", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🚀 Text to Code Generator - Model Comparison
        
        Compare code generation from two different Gemma models:
        - **Model 1**: akhaliq/MyGemmaGradioCoder
        - **Model 2**: google/gemma-3-270m-it
        
        Simply describe what you want to build, and see how each model responds!
        """
    )
    
    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 from Both Models", 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=2):
            # Output section - Two columns for comparison
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Model 1: MyGemmaGradioCoder")
                    output_code1 = gr.Code(
                        label="Generated Code (Model 1)",
                        language="python",
                        lines=15,
                        interactive=True,
                        show_line_numbers=True,
                        wrap_lines=True,
                        autocomplete=True
                    )
                    copy_btn1 = gr.Button("📋 Copy Code", scale=1)
                
                with gr.Column():
                    gr.Markdown("### Model 2: gemma-3-270m-it")
                    output_code2 = gr.Code(
                        label="Generated Code (Model 2)",
                        language="python",
                        lines=15,
                        interactive=True,
                        show_line_numbers=True,
                        wrap_lines=True,
                        autocomplete=True
                    )
                    copy_btn2 = gr.Button("📋 Copy Code", scale=1)
            
    # Add event handlers
    generate_btn.click(
        fn=generate_both,
        inputs=input_text,
        outputs=[output_code1, output_code2],
        api_name="generate"
    )
    
    input_text.submit(
        fn=generate_both,
        inputs=input_text,
        outputs=[output_code1, output_code2]
    )
    
    # Copy functionality for both outputs
    copy_btn1.click(
        fn=None,
        inputs=output_code1,
        outputs=None,
        js="""
        (code) => {
            navigator.clipboard.writeText(code);
            alert('Code from Model 1 copied to clipboard!');
            return null;
        }
        """
    )
    
    copy_btn2.click(
        fn=None,
        inputs=output_code2,
        outputs=None,
        js="""
        (code) => {
            navigator.clipboard.writeText(code);
            alert('Code from Model 2 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
        
        **Models:**
        - [akhaliq/MyGemmaGradioCoder](https://huggingface.co/akhaliq/MyGemmaGradioCoder)
        - [google/gemma-3-270m-it](https://huggingface.co/google/gemma-3-270m-it)
        """
    )

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