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

# Initialize the Together client
api_key = os.environ.get('TOGETHER_API_KEY')
client = Together(api_key=api_key)

def generate_gradio_app(image):
    if not api_key:
        return "Error: TOGETHER_API_KEY not set. Please check your API key."

    try:
        # 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
        stream = 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=True
        )
        
        # Collect the streamed response
        generated_code = ""
        for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                generated_code += chunk.choices[0].delta.content
                yield f"Generating Gradio code... (Current length: {len(generated_code)} characters)\n\n{generated_code}"
        
        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 try again or check your API key and connection."

with gr.Blocks() as demo:
    gr.Markdown("# Generate Gradio App from Wireframe")
    gr.Markdown("Upload an image of your UI design, and we'll generate Gradio code to recreate it.")
    
    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=[image_input],
        outputs=[code_output]
    )

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