import gradio as gr
import os
from together import Together
import base64
from io import BytesIO

# 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_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 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
        messages = [
            {"role": "system", "content": "You are an AI assistant that can analyze wireframe images and generate React code based on their content."},
            {"role": "user", "content": [
                {"type": "image_url", "image_url": f"data:image/png;base64,{img_str}"},
                {"type": "text", "text": "Analyze this wireframe image and generate React code with Tailwind CSS classes that could recreate the main elements seen in the image."}
            ]}
        ]
        
        # Make the API call
        response = client.chat.completions.create(
            model="meta-llama/Llama-2-70b-chat",
            messages=messages,
            max_tokens=1024,
            temperature=0.7,
            top_p=0.7,
            top_k=50,
            repetition_penalty=1,
            stop=["<|im_end|>"]
        )
        
        # Extract the generated code from the response
        generated_code = response.choices[0].message.content
        
        return generated_code
    except Exception as e:
        return f"An error occurred: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown("# Turn your wireframe into an app")
    gr.Markdown("Upload an image of your website design and we'll build it for you with React + Tailwind.")
    
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Upload a screenshot", elem_id="image_upload")
            example_link = gr.Markdown("Need an example image? [Try ours](https://example.com/wireframe.png).")
            
            model_dropdown = gr.Dropdown(
                choices=["Llama 2 70B Vision"],
                value="Llama 2 70B Vision",
                label="AI Model"
            )
            
            generate_button = gr.Button("Generate app", variant="primary")
    
    code_output = gr.Code(language="javascript", label="Generated React Code", lines=10)
    
    generate_button.click(
        fn=generate_app,
        inputs=[image_input],
        outputs=[code_output]
    )

demo.launch()