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 = 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
        # Using the format: <image>{base64_image}</image>\nHuman: {prompt}\nAssistant:
        message = f"""<image>{img_str}</image>
Human: 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.
Assistant: Certainly! I'll analyze the wireframe image and suggest a simple Gradio app layout based on what I see. Let me describe the main elements and how they could be implemented using Gradio components."""

        # Make the API call
        response = client.chat.completions.create(
            model="meta-llama/Llama-Vision-Free",
            messages=[{"role": "user", "content": message}],
            max_tokens=512,
            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_text = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                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()