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):
    try:
        with Image.open(image_path) as img:
            buffered = io.BytesIO()
            img.save(buffered, format="PNG")
            return base64.b64encode(buffered.getvalue()).decode('utf-8')
    except Exception as e:
        print(f"Error encoding image: {e}")
        raise e

def bot_streaming(message, history, together_api_key, max_new_tokens=250, temperature=0.7):
    # Initialize history if it's None
    if history is None:
        history = []
    
    # Initialize the Together client if not already done
    if client is None:
        try:
            initialize_client(together_api_key)
        except Exception as e:
            # Append error to history and yield
            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 the API
    for idx, (user_msg, assistant_msg) in enumerate(history):
        # Append user messages
        messages.append({
            "role": "user",
            "content": [
                {"type": "text", "text": user_msg}
            ]
        })
        # Append assistant messages
        messages.append({
            "role": "assistant",
            "content": [
                {"type": "text", "text": assistant_msg}
            ]
        })

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

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

                # Encode the image to base64
                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:
            # If message is a string
            user_text = message
            content.append({"type": "text", "text": user_text})
    except Exception as e:
        # If there's an error processing the input, append it to history and yield
        error_message = f"An error occurred while processing your input: {str(e)}"
        print(error_message)  # Debug statement
        history.append([user_text or "[Invalid input]", error_message])
        yield history
        return

    # Append the new user message with an empty assistant response
    history.append([user_text, ""])
    yield history  # Yield the updated history to show the user's message immediately

    # Append the current user message to the API messages
    messages.append({"role": "user", "content": content})

    try:
        # Call the Together AI API with streaming
        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:
            # Extract the content from the API response
            chunk_content = chunk.choices[0].delta.content or ""
            response += chunk_content
            # Update the last assistant message in history
            if history:
                history[-1][1] = response
                yield history
            else:
                # If history is somehow empty, append the response
                history.append(["", response])
                yield history

        if not response:
            # If no response was generated, notify the user
            history[-1][1] = "No response generated. Please try again."
            yield history

    except Exception as e:
        # Handle exceptions from the API call
        error_message = ""
        if "Request Entity Too Large" in str(e):
            error_message = "The image is too large. Please try with a smaller image or compress the existing one."
        else:
            error_message = f"An error occurred: {str(e)}"
        
        print(error_message)  # Debug statement

        if history:
            history[-1][1] = error_message
        else:
            history.append(["", error_message])
        
        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)