import os

import deepl
import gradio as gr

from main import retrieve, generate, State

deepl_client = deepl.DeepLClient(os.getenv("DEEPL_AUTH_KEY"))

LANGUAGES = {"中文": "ZH-Hans", "英语": "EN-US"}


def query_model(question: str, target_lang: str) -> str:
    trans_question = str(deepl_client.translate_text(question, target_lang="EN-US"))
    state = State(question=trans_question, context=[], answer="")
    print(f"Question: {state['question']}\n")

    retrieved_state = retrieve(state)
    state.update(retrieved_state)
    print(f"Retrieved {len(state['context'])} documents\n")
    for doc in state["context"]:
        print(doc.page_content, "\n")

    generated_state = generate(state)
    state.update(generated_state)

    if target_lang == "EN-US":
        return state["answer"]
    else:
        return str(deepl_client.translate_text(state["answer"], target_lang=target_lang))


iface = gr.Interface(fn=query_model,
                     inputs=[gr.Textbox(label="输入你的问题", lines=2, placeholder="我想学机器学习，应该学哪些课程？"),
                             gr.Dropdown(choices=[(k, v) for k, v in LANGUAGES.items()], label="小克的回复语言...",
                                         allow_custom_value=False, value=LANGUAGES["英语"]), ],
                     outputs=gr.Textbox(label="小克的回答", lines=9), title="RIC课程小助手",
                     description="问一些关于课程的问题，然后等待小克的回答！", flagging_mode="never")

iface.launch(share=False)