import time
import os
import gradio as gr
import requests as requests


URZEDASEK_ENDPOINT = os.getenv("URZEDASEK_ENDPOINT")
ANONYMOUS_TOKEN = os.getenv("ANONYMOUS_TOKEN")


class UrzedasekInferenceClient():

    def chat_completion(
        self,
        messages,
    ):
        auth_header = {
            "Authorization": f"Token {ANONYMOUS_TOKEN}"
            }
        ret = requests.post(
            URZEDASEK_ENDPOINT,
            headers=auth_header,
            json={
                "messages": messages,
            },
        )
        return ret.json()["result"]["urzedas_answer_builder"]["replies"]


client = UrzedasekInferenceClient()


def chunker(result):
    for i in range(0, len(result), 40):
        time.sleep(0.1)
        yield result[i:i + 40]


def respond(
    message,
    history: list[tuple[str, str]],
):
    messages = []

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(messages,):
        for chunk in chunker(message):
            response += chunk
            yield response


with gr.Blocks(fill_height=True) as demo:
    gr.Markdown(
    """
# Urzedasek.pl - Wersja eksperymentalna
Czy zmagasz się z biurokracją, która zdaje się wciągać jak czarna dziura? "Urzedasek.pl" przychodzi z pomocą! Nasza aplikacja to Twój biurokratyczny GPS 📍, który poprowadzi Cię przez meandry urzędowych formularzy, decyzji i nie tylko. Od zdobycia dofinansowania 💸, przez rejestrację działalności gospodarczej 💼, aż po wyrabianie paszportów 📔 – jedyne, do czego potrzebujesz, to kilka kliknięć!
Świat, w którym zdobycie zasiłku czy odwołanie się od mandatu nie musi oznaczać straconego popołudnia i nerwów. A wszystko to osiągalne bez wychodzenia z domu 🏡!
Niech każde "muszę to załatwić" stanie się "już to załatwiłem" z "Urzedasek.pl"!
    """)

    chat = gr.ChatInterface(
        respond,
        examples=[
            ["Chcę skorzystać z dofinansowania, o jakie dofinansowanie mogę się starać?", ""],
            ["Ukradli mi tablice rejestracyjną co mogę zrobić?", ""],
            ["Jak mogę założyć działalność gospodarczą?", ""],
            ["Chce wstąpić do wojska jak to mogę zrobić?", ""],
            ["Chce uzyskać paszport?", ""],
            ["Czy mogę zdobyć zasiłek na dziecko z autyzmem bez orzeczenia?", ""],
            ["Jakie dokumenty należy załatwić po urodzeniu dziecka?", ""],
            ["Jak mogę odwołać się od mandatu?", ""],
            ["Gdzie się wyrabia świadectwo urodzenia? ", ""],

        ],
        additional_inputs=[],
        retry_btn=None,
    )


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