import gradio as gr
from PIL import Image
import requests
import os
from together import Together
import base64
import io

# Initialize Together client
client = None

def initialize_client(api_key=None):
    global client
    if api_key:
        os.environ["TOGETHER_API_KEY"] = api_key
    if "TOGETHER_API_KEY" in os.environ:
        client = Together()
    else:
        raise ValueError("Please provide a Together API Key")

def encode_image(image_path):
    with Image.open(image_path) as img:
        buffered = io.BytesIO()
        img.save(buffered, format="PNG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')

def bot_streaming(message, history, together_api_key, max_new_tokens=250, temperature=0.7):
    if history is None:
        history = []
    if client is None:
        try:
            initialize_client(together_api_key)
        except Exception as e:
            # Initialize history with error message
            history.append(["Error initializing client", str(e)])
            yield history
            return

    prompt = "You are a helpful AI assistant. Analyze the image provided (if any) and respond to the user's query or comment."

    messages = [{"role": "system", "content": prompt}]

    # Build the conversation history
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": [{"type": "text", "text": user_msg}]})
        messages.append({"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]})

    # Prepare the current message
    content = []
    user_text = ""

    try:
        if isinstance(message, dict):
            if 'text' in message and message['text']:
                user_text = message['text']
                content.append({"type": "text", "text": user_text})
            if 'files' in message and len(message['files']) > 0:
                file_info = message['files'][0]
                if isinstance(file_info, dict) and 'name' in file_info:
                    image_path = file_info['name']
                elif isinstance(file_info, str):
                    image_path = file_info
                else:
                    raise ValueError("Invalid file information.")

                image_base64 = encode_image(image_path)
                content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}})
                user_text += "\n[User uploaded an image]"
        else:
            user_text = message
            content.append({"type": "text", "text": user_text})
    except Exception as e:
        # Update history before yielding
        history.append([user_text, f"An error occurred while processing your input: {str(e)}"])
        yield history
        return

    messages.append({"role": "user", "content": content})

    # Update the history with the new user message (with empty assistant response)
    history.append([user_text, ""])
    yield history

    try:
        stream = client.chat.completions.create(
            model="meta-llama/Llama-Vision-Free",
            messages=messages,
            max_tokens=max_new_tokens,
            temperature=temperature,
            stream=True,
        )

        response = ""
        for chunk in stream:
            response += chunk.choices[0].delta.content or ""
            history[-1][1] = response
            yield history

        if not response:
            history[-1][1] = "No response generated. Please try again."
            yield history

    except Exception as e:
        if "Request Entity Too Large" in str(e):
            history[-1][1] = "The image is too large. Please try with a smaller image or compress the existing one."
        else:
            history[-1][1] = f"An error occurred: {str(e)}"
        yield history

with gr.Blocks() as demo:
    gr.Markdown("# Meta Llama-3.2-11B-Vision-Instruct (FREE)")
    gr.Markdown("Try the new Llama 3.2 11B Vision API by Meta for free through Together AI. Upload an image, and start chatting about it. Just paste in your Together AI API key and get started!")
    
    with gr.Row():
        together_api_key = gr.Textbox(
            label="Together API Key",
            placeholder="Enter your TOGETHER_API_KEY here",
            type="password"
        )
    
    with gr.Row():
        max_new_tokens = gr.Slider(
            minimum=10,
            maximum=500,
            value=250,
            step=10,
            label="Maximum number of new tokens",
        )
        temperature = gr.Number(
            value=0.7,
            minimum=0,
            maximum=1,
            step=0.1,
            label="Temperature"
        )
    
    chatbot = gr.Chatbot()
    msg = gr.MultimodalTextbox(label="Enter text or upload an image")
    clear = gr.Button("Clear")

    msg.submit(
        bot_streaming, 
        inputs=[msg, chatbot, together_api_key, max_new_tokens, temperature], 
        outputs=chatbot
    )
    clear.click(lambda: [], None, chatbot, queue=False)

if __name__ == "__main__":
    demo.launch(debug=True)