import gradio as gr


from llm.qa_agent import QnAAgent
from llm.calculator_agent import CalculatorAgent
from llm.orchestrator import Orchestrator


if __name__ == "__main__":
    orchestrator = Orchestrator()
    qna_agent = QnAAgent()
    calculator_agent = CalculatorAgent()

    # question = input("Question - ")
    def get_answer(question: str) -> [str, str]:

        api_name, parameters = orchestrator.get_API_call(question)

        print(f"Using the {api_name} Agent")
        print(api_name, parameters)
        if api_name == "QnA":
            answer, wiki_page = qna_agent.get_answer(parameters)

        # elif api_name == "calculator":
        #     operand, op1, op2 = parameters.split(",")
        #     answer = calculator_agent.calculate(operand, op1, op2)

        print(answer)
        return [answer, wiki_page]

    demo = gr.Interface(
        fn=get_answer,
        inputs=gr.Textbox(
            placeholder="Enter your question...[Who won the Cricket World Cup in 2023?]"
        ),
        # outputs=[gr.Textbox(label=f'Document {i+1}') for i in range(TOP_K)],
        outputs=[gr.Textbox(label="Answer"), gr.Textbox(label="Wikipedia Page")],
        title="Real time Question Answering",
        description="Try asking questions beyond the LLM's cutoff date. Ex: Who is the Prime Minister of France? [The new Prime Minister assumed office in January 2024]"
    )

    demo.launch(share=True)