import os
import gradio as gr
from gradio.components import Textbox, Button
# from AinaTheme import theme
from urllib.error import HTTPError

from rag import RAG
from utils import setup

setup()


rag = RAG(
    hf_token=os.getenv("HF_TOKEN"),
    embeddings_model=os.getenv("EMBEDDINGS"), 
    model_name=os.getenv("MODEL"),   
    

)


def generate(prompt):
    try:
        output = rag.get_response(prompt)
        return output
    except HTTPError as err:
        if err.code == 400:
            gr.Warning(
                "The inference endpoint is only available Monday through Friday, from 08:00 to 20:00 CET."
            )
    except:
        gr.Warning(
            "Inference endpoint is not available right now. Please try again later."
        )


def submit_input(input_):
    if input_.strip() == "":
        gr.Warning("Not possible to inference an empty input")
        return None

    output = generate(input_)

    return output


def change_interactive(text):
    if len(text) == 0:
        return gr.update(interactive=True), gr.update(interactive=False)
    return gr.update(interactive=True), gr.update(interactive=True)


def clear():
    return (
        None,
        None,
    )


def gradio_app():
    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column(scale=0.1):
                gr.Image("rag_image.jpg", elem_id="flor-banner", scale=1, height=256, width=256, show_label=False, show_download_button = False, show_share_button = False)
            with gr.Column():
                gr.Markdown(
                    """# Retrieval-Augmented Generation (experimental)
                    🔍 **Retrieval-Augmented Generation** (RAG) is an AI framework for improving the quality of LLM-generated responses
                        by grounding the model on external sources of knowledge to supplement the LLM's internal representation of 
                        information. Implementing RAG in an LLM-based question answering system has two main benefits: It ensures 
                        that the model has access to the most current, reliable facts, and that users have access to the model's 
                        sources, ensuring that the information can be checked for accuracy and ultimately trusted. 
                        
                    🎯 **Purpose:** The main purpose of this RAG is answering questions related to the [AI ACT](https://artificialintelligenceact.eu/wp-content/uploads/2024/01/AI-Act-FullText.pdf). 
                        By incorporating external knowledge sources, RAG enables the LLM to provide more informed and reliable 
                        responses specifically tailored to inquiries about it. 
                    ⚠️ **Limitations**: This version is for beta testing only. The content generated by these models is unsupervised 
                        and might be wrong. Please bear this in mind when exploring this resource.                 
                    """
                )
        with gr.Row(equal_height=True):
            with gr.Column(variant="panel"):
                input_ = Textbox(
                    lines=11,
                    label="Input",
                    placeholder="e.g. What is the AI Act?",
                    # value = "Quina és la finalitat del Servei Meteorològic de Catalunya?"
                )

            with gr.Column(variant="panel"):
                output = Textbox(
                    lines=11, label="Output", interactive=False, show_copy_button=True
                )
                with gr.Row(variant="panel"):
                    clear_btn = Button(
                        "Clear",
                    )
                    submit_btn = Button("Submit", variant="primary", interactive=False)

        input_.change(
            fn=change_interactive,
            inputs=[input_],
            outputs=[clear_btn, submit_btn],
            api_name=False,
        )

        input_.change(
            fn=None,
            inputs=[input_],
            api_name=False,
            js="""(i, m) => {
            document.getElementById('inputlenght').textContent = i.length + '  '
            document.getElementById('inputlenght').style.color =  (i.length > m) ? "#ef4444" : "";
        }""",
        )

        clear_btn.click(
            fn=clear, inputs=[], outputs=[input_, output], queue=False, api_name=False
        )
        submit_btn.click(
            fn=submit_input, inputs=[input_], outputs=[output], api_name="get-results"
        )

        demo.launch(show_api=True)


if __name__ == "__main__":
    gradio_app()