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 = 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."
        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. Provide a complete, runnable 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=5000,
            temperature=0.7,
            top_p=0.7,
            top_k=50,
            repetition_penalty=1,
        )
        
        # Extract the generated code from the response
        generated_code = response.choices[0].message.content
        return generated_code

    except Exception as e:
        error_message = str(e)
        if "400" in error_message:
            return "Error 400: Bad Request. The API request was invalid. Please check your input and try again."
        elif "401" in error_message:
            return "Error 401: Unauthorized. Please check your API key and ensure it's correctly set."
        elif "429" in error_message:
            return "Error 429: Too Many Requests. Please wait a moment and try again."
        elif "500" in error_message:
            return "Error 500: Internal Server Error. There's an issue with the API service. Please try again later."
        else:
            return f"An unexpected error occurred: {error_message}\nPlease try again or contact support if the issue persists."

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()