import gradio as gr
import pandas as pd
import random
import json

with open('pokemon.json', 'r') as f:
    pokemons = json.load(f)
pokemons_types = ["모든 타입"] + sorted(set([t for poke in pokemons for t in poke['types']]))
df = pd.DataFrame(pokemons)
GEN_RANGE = {
    "모든 세대": [1, 1017],
    "1세대": [1, 151],
    "2세대": [152, 251],
    "3세대": [252, 386],
    "4세대": [387, 493],
    "5세대": [494, 649],
    "6세대": [650, 721],
    "7세대": [722, 809],
    "8세대": [810, 905],
    "9세대": [906, 1017]
}

QUESTION_TEMPLATE = {"question": "다음 포켓몬의 이름은 뭘까요?![]({img_url})", "answer": "{name}"}

def get_question_answer(pokemons_set):
    chosen = random.choice(pokemons_set)
    name = chosen['name']
    image_path = chosen['image_path']
    answer.value = QUESTION_TEMPLATE['answer'].format(name=name)
    img_url = f"https://huggingface.co/spaces/yoon-gu/pokemon/resolve/main/{image_path}"
    q = QUESTION_TEMPLATE["question"].format(img_url=img_url)
    a = QUESTION_TEMPLATE['answer'].format(name=name)
    return q, a

MD = """# 포켓몬 퀴즈
"""
with gr.Blocks() as demo:
    quiz_start = gr.State(value=False)
    score = gr.State(value=0)
    answer = gr.State(value="")
    with gr.Row():
        with gr.Column():
            gr.Markdown(MD)
        with gr.Column():
            with gr.Row():
                generation = gr.Dropdown(
                            [f"{k}세대" for k in range(1, 10)] + ["모든 세대"], value="1세대", label="포켓몬 세대", info="원하는 포켓몬 세대를 선택하세요."
                        )
                poke_types = gr.Dropdown(
                            pokemons_types, value="모든 타입", label="포켓몬 타입", info="원하는 포켓몬 타입을 선택하세요."
                    )
            chatbot = gr.Chatbot(bubble_full_width=False)
            msg = gr.Textbox(value="퀴즈 시작")
            clear = gr.ClearButton([msg, chatbot, quiz_start])

    def respond(gen, types, message, chat_history, request: gr.Request):
        start, end = GEN_RANGE[gen]
        sdf = df[start:end]
        pokemons_set = sdf[sdf['types'].apply(lambda x: (types in x)) | (types == "모든 타입")]
        pokemons_set = pokemons_set.to_dict("records")
        if "퀴즈 시작" == message:
            quiz_start.value = True
        if not quiz_start.value:
            if "퀴즈 시작" == message:
                q, a = get_question_answer(pokemons_set)
                bot_message = f"퀴즈를 시작합니다.\n{q}"
                answer.value = a
                quiz_start.value = True
            else:
                bot_message = "퀴즈를 시작하고 싶으시면, **퀴즈 시작**이라고 말씀해주세요."
        else:
            if answer.value == message:
                q, a = get_question_answer(pokemons_set)
                answer.value = a
                bot_message = f"🎉정답입니다! 다음 문제입니다.\n{q}"
            else:
                bot_message = f"***{message}***!? 🧐 다시 한번 생각해보세요."

        chat_history.append((message, bot_message))
        return "", chat_history
    
    msg.submit(respond, [generation, poke_types, msg, chatbot], [msg, chatbot])

demo.queue(concurrency_count=3)
demo.launch()