import gradio as gr
from huggingface_hub import InferenceClient
import google.generativeai as genai
import numpy as np
import PIL.Image
import io, os
"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
genai.configure(api_key='AIzaSyAtdSFdZ2WZv2TmJYijVz286JBX7HpddQk')


def respond(
    message,
    history: list[tuple[str, str]],
    image,
):
    messages = ""

    for val in history:
        if val[0]:
            messages += "{'role': 'user', 'content': "+val[0]+"}"
        if val[1]:
            messages += "{'role': 'user', 'content': "+val[1]+"}"

    messages += "{'role': 'user', 'content': "+message+"}"
    print (messages,image)
    ## for image
    if isinstance(image, np.ndarray):
        img = PIL.Image.fromarray(image)
    else:
        try:
            img = PIL.Image.open(image)
        except (AttributeError, IOError) as e:
            return f"Invalid image provided. Please provide a valid image file. Error: {e}"
    
    # Load model
    model = genai.GenerativeModel("gemini-pro-vision")
    ## for image
    response = model.generate_content([messages, img])
    try:
        response = model.generate_content([messages, img])
        if not response or not response.text:
            return "No valid response received. The response might have been blocked."
        
        # Formatting the response
        formatted_response = ""
        for line in response.text.split("\n"):
            if line.strip().endswith(":"):
                formatted_response += f"**{line.strip()}**\n"
            else:
                formatted_response += line + "\n"
        
        return formatted_response
    except ValueError as e:
        return f"Error in generating response: {e}"
    # response = ""

    # for message in client.chat_completion(
    #     messages,
    #     max_tokens=max_tokens,
    #     stream=True,
    #     temperature=temperature,
    #     top_p=top_p,
    # ):
    #     token = message.choices[0].delta.content

    #     response += token
    #     yield response

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Image(show_label=False)
    ],
    additional_inputs_accordion=gr.Accordion(open=True),
    theme=gr.themes.Soft(),
    retry_btn=None,
    undo_btn=None,
    clear_btn=None;
    css="""#component-2{
    flex-direction:row !important;
}
#component-3{
    min-height: 100% !important;
    width: 50% !important;
}
#component-6{
    width: 50% !important;
    position: absolute !important;
    left: 50% !important;
    z-index: 100 !important;
    bottom: 45% !important;
}
#component-14{
    width: 50% !important;
    height: 50% !important;
}"""
)


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