import gradio as gr
from PIL import Image
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 client is None:
        try:
            initialize_client(together_api_key)
        except Exception as e:
            history.append((message, f"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}]

    # Add history to messages
    for user_msg, assistant_msg in history:
        if isinstance(user_msg, str):  # Text message
            messages.append(
                {"role": "user", "content": [{"type": "text", "text": user_msg}]}
            )
        elif isinstance(user_msg, dict):  # Image message
            image_base64 = encode_image(user_msg["image_path"])
            messages.append(
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": user_msg.get("text", "")},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{image_base64}"
                            },
                        },
                    ],
                }
            )
        messages.append(
            {"role": "assistant", "content": [{"type": "text", "text": assistant_msg}]}
        )

    # Prepare the current message
    user_message_content = []
    if isinstance(message, dict):
        if message.get("text"):
            user_message_content.append({"type": "text", "text": message["text"]})
        if message.get("files") and len(message["files"]) > 0:
            image_path = message["files"][0]
            image_base64 = encode_image(image_path)
            user_message_content.append(
                {
                    "type": "image_url",
                    "image_url": {"url": f"data:image/png;base64,{image_base64}"},
                }
            )
    elif isinstance(message, str):
        user_message_content.append({"type": "text", "text": message})

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

    # Add the user's message to the history
    user_display_message = message["text"] if isinstance(message, dict) else message
    history = history + [(user_display_message, "")]

    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:
            if (
                chunk.choices
                and chunk.choices[0].delta
                and chunk.choices[0].delta.content is not None
            ):
                response += chunk.choices[0].delta.content
                # Update the assistant's response in the history
                history[-1] = (user_display_message, response)
                yield history

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

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


# The rest of your Gradio interface code remains the same
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=10000,
            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="")
    clear = gr.Button("Clear")

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

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