from pathlib import Path
import gradio as gr
from src.rag import Rag
from src.model_huggingface import HuggingFaceModel

from src.amodel import AModel
AModel.load_env_variables()

STORE_DIR = "./db/rag_app" # Le répertoire de la base
# STORE_DIR = None # Store éphémère
MAX_DOCS = 6 # Le nombre max de documents dans la base

# global State https://www.gradio.app/guides/state-in-blocks
rag:Rag = Rag(
    HuggingFaceModel("meta-llama/Meta-Llama-3-8B-Instruct", None, 0),
    HuggingFaceModel(None, "sentence-transformers/all-MiniLM-l6-v2", 0),
    STORE_DIR
)
rag.reset_store() # Reset de la base à chaque démarrage du serveur
print("rag created, store reseted")

# UI
with gr.Blocks(title="RAGnar", 
            #    theme="Yntec/HaleyCH_Theme_Orange_Green",
               fill_height=True,
               analytics_enabled=False,
               css="footer {visibility: hidden}",
               ) as demo:
    def upload_file(file_path):
        name:str = Path(file_path).name
        names  = rag.emb_store.get_collection_names()
        count = len(names)
        if name in names:
            rag.delete_collection(name) 
            print("collection", name, "deleted because already exists") 
        names  = rag.emb_store.get_collection_names()          
        if count >= MAX_DOCS:
            print("collection", names[0], "deleted because too many collections") 
            rag.delete_collection(names[0])
        rag.add_pdf_to_store(file_name=file_path, collection_name=name)
        return gr.Dropdown(
                choices=rag.emb_store.get_collection_names(),
                value=rag.emb_store.collections[-1].name,
                show_label=False,
                container=False,
                interactive=True
            ) 

    def ask_rag(question:str, col_name:str):
        if col_name == "Aucun fichier":
            return "Aucun pdf actif, veuillez en uploader un !"
        if question.strip() == "":
            return "Veuillez poser une question."
        names  = rag.emb_store.get_collection_names()
        if not col_name in names:
            return "'{name}' n'est plus sur le serveur, veuillez le recharger".format(name=col_name)
        prompt, resp, sources, ids = rag.ask_rag(question, col_name)
        return resp

    def on_temperature_change(temp):
        rag.set_temperature(temp) 

    def on_refresh():
        print("on_refresh")
        choices=rag.emb_store.get_collection_names() if len(rag.emb_store.collections) > 0 else ["Aucun fichier"]
        value = rag.emb_store.collections[-1].name if len(rag.emb_store.collections) > 0 else "Aucun fichier"
        return gr.Dropdown(
                choices=choices,
                value=value,
                show_label=False,
                container=False,
                interactive=True
            )     
        

    with gr.Row():
        gr.Image("./files/drane.png",  show_download_button=False, 
                show_fullscreen_button=False, show_label=False, show_share_button=False,
                interactive=False, container=False)
        # https://www.svgrepo.com/svg/483648/viking-face
        gr.Image("./files/viking.png",  show_download_button=False, 
                show_fullscreen_button=False, show_label=False, show_share_button=False,
                interactive=False, container=False)       
    with gr.Tab("RAG naïf"):
        with gr.Row():
            # Le button qui permet d'updater le combo des collections
            refresh = gr.Button("Refresh", scale=1)
            # Le combo qui affiche toutes les collections du store
            choices=rag.emb_store.get_collection_names() if len(rag.emb_store.collections) > 0 else ["Aucun fichier"]
            value = rag.emb_store.collections[-1].name if len(rag.emb_store.collections) > 0 else "Aucun fichier"
            cols = gr.Dropdown(
                choices=choices,
                value=value,
                show_label=False,
                container=False,
                interactive=True,
                scale=10
            )  
            # Le button qui permet d'uploader un pdf
            upload_button = gr.UploadButton(
                "Clique pour ajouter un pdf", 
                file_types=[".pdf"], 
                file_count="single",
                scale=10)
            
        # La zone où on pose une question au RAG
        ask_input = gr.Text(placeholder="Pose une question à ton pdf", show_label=False, container=False)
        # La réponse du RAG (Markdown pour afficher les formules .tex)
        rag_output = gr.Textbox("", show_copy_button=False,
                    show_label=False,
                    container=False,
                    max_lines=15)
        
    with gr.Tab("Réglages"):
        gr.Markdown("## Modèles:")
        gr.Markdown("- " + rag.get_llm_name())
        gr.Markdown("- " + rag.get_feature_name())
        temperature_slider = gr.Slider(minimum=0,
                    maximum=1.0,
                    value=0.0,
                    step=0.1,
                    label="Température")
        
    # Réponses aux évènements
    refresh.click(fn=on_refresh, inputs=[], outputs=[cols])
    upload_button.upload(fn=upload_file, inputs=upload_button, outputs=[cols], show_progress=True)
    ask_input.submit(fn=ask_rag, inputs=[ask_input, cols], outputs=rag_output, show_progress=True)
    temperature_slider.change(fn=on_temperature_change, inputs=temperature_slider)
    demo.load(fn=on_refresh, inputs=[], outputs=[cols])


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