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 json

# 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 message for the API call
        message_content = f"""<image>{img_str}</image>
Analyze this wireframe image and suggest a simple Gradio app layout based on it. Describe the main elements you see and how they could be implemented using Gradio components."""

        # Prepare the API request payload
        payload = {
            "model": "meta-llama/Llama-Vision-Free",
            "messages": [{"role": "user", "content": message_content}],
            "max_tokens": 512,
            "temperature": 0.7,
            "top_p": 0.7,
            "top_k": 50,
            "repetition_penalty": 1,
            "stop": ["<|eot_id|>", "<|eom_id|>"],
            "stream_tokens": True
        }

        # Make the API call
        response = client.chat.completions.create(**payload)
        
        # Collect the streamed response
        generated_text = ""
        for chunk in response:
            if chunk.choices and chunk.choices[0].delta.content:
                generated_text += chunk.choices[0].delta.content
        
        return generated_text

    except Exception as e:
        error_message = str(e)
        return f"An error occurred: {error_message}\nPlease try again or check your API key and connection."

with gr.Blocks() as demo:
    gr.Markdown("# Analyze wireframe and suggest Gradio app layout")
    gr.Markdown("Upload an image of your UI design for analysis and suggestions.")
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(label="Upload a screenshot", elem_id="image_upload")
            generate_button = gr.Button("Analyze and Suggest", variant="primary")
        
        with gr.Column(scale=2):
            text_output = gr.Textbox(label="Analysis and Suggestions", lines=10)
    
    generate_button.click(
        fn=generate_gradio_app,
        inputs=[image_input],
        outputs=[text_output]
    )

demo.launch()