import urllib.request
import gradio as gr
from huggingface_hub import get_token

from chatbot import get_retrieval_qa
from flagging import HuggingFaceDatasetSaver


# get the html data and save it to a file
def download_html(_url: str, _filename: str):
    html = urllib.request.urlopen(_url).read()
    with open(_filename, "wb") as f:
        f.write(html)


url = "https://sea.ai/faq"
filename = "FAQ_SEA.AI.html"
download_html(url, filename)

# load the retrieval QA model
qa = get_retrieval_qa(filename)

# dataset callback
dataset_name = "SEA-AI/seadog-chat-history"
hf_writer = HuggingFaceDatasetSaver(get_token(), dataset_name)


def answer_question(message, history, system):
    print(f"{message=}, {history=}, {system=}")
    # concatenate the history, message and system
    query = " ".join([message, system])
    retrieval_qa = qa.invoke(query)
    result = retrieval_qa["result"]  # "query" and "source_documents" are also available
    result = result.replace('"', "").strip()  # clean up the result

    # save the query and result to the dataset
    hf_writer.flag(flag_data=[query, [dict(role="assistant", content=result)]])
    return result


title = "✨ SEA Dog"
description = """
<p align="center">
DISCLAIMERS
<br>
I can't remember conversations yet, be patient with me.
<br>
Your queries will be saved to
<a href='https://huggingface.co/datasets/SEA-AI/seadog-chat-history'>this dataset</a>
for analytics purposes.
</p>
"""

css = """
h1 {
    text-align: center;
    display: block;
}
"""

theme = gr.themes.Default(primary_hue=gr.themes.colors.indigo)

chatbot = gr.Chatbot(
    value=[
        gr.ChatMessage(
            role="assistant",
            content="I have memorized the entire SEA.AI FAQ page. Ask me anything about it! 🧠",
        ),
    ],
    label="SEA Dog",
    type="messages",
    show_label=False,
    show_copy_button=True,
)


def on_like(evt: gr.LikeData):
    print(f"{evt.index=}, {evt.value=}, {evt.liked=}")


with gr.ChatInterface(
    answer_question,
    type=chatbot.type,
    chatbot=chatbot,
    title=title,
    description=description,
    additional_inputs=[gr.Textbox("", label="SYSTEM")],
    # examples=[
    #     ["Can SEA.AI see at night?", "You are a helpful assistant."],
    #     ["Can SEA.AI see at night?", "Reply with sailor slang."],
    # ],
    cache_examples=False,
    submit_btn=None,
    css=css,
    theme=theme,
) as demo:
    # on page load, download the html and save it to a file
    demo.load(lambda: download_html(url, filename))
    # This needs to be called prior to the first call to callback.flag()
    hf_writer.setup([demo.textbox, demo.chatbot], "flagged")
    print("flagging setup done!")
    # like callback
    chatbot.like(on_like, None, None)


if __name__ == "__main__":
    demo.launch()