import os

import google.generativeai as genai
import gradio as gr
from fastai.learner import load_learner
from PIL import Image

API_KEY = os.getenv("GENAI_API_KEY")

genai.configure(api_key=API_KEY)

model = genai.GenerativeModel(model_name="gemini-pro-vision")

learn = load_learner("pets.pkl")
categories = learn.vocab

prompt = (
    "You are and animal expert and a veterinarian.\n"
    "Give your expert opinion on the following questions based on"
    "the image of the animal.\n"
    "Don't give anything besides the answer."
)


def classify_image(img):
    pred, idx, probs = learn.predict(img)
    return dict(zip(categories, map(float, probs)))


def random_response(message, history, image=None):
    if image is not None:
        image = Image.fromarray(image)
        message = "Q: " + message.strip()
        history = history[-5:] if len(history) > 0 else ""
        history = "\n".join([f"Q: {i[0]}\nA:{i[1]}\n" for i in history])
        message = prompt + '\n\n' + history + '\n' + message + "\nA: "
        print(f"The new message is : \n{message}")
        return model.generate_content([message, image]).text
    else:
        return "Please provide an image."


with gr.Blocks() as demo:
    image = gr.Image()
    label = gr.Label(num_top_classes=5)
    gr.Interface(
        classify_image,
        inputs=image,
        outputs=label,
        title="Pet Classifier",
        description="Classify an image of a pet into different categories.",
    )
    gr.ChatInterface(random_response, additional_inputs=[image])

demo.launch()