import gradio as gr
import random
from transformers import pipeline
import pathlib

model = pipeline(model="declare-lab/flan-alpaca-large")


class Game:
    def __init__(self):
        self.words = pathlib.Path('solutions.txt').read_text().splitlines()
        self.word_list = random.sample(self.words, 5)
        self.secret_word = random.choice(self.word_list)

    def reset(self):
        self.word_list = random.sample(self.words, 5)
        self.secret_word = random.choice(self.word_list)

    @property
    def prompt(self):
        return f"Try to guess the word! Either enter the word or ask a hint. The word will be one of {self.word_list}"


with gr.Blocks(theme='gstaff/xkcd') as demo:
    game = Game()
    title = gr.Markdown("# Guessing Game")
    chatbot = gr.Chatbot(value=[(None, game.prompt)])
    msg = gr.Textbox()
    restart = gr.Button("Restart")

    def user(user_message, history):
        return "", history + [[user_message, None]]

    def bot(history):
        user_input = history[-1][0]
        if game.secret_word == user_input.strip().lower():
            history[-1][1] = f"You win, the word was {game.secret_word}!"
            return history
        if user_input.strip().lower() in game.word_list:
            history[-1][1] = "Wrong guess, try again."
            return history
        instructions = f"The secret word is {game.secret_word}. Answer this: {user_input}"
        bot_message = model(instructions, max_length=256, do_sample=True)[0]['generated_text']
        response = bot_message.replace(game.secret_word, "?????").replace(game.secret_word.title(), "?????")
        history[-1][1] = response
        return history

    def restart_game():
        game.reset()
        return [(None, game.prompt)]

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    restart.click(restart_game, None, chatbot, queue=False)

demo.launch()