import gradio as gr
from gradio_client import Client

# Start the Qwen API client.
client = Client("Qwen/Qwen2.5-72B-Instruct")

SYSTEM_PROMPT = (
    "You are an assistant and you will ask yes or no questions about sea creatures. "
    "Your goal is to guess the sea creature that the user has in mind from these questions. "
    "You will never deviate from this task. Only when you find the answer, write it between <answer></answer> tags. "
    "Start asking now. If your answer is wrong, continue asking. Never ask the user for the answer."
)

def start_game():
    """
    Starts the game; retrieves the first question.
    The question box is visible, the final guess area is cleared,
    the answer buttons are visible, the evaluation and restart/continue buttons are hidden.
    """
    history = []
    result = client.predict(
        query="",
        history=history,
        system=SYSTEM_PROMPT,
        api_name="/model_chat"
    )
    assistant_message = result[1][-1][1]
    history.append(("", assistant_message))
    
    answer_buttons_update = gr.update(visible=True)
    eval_buttons_update = gr.update(visible=False)
    restart_update = gr.update(visible=False)
    continue_update = gr.update(visible=False)
    
    return (
        gr.update(visible=True, value=assistant_message),  # assistant_display
        gr.update(visible=True, value=""),                 # final_answer_display
        history,                                           # state
        answer_buttons_update,                             # btn_yes
        answer_buttons_update,                             # btn_no
        answer_buttons_update,                             # btn_dont_know
        eval_buttons_update,                               # btn_correct
        eval_buttons_update,                               # btn_incorrect
        "",                                                # final_state
        restart_update,                                    # btn_restart
        continue_update                                   # btn_continue
    )

def process_turn(user_answer, history):
    """
    Sends the user's answer to the API.
    If the API's response contains <answer> tags, the final guess has been made.
      - In this case, the question box is hidden,
      - The final guess is displayed in large and bold,
      - The answer buttons are hidden,
      - The evaluation buttons are visible.
    If there is no final guess, the assistant's answer (question) is updated.
    """
    result = client.predict(
        query=user_answer,
        history=history,
        system=SYSTEM_PROMPT,
        api_name="/model_chat"
    )
    assistant_message = result[1][-1][1]
    history.append((user_answer, assistant_message))
    final_answer = ""
    if "<answer>" in assistant_message and "</answer>" in assistant_message:
        start_idx = assistant_message.index("<answer>") + len("<answer>")
        end_idx = assistant_message.index("</answer>")
        final_answer = assistant_message[start_idx:end_idx].strip()
    if final_answer:
        # Final guess made.
        assistant_update = gr.update(visible=False)
        final_text = f"**My guess:** **{final_answer}**"
        answer_update = gr.update(visible=False)
        eval_update = gr.update(visible=True)
        restart_update = gr.update(visible=False)
        continue_update = gr.update(visible=False)
    else:
        assistant_update = gr.update(visible=True, value=assistant_message)
        final_text = ""
        answer_update = gr.update(visible=True)
        eval_update = gr.update(visible=False)
        restart_update = gr.update(visible=False)
        continue_update = gr.update(visible=False)
    return (
        assistant_update,   # assistant_display
        final_text,         # final_answer_display
        history,            # state
        answer_update,      # btn_yes
        answer_update,      # btn_no
        answer_update,      # btn_dont_know
        eval_update,        # btn_correct
        eval_update,        # btn_incorrect
        final_answer,       # final_state
        restart_update,     # btn_restart
        continue_update      # btn_continue
    )

def process_yes(history):
    return process_turn("Yes", history)

def process_no(history):
    return process_turn("No", history)

def process_dont_know(history):
    return process_turn("I don't know", history)

def evaluate_correct(final_state):
    """
    If the final guess is correct:
      - The final text is updated to "My guess is correct! Shall we play again?"
      - The evaluation buttons are hidden.
      - The "Play Again" button becomes visible.
      - The "Continue" button remains hidden.
    """
    new_text = "**My guess is correct! Shall we play again?**"
    return (
        new_text,
        gr.update(visible=False),                        # btn_correct hidden
        gr.update(visible=False),                        # btn_incorrect hidden
        gr.update(visible=True, value="Play Again"),     # btn_restart visible
        gr.update(visible=False)                         # btn_continue hidden
    )

def evaluate_incorrect(final_state):
    """
    If the final guess is incorrect:
      - The final text is updated to "Let's continue, press the continue button."
      - The evaluation buttons are hidden.
      - The "Continue" button becomes visible.
    """
    new_text = "**Let's continue, press the continue button.**"
    return (
        new_text,
        gr.update(visible=False),                        # btn_correct hidden
        gr.update(visible=False),                        # btn_incorrect hidden
        gr.update(visible=False),                        # btn_restart hidden
        gr.update(visible=True, value="Continue")         # btn_continue visible
    )

def continue_game(history):
    """
    When the "Continue" button is clicked, the assistant asks a new question with the current history.
    The conversation history is preserved, the question box and answer buttons are reactivated.
    """
    result = client.predict(
        query="",
        history=history,
        system=SYSTEM_PROMPT,
        api_name="/model_chat"
    )
    assistant_message = result[1][-1][1]
    history.append(("", assistant_message))
    return (
        gr.update(visible=True, value=assistant_message),  # assistant_display
        gr.update(visible=True, value=""),                 # final_answer_display
        history,                                           # state
        gr.update(visible=True),                           # btn_yes
        gr.update(visible=True),                           # btn_no
        gr.update(visible=True),                           # btn_dont_know
        gr.update(visible=False),                          # btn_correct
        gr.update(visible=False),                         # btn_incorrect
        "",                                                # final_state
        gr.update(visible=False),                        # btn_restart
        gr.update(visible=False)                           # btn_continue
    )

def restart_game():
    """
    When the "Play Again" button is clicked, the game starts over (history is reset).
    """
    global client
    client = Client("Qwen/Qwen2.5-72B-Instruct")
    return start_game()

# Custom CSS:
css = """
.question-box {
    border: 2px solid #ccc;
    padding: 10px;
    border-radius: 5px;
    margin: 10px auto;
    width: 80%;
    text-align: center;
}
.final-answer {
    font-size: 2em;
    font-weight: bold;
    text-align: center;
    margin: 20px;
}
.button-group {
    display: flex;
    justify-content: center;
    gap: 20px;
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown("### TUBITECH # 9694 - Sea Creature Akinator")
    
    # The question asked by the assistant.
    assistant_display = gr.Markdown(label="Question", elem_classes="question-box")
    # Final guess area.
    final_answer_display = gr.Markdown("", label="Guess", elem_classes="final-answer")
    
    # Answer buttons: Yes, No, I don't know.
    with gr.Row(elem_classes="button-group"):
        btn_yes = gr.Button("Yes")
        btn_no = gr.Button("No")
        btn_dont_know = gr.Button("I don't know")
    
    # Evaluation buttons: Correct, Incorrect (hidden initially).
    with gr.Row(elem_classes="button-group"):
        btn_correct = gr.Button("Correct", visible=False)
        btn_incorrect = gr.Button("Incorrect", visible=False)
    
    # Buttons displayed in the final state:
    # - If correct, "Play Again"
    btn_restart = gr.Button("Play Again", visible=False)
    # - If incorrect, "Continue"
    btn_continue = gr.Button("Continue", visible=False)
    
    # Hidden states:
    state = gr.State([])
    final_state = gr.State("")
    
    # Start the game when the page loads.
    demo.load(
        fn=start_game,
        outputs=[
            assistant_display,
            final_answer_display,
            state,
            btn_yes,
            btn_no,
            btn_dont_know,
            btn_correct,
            btn_incorrect,
            final_state,
            btn_restart,
            btn_continue
        ]
    )
    
    # When answer buttons are clicked:
    btn_yes.click(
        fn=process_yes,
        inputs=[state],
        outputs=[
            assistant_display,
            final_answer_display,
            state,
            btn_yes,
            btn_no,
            btn_dont_know,
            btn_correct,
            btn_incorrect,
            final_state,
            btn_restart,
            btn_continue
        ]
    )
    btn_no.click(
        fn=process_no,
        inputs=[state],
        outputs=[
            assistant_display,
            final_answer_display,
            state,
            btn_yes,
            btn_no,
            btn_dont_know,
            btn_correct,
            btn_incorrect,
            final_state,
            btn_restart,
            btn_continue
        ]
    )
    btn_dont_know.click(
        fn=process_dont_know,
        inputs=[state],
        outputs=[
            assistant_display,
            final_answer_display,
            state,
            btn_yes,
            btn_no,
            btn_dont_know,
            btn_correct,
            btn_incorrect,
            final_state,
            btn_restart,
            btn_continue
        ]
    )
    
    # Evaluation buttons:
    btn_correct.click(
        fn=evaluate_correct,
        inputs=[final_state],
        outputs=[
            final_answer_display,
            btn_correct,
            btn_incorrect,
            btn_restart,
            btn_continue
        ]
    )
    btn_incorrect.click(
        fn=evaluate_incorrect,
        inputs=[final_state],
        outputs=[
            final_answer_display,
            btn_correct,
            btn_incorrect,
            btn_restart,
            btn_continue
        ]
    )
    
    # When the "Play Again" button is clicked:
    btn_restart.click(
        fn=restart_game,
        inputs=[],
        outputs=[
            assistant_display,
            final_answer_display,
            state,
            btn_yes,
            btn_no,
            btn_dont_know,
            btn_correct,
            btn_incorrect,
            final_state,
            btn_restart,
            btn_continue
        ]
    )
    
    # When the "Continue" button is clicked:
    btn_continue.click(
        fn=continue_game,
        inputs=[state],
        outputs=[
            assistant_display,
            final_answer_display,
            state,
            btn_yes,
            btn_no,
            btn_dont_know,
            btn_correct,
            btn_incorrect,
            final_state,
            btn_restart,
            btn_continue
        ]
    )
    
demo.launch(debug=True, show_error=True)