import re
import gradio as gr
import os
import asyncio
from transformers import pipeline
import time


class TaskClassifier:
    
    def __init__(self):
        self.classifier = pipeline("zero-shot-classification",
                      model="MoritzLaurer/mDeBERTa-v3-base-xnli-multilingual-nli-2mil7")
    
    def __call__(self, client_input: str, task_types: str):
        """Classify tasks for LLM-based gent"""
        candidate_labels = [label.strip() for label in task_types.split(",")]
        time_execution = time.time()
        output = self.classifier(str(client_input), candidate_labels, multi_label=False)
        # output = classifier(input, candidate_labels, multi_label=False)
        time_execution = round(time.time() - time_execution, 2)
        # return {"task_type": output['labels'][0], "confidence": round(output['scores'][0],2), "inference_time": time_execution}
        return f"Task Type : {output['labels'][0]}\nScore : {round(output['scores'][0],2)}\nInference Time : {time_execution}"
    



def load_classifier(client_input, task_types):
    global classifier
    return classifier(client_input, task_types)


def question_answer(client_input, task_types):
    if client_input.strip()=='':
        return '''[ERROR]: Please enter client input (e.g., 'Find the top products for a given category').'''
    if task_types.strip() == '':
        return '''[ERROR]: Please enter list of task type of LLM-based agents (e.g., 'Greeting, Information retrieval, Sentiment analysis, Text generation, Code generation, Q&A, Summarize'): '''
    return load_classifier(client_input, task_types)


classifier = TaskClassifier()

title = 'Task Clarity for LLM-based Agents'
description = """ Task Clarity for LLM-based Agents is a powerful tool that assists developers in crafting precise task instructions, identifies task types (e.g., Q&A, Text generation) for your LLM-based Agents."""


with gr.Blocks() as demo:

    gr.Markdown(f'<center><h1>{title}</h1></center>')
    gr.Markdown(description)

    with gr.Row():
        
        with gr.Group():
            gr.Markdown(f'<p style="text-align:center">Report about the model: <a href="https://sinh-nguyen.notion.site/Report-Solving-Task-Clarity-for-LLM-based-Agents-4b49b5229a3f423984743b11f3c2bec8">here</a></p>')
            client_input=gr.Textbox(label='''Please enter client's input (e.g., 'Hello?'): ''')
            task_types = gr.Textbox(label='''Please enter list of task type of LLM-based agents (e.g., 'Greeting, Information retrieval, Sentiment analysis, Text generation, Code generation, Q&A, Summarization'): ''')
            btn = gr.Button(value='Submit')
            btn.style(full_width=True)
#openai.api_key = os.getenv('Your_Key_Here') 
        with gr.Group():
            answer = gr.Textbox(label='The answer to your question is :')

        btn.click(question_answer, inputs=[client_input, task_types], outputs=[answer])

demo.launch(share=True)
# demo.launch(server_name="0.0.0.0", server_port=7860)