import gradio as gr

from src.constants import MAX_ATTEMPTS
from src.handler import all_tokens
from src.handler import handle_guess
from src.state import STATE
from src.state import tokenizer


def build_demo():
    """
    Builds and returns the gradio app interface
    """

    with gr.Blocks() as demo:
        with gr.Row():
            gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
        with gr.Row():
            gr.Markdown(
                "Can you beat language models at their own game?\n"
                "In this game you're pitted against a language model in the task of, you guessed it, laungage modelling.\n"
                "Your task is to predict the next word given the previous sequence. You will get 3 attempts to guess.\n"
                "The one with the fewest guesses for a given word gets a point."
            )
        with gr.Row():
            prompt_text = gr.Textbox(
                value=tokenizer.decode(all_tokens[: STATE.current_word_index]),
                label="Context",
                interactive=False,
            )
        with gr.Row():
            with gr.Column():
                player_points = gr.Number(label="your points", interactive=False)
            with gr.Column():
                lm_points = gr.Number(label="LM points", interactive=False)
        with gr.Row():
            with gr.Column():
                remaining_attempts = gr.Number(
                    value=MAX_ATTEMPTS,
                    label="Remaining attempts",
                    precision=0,
                    interactive=False,
                )
                current_guesses = gr.Textbox(label="Your guesses")
            with gr.Column():
                lm_guesses = gr.Textbox(label="LM guesses")

        with gr.Row():
            with gr.Column():
                guess_field = gr.Textbox(label="")
                guess_button = gr.Button(value="Guess!")

        with gr.Row():
            bottom_html = gr.HTML()

        guess_button.click(
            handle_guess,
            inputs=[
                guess_field,
            ],
            outputs=[
                prompt_text,
                player_points,
                lm_points,
                current_guesses,
                lm_guesses,
                remaining_attempts,
                guess_field,
                guess_button,
                bottom_html,
            ],
        )

    return demo


def wip_sign():
    with gr.Blocks() as demo:
        gr.Markdown("<h1><center>Can you beat a language model?</center></h1>")
        with gr.Row():
            gr.Markdown("<h1><center>⛔👷‍♂️ Work in progress, come back later </center></h1>")

    return demo


def get_demo(wip=False):
    if wip:
        return wip_sign()
    else:
        return build_demo()