import gradio as gr
from together import Together
import base64
from io import BytesIO
from PIL import Image
import numpy as np
import traceback
import json

def generate_gradio_app(api_key, image):
    if not api_key:
        return "Error: API key not provided. Please enter your Together API key."

    try:
        # Initialize the Together client with the provided API key
        client = Together(api_key=api_key)

        # Convert numpy array to PIL Image
        if isinstance(image, np.ndarray):
            image = Image.fromarray(image.astype('uint8'), 'RGB')
        
        # Convert the image to base64
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        # Prepare the prompt
        prompt = """You are an AI assistant specialized in UI/UX design and Gradio app development. Analyze the attached screenshot or UI mockup and generate a complete Gradio code that recreates this design. Follow these steps:

1. Describe the main elements of the UI, including layout, colors, and components.
2. Generate a complete Gradio Python code that recreates this UI as closely as possible.
3. Use appropriate Gradio components for each element in the UI.
4. Include all necessary imports at the beginning of the code.
5. Implement placeholder functions for any interactive elements (buttons, inputs, etc.).
6. Use gr.Blocks() to create a layout that matches the screenshot.
7. Add appropriate labels and descriptions for all components.
8. Include the gr.Blocks().launch() call at the end of the code.
9. Provide a complete, runnable Gradio application that can be executed as-is.
10. Add comments explaining the purpose of each major section or component.

Please generate the entire Gradio code based on the provided image."""

        # Make the API call
        response = client.chat.completions.create(
            model="meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo",
            messages=[
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": prompt},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{img_str}",
                            },
                        },
                    ],
                }
            ],
            max_tokens=4096,
            temperature=0.7,
            top_p=0.7,
            top_k=50,
            repetition_penalty=1,
            stop=["<|eot_id|>", "<|eom_id|>"],
            stream=False  # Changed to False for easier debugging
        )
        
        # Debug: Print the entire response
        print("API Response:", json.dumps(response, indent=2))

        # Extract the generated code from the response
        if response.choices and response.choices[0].message:
            generated_code = response.choices[0].message.content
        else:
            return "Error: Unexpected response structure from the API."
        
        if not generated_code:
            return "Error: No code generated from the model. Please try again."
        
        return generated_code

    except Exception as e:
        error_message = str(e)
        stack_trace = traceback.format_exc()
        return f"An error occurred: {error_message}\n\nStack trace:\n{stack_trace}\n\nPlease check your API key and try again."

with gr.Blocks() as demo:
    gr.Markdown("# Generate Gradio App from Wireframe")
    gr.Markdown("Enter your Together API key, upload an image of your UI design, and we'll generate Gradio code to recreate it.")
    
    api_key_input = gr.Textbox(label="Enter your Together API Key", type="password")
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="Upload a screenshot", elem_id="image_upload")
            generate_button = gr.Button("Generate Gradio Code", variant="primary")
        
        with gr.Column(scale=2):
            code_output = gr.Code(language="python", label="Generated Gradio Code", lines=30)
    
    generate_button.click(
        fn=generate_gradio_app,
        inputs=[api_key_input, image_input],
        outputs=[code_output]
    )

if __name__ == "__main__":
    demo.launch(debug=True)