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

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 concise Gradio code that recreates the main elements of this design. Follow these steps:

1. Briefly describe the main elements of the UI.
2. Generate a compact Gradio Python code that recreates the essential UI elements.
3. Use appropriate Gradio components for key elements in the UI.
4. Include necessary imports at the beginning of the code.
5. Implement minimal placeholder functions for interactive elements.
6. Use gr.Blocks() to create a basic layout that captures the essence of the screenshot.
7. Include the gr.Blocks().launch() call at the end of the code.
8. Provide a runnable Gradio application focusing on the most important aspects of the UI.
9. Keep the code concise, aiming for no more than 2000 tokens.

Please generate the Gradio code based on the provided image, focusing on the most crucial elements to fit within the token limit."""

        # Make the API call with the corrected message format
        response = client.chat.completions.create(
            model="meta-llama/Llama-3.2-90B-Vision-Instruct-Turbo",
            messages=[
                {
                    "role": "user",
                    "content": prompt + f"\n\n<image>\ndata:image/png;base64,{img_str}\n</image>"
                }
            ],
            max_tokens=2000,
            temperature=0.7,
            top_p=0.7,
            top_k=50,
            repetition_penalty=1,
            stop=["<|eot_id|>", "<|eom_id|>"],
            stream=False
        )
        
        # Debug: Print relevant parts of the response
        print("API Response received. Content available:", bool(response.choices))

        # Extract the generated code from the response
        if response.choices and response.choices[0].message:
            generated_code = response.choices[0].message.content
            print("Generated code length:", len(generated_code))
        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 Concise Gradio App from Wireframe")
    gr.Markdown("Enter your Together API key, upload an image of your UI design, and we'll generate a compact Gradio code to recreate its main elements.")
    
    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)