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)

# Update the template to extract topic information
template_classify = '''
Please read the following text written in {LANG} language and extract the main topics discussed in it.
You can list more than one topic or topics sentence by sentence. List the topics clearly.

<text>  
{TEXT}  
</text>
'''

template_json = '''
Your task is to read the following extracted topics and convert them into JSON format using 'Topics' as the key.

<text>
{RESPONSE}
</text>

The final response should be in this format:
{{"Topics": ["Topic1", "Topic2", ...]}}
'''
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

# 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 to Classify")
        output_text = gr.Textbox(label="Extracted Topics")
        time_taken = gr.Textbox(label="Time Taken (seconds)")
        submit_btn = gr.Button("Extract Topics")

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

    iface.launch()

if __name__ == "__main__":
    create_gradio_interface()