import gradio as gr
from transformers import pipeline

# Load a pre-trained code generation model (replace with your actual model)
code_generator = pipeline("text-generation", model="gpt-3.5-turbo")  # Example model

# Function to generate code
def generate_code(prompt, file_type):
    if file_type in ["Gradio", "Vercel", "Streamlit"]:
        # Handle platform-specific code generation
        prompt_with_file_type = f"Write a configuration or setup code for {file_type} to: {prompt}"
    else:
        # Handle programming language code generation
        prompt_with_file_type = f"Write a {file_type} code for: {prompt}"
    
    # Generate code using the AI model
    generated_code = code_generator(prompt_with_file_type, max_length=200, num_return_sequences=1)
    return generated_code[0]['generated_text']

# Function to update code
def update_code(existing_code, update_prompt):
    # Combine the existing code and the update prompt
    prompt_with_update = f"Rewrite the following code to: {update_prompt}\n\nExisting Code:\n{existing_code}"
    
    # Generate updated code using the AI model
    updated_code = code_generator(prompt_with_update, max_length=200, num_return_sequences=1)
    return updated_code[0]['generated_text']

# Examples for the interface
examples = [
    ["Write a function to calculate factorial", "Python"],
    ["Create a simple interface for a calculator", "Gradio"],
    ["Deploy a Next.js app", "Vercel"],
    ["Create a data visualization app", "Streamlit"],
    ["Write a program to reverse a string", "JavaScript"],
    ["Create a responsive navbar", "HTML"],
]

# Gradio interface
with gr.Blocks(theme='Nymbo/Nymbo_Theme') as demo:
    gr.Markdown("# AI Code Generator with Update Feature")
    gr.Markdown("Enter a prompt and select the file type or platform to generate code. You can also update the generated code with a new prompt.")
    
    with gr.Row():
        input_prompt = gr.Textbox(label="Input Prompt", placeholder="e.g., Write a function to calculate factorial...")
        file_type = gr.Dropdown(
            label="File Type / Platform",
            choices=["Python", "JavaScript", "HTML", "CSS", "Java", "C++", "Gradio", "Vercel", "Streamlit"],
            value="Python"
        )
    
    output_code = gr.Textbox(label="Generated Code", lines=10, interactive=False)
    
    generate_button = gr.Button("Generate Code")
    generate_button.click(fn=generate_code, inputs=[input_prompt, file_type], outputs=output_code)

    # Update Code Section
    with gr.Row():
        update_prompt = gr.Textbox(label="Update Prompt", placeholder="e.g., Add error handling to the code...")
        update_button = gr.Button("Update Code")
    
    update_button.click(fn=update_code, inputs=[output_code, update_prompt], outputs=output_code)

    # Add examples
    gr.Examples(
        examples=examples,
        inputs=[input_prompt, file_type],
        outputs=output_code,
        fn=generate_code,
        cache_examples=True,  # Cache results for faster loading
        label="Click on an example to get started!"
    )

# Launch the interface
demo.launch()