import gradio as gr
from langchain.prompts import PromptTemplate
from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
from langchain_core.output_parsers import JsonOutputParser
from langdetect import detect
import time

# Initialize the LLM and other components
llm = HuggingFaceEndpoint(
    repo_id="mistralai/Mistral-7B-Instruct-v0.3",
    task="text-generation",
    max_new_tokens=4096,
    temperature=0.5,
    do_sample=False,
)
llm_engine_hf = ChatHuggingFace(llm=llm)

template_classify = '''
Please carefully read the following text. The text is written in {LANG} language:

<text>  
{TEXT}  
</text>

After reading it, I want you to extract topic informations from text.
'''

template_json = '''
Your task is to read the following text, convert it to json format using 'Answer' as key and return it.
<text>
{RESPONSE}
</text>

Your final response MUST contain only the response, no other text.
Example:
{{"Answer":"Positive"}}
'''
json_output_parser = JsonOutputParser()

# Define the classify_text function
def classify_text(text):
    global llm

    start = time.time()
    lang = detect(text)

    prompt_classify = PromptTemplate(
        template=template_classify,
        input_variables=["LANG", "TEXT"]
    )
    formatted_prompt = prompt_classify.format(TEXT=text, LANG=lang)
    classify = llm.invoke(formatted_prompt)

    prompt_json = PromptTemplate(
        template=template_json,
        input_variables=["RESPONSE"]
    )

    formatted_prompt = template_json.format(RESPONSE=classify)
    response = llm.invoke(formatted_prompt)

    parsed_output = json_output_parser.parse(response)
    end = time.time()
    duration = end - start
    return parsed_output, duration #['Answer']

# Create the Gradio interface
def gradio_app(text):
    classification, time_taken = classify_text(text)
    return classification, f"Time taken: {time_taken:.2f} seconds"

def create_gradio_interface():
    with gr.Blocks() as iface:
        text_input = gr.Textbox(label="Text")
        output_text = gr.Textbox(label="Classify")
        time_taken = gr.Textbox(label="Time Taken (seconds)")
        submit_btn = gr.Button("Classify")

        submit_btn.click(fn=classify_text, inputs=text_input, outputs=[output_text, time_taken])

    iface.launch()

if __name__ == "__main__":
    create_gradio_interface()