from dotenv import load_dotenv
import time
import gradio as gr
from PIL import Image
from ChatCohere import chat_completion, summarize
from PokemonCards import choose_random_cards
from QdrantRag import NeuralSearcher, SemanticCache, qdrant_client, encoder, image_encoder, processor, sparse_encoder
load_dotenv()


searcher = NeuralSearcher("pokemon_texts", "pokemon_images", qdrant_client, encoder, image_encoder, processor, sparse_encoder)
semantic_cache = SemanticCache(qdrant_client, encoder, "semantic_cache", 0.75)

def chat_pokemon(message: str, history):
    answer = semantic_cache.search_cache(message)
    if answer != "":
        r = ""
        for c in answer:
            r += c
            time.sleep(0.001)
            yield r
    else:
        context_search = searcher.search_text(message)
        reranked_context = searcher.reranking(message, context_search)
        context = "\n\n-----------------\n\n".join(reranked_context)
        final_prompt = f"USER QUERY:\n\n{message.content}\n\nCONTEXT:\n\n{context}"
        response = chat_completion(final_prompt)
        semantic_cache.upload_to_cache(message, response)
        r = ""
        for c in response:
            r += c
            time.sleep(0.001)
            yield r
    
def what_pokemon(image_input):
    save_path = Image.fromarray(image_input)
    result = searcher.search_image(save_path)
    return "You Pokemon might be: " + result[0]

def card_package(n_cards:int=5):
    description, cards = choose_random_cards(n_cards)
    package = [f"![Card {i+1}]({cards[i]})" for i in range(len(cards))]
    cards_message = "\n\n".join(package)
    natural_lang_description = chat_completion(f"Can you enthusiastically describe the cards in this package?\n\n{description}")
    return "## Your package:\n\n" + cards_message + "\n\n## Description:\n\n" + summarize(natural_lang_description)


iface1 = gr.ChatInterface(fn=chat_pokemon, title="Pokemon Chatbot", description="Ask any question about Pokemon and get an answer!")
iface2 = gr.Interface(fn=what_pokemon, title="Pokemon Image Classifier", description="Upload an image of a Pokemon and get its name!", inputs="image", outputs="text")
iface3 = gr.Interface(fn=card_package, title="Pokemon Card Package", description="Get a package of random Pokemon cards!", inputs=gr.Slider(5,10,step=1, label="Number of cards"), outputs=gr.Markdown(value="Your output will be displayed here", label="Card Package"))

iface = gr.TabbedInterface([iface1, iface2, iface3], ["PokemonChat", "Identify Pokemon", "Card Package"])

iface.launch()