import gradio as gr
from langchain.prompts import PromptTemplate
from langchain_huggingface import 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,
)

template_classify = '''
What is the main topic of given text?:

<text>  
{TEXT}  
</text>

Answer generally not specifically, your response should not contain specific information from given text just general topic.
Answer shortly with two or three word phrases not with long sentence
'''

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":["General"]}}
'''

json_output_parser = JsonOutputParser()

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

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

    language_map = {"tr": "turkish",
                    "en": "english",
                    "ar": "arabic",
                    "es": "spanish",
                    "it": "italian",
                    }
    lang = language_map[lang]

    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="Topics")
        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()