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

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

if api_key:
    try:
        client = Together(api_key=api_key)
    except Exception as e:
        print(f"Error initializing Together client: {e}")

def generate_gradio_app(image):
    if not api_key or not client:
        return "Error: TOGETHER_API_KEY not set or client initialization failed. 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 messages for the API call
        system_message = """You are an AI assistant that can analyze wireframe images and generate detailed Gradio code based on their content. Your task is to provide a complete, runnable Gradio application that recreates the UI elements seen in the wireframe."""

        user_message = f"""
        <image>
        data:image/png;base64,{img_str}
        </image>
        
        Analyze this wireframe image and generate a complete Python code using Gradio that recreates all the main elements seen in the image. Follow these guidelines:
        1. Use appropriate Gradio components that best represent each UI element in the wireframe.
        2. Include all necessary imports at the beginning of the code.
        3. Implement placeholder functions for any interactive elements (buttons, inputs, etc.).
        4. Use gr.Blocks() to create a layout that matches the wireframe as closely as possible.
        5. Add appropriate labels and descriptions for all components.
        6. Include the gr.Blocks().launch() call at the end of the code.
        7. Provide a complete, runnable Gradio application that can be executed as-is.
        8. Add comments explaining the purpose of each major section or component.

        Please generate the entire code, including all necessary parts to make it a fully functional Gradio application.
        """
        
        messages = [
            {"role": "system", "content": system_message},
            {"role": "user", "content": user_message}
        ]
        
        # Make the API call
        response = client.chat.completions.create(
            model="meta-llama/Llama-Vision-Free",
            messages=messages,
            max_tokens=2048,
            temperature=0.7,
            top_p=0.7,
            top_k=50,
            repetition_penalty=1,
            stop=["<|im_end|>"],
            stream=True
        )
        
        # Collect the streamed response
        generated_code = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                generated_code += chunk.choices[0].delta.content
        
        return generated_code
    except Exception as e:
        return f"An error occurred: {str(e)}\n\nPlease try again or check your API key and connection."

with gr.Blocks() as demo:
    gr.Markdown("# Turn your wireframe into a Gradio app")
    gr.Markdown("Upload an image of your UI design and we'll build a Gradio app for you.")
    
    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 app", 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]
    )

demo.launch()