import gradio as gr
import pandas as pd
import random
from datasets import load_dataset

# Load dataset
dataset = load_dataset("tejasashinde/a2z-multidomain-glossary", split="train")
df = pd.DataFrame(dataset)

# Prepare dataset
df = (
    df.dropna(subset=["word", "description"])
    .sample(frac=1)
    .reset_index(drop=True)
)

# Game settings
GAME_LENGTHS = {
    "Quick Game (10 Qs)": 10,
    "Medium Game (20 Qs)": 20,
    "Endless Mode": None,
}

# Session state
session_state = {
    "score": 0,
    "question_idx": 0,
    "game_length": 10,
    "used_indices": [],
    "current_correct_word": "",
    "current_options": [],
    "mode": None,
}

def generate_question():
    while True:
        idx = random.randint(0, len(df) - 1)
        if idx not in session_state["used_indices"]:
            session_state["used_indices"].append(idx)
            break
    row = df.loc[idx]
    correct = row["word"]
    description = row["description"]
    domain = row["domain"]
    prev = df[(df["domain"] == domain) & (df["word"] != correct)]
    distractors = (
        prev["word"].sample(3).tolist()
        if len(prev) >= 3
        else df[df["word"] != correct]["word"].sample(3).tolist()
    )
    options = distractors + [correct]
    random.shuffle(options)
    session_state.update(current_correct_word=correct, current_options=options)
    return description, options

def show_mode_selection():
    return (
        gr.update(visible=False),  # start_btn
        gr.update(visible=True),   # mode_selector
        gr.update(visible=True),   # proceed_btn
        gr.update(visible=False),  # level_text
        gr.update(value="", visible=False),  # question_display
        *(gr.update(visible=False),)*4,  # buttons b0,b1,b2,b3
        "",  # feedback cleared
        "",  # summary cleared
        gr.update(visible=False),  # play_again_btn
        gr.update(visible=False),  # end_game_btn
    )

def start_game(mode):
    session_state.update(
        score=0,
        question_idx=0,
        used_indices=[],
        game_length=GAME_LENGTHS[mode],
        mode=mode,
    )
    desc, options = generate_question()
    return (
        gr.update(visible=False),  # mode_selector
        gr.update(visible=False),  # proceed_btn
        gr.update(value=f"Question 1", visible=True),  # level_text
        gr.update(value=desc, visible=True),  # question_display
        *[gr.update(value=opt, visible=True) for opt in options],  # buttons
        "",  # feedback cleared
        "",  # summary cleared
        gr.update(visible=False),  # play_again_btn
        gr.update(visible=(mode == "Endless Mode")),  # end_game_btn only visible in endless mode
    )

def answer_question(choice):
    session_state["question_idx"] += 1
    correct = session_state["current_correct_word"]

    if choice == correct:
        session_state["score"] += 1
        feedback = '<div class="alert alert-success">✅ Correct!</div>'
    else:
        feedback = f'<div class="alert alert-danger">❌ Wrong! Correct was <b>{correct}</b></div>'

    # Check if game ended (for finite mode)
    if session_state["game_length"] and session_state["question_idx"] >= session_state["game_length"]:
        summary_html = f'''
        <div class="game-over-box">
            <h2>🏁 Game Over</h2>
            <p>You scored <strong>{session_state["score"]}</strong> out of <strong>{session_state["game_length"]}</strong></p>
        </div>
        '''
        return (
            gr.update(value=feedback, visible=True),   # Show last feedback
            gr.update(value="", visible=False),        # Hide question_display
            *(gr.update(visible=False),)*4,            # Hide buttons
            gr.update(value=summary_html, visible=True),  # Show summary
            gr.update(visible=False),                   # Hide level_text
            gr.update(visible=True),                    # Show play_again_btn
            gr.update(visible=False),                   # Hide end_game_btn
        )

    # For ongoing game (including endless mode)
    desc, options = generate_question()
    return (
        gr.update(value=feedback, visible=True),  # Show feedback
        gr.update(value=desc, visible=True),
        *[gr.update(value=o, visible=True) for o in options],
        gr.update(value="", visible=False),        # Clear summary
        f"Question {session_state['question_idx'] + 1}",
        gr.update(visible=False),                   # Hide play_again_btn during game
        gr.update(visible=(session_state["mode"] == "Endless Mode")),  # Show end_game_btn only in endless mode
    )

def end_game():
    summary_html = f'''
    <div class="game-over-box">
        <h2>🏁 Game Ended</h2>
        <p>You scored <strong>{session_state["score"]}</strong> points.</p>
    </div>
    '''
    return (
        gr.update(value="", visible=False),        # Clear feedback
        gr.update(value="", visible=False),        # Hide question_display
        *(gr.update(visible=False),)*4,            # Hide buttons
        gr.update(value=summary_html, visible=True),
        gr.update(visible=False),                   # Hide level_text
        gr.update(visible=True),                    # Show play_again_btn
        gr.update(visible=False),                   # Hide end_game_btn
    )

# Styling (unchanged)
custom_css = """
.alert { padding:10px; border-radius:4px; margin-top:10px; font-weight:bold; }
.alert-success { background: #d4edda; color: #155724; }
.alert-danger { background: #f8d7da; color: #721c24; }
.alert b { color: black; }

body {
    background: radial-gradient(ellipse at center, #001f3f 0%, #000 100%);
    color: white; font-family: 'Open Sans', sans-serif;
    margin:0; padding:0;
    height: 100vh;
    display: flex;
    justify-content: center;
    align-items: center;
}

.tv-screen {
    position: relative;
    display: flex;
    flex-direction: column;
    background: #000;
    border: 20px solid #444;
    border-radius: 30px;
    width: 650px;
    height: 600px;
    padding: 25px;
    box-shadow: inset 0 0 30px #0ff, 0 0 20px #0ff;
    overflow: hidden;
}

.gr-button.option {
    margin: 10px 5px;
    background:#003366;
    color:white;
    border-radius:8px;
    min-width: 140px;
    font-weight: bold;
    font-size: 16px;
}
.gr-button.option:hover { background:#0052cc; }

.question-header {
    font-size: 28px;
    font-weight: bold;
    background: rgba(0, 0, 0, 0.5);
    padding: 20px;
    text-shadow: 1px 1px 5px gold;
    border: 2px solid #ffaa00;
    border-radius: 10px;
    text-align: center;
    margin-bottom: 5px;
}

.level-tracker {
    text-align: center;
    color: gold;
    font-weight: bold;
    font-size: 18px;
    margin-bottom: 8px;
}

.game-over-box {
    top: 0; left: 0; right: 0; bottom: 0;
    background-color: rgba(255, 255, 255, 0.95);
    color: #000;
    border-radius: 30px;
    padding: 40px;
    display: flex;
    flex-direction: column;
    justify-content: center;
    align-items: center;
    font-family: 'Open Sans', sans-serif;
    font-size: 20px;
    text-align: center;
    box-shadow: 0 4px 20px rgba(0, 0, 0, 0.5);
    z-index: 10;
}
.game-over-box h2 {
    font-size: 28px;
    margin-bottom: 10px;
    color: black;
}
.game-over-box p,strong {
    font-size: 18px;
    margin-bottom: 20px;
    color: black;
}
"""

with gr.Blocks(css=custom_css, title="Word Quiz") as demo:
    with gr.Column(elem_classes="tv-screen"):
        gr.Markdown("## 🧐 Guess the Word!")
        gr.HTML('<p style="color: white;">This space uses <a href="https://huggingface.co/datasets/tejasashinde/a2z-multidomain-glossary" target="_blank">tejasashinde/a2z-multidomain-glossary</a> dataset.</p>')

        start_btn = gr.Button("🎮 Start Game")
        mode_selector = gr.Radio(
            choices=list(GAME_LENGTHS.keys()), label="Choose Mode", visible=False
        )
        proceed_btn = gr.Button("Proceed", visible=False)

        level_text = gr.Markdown("Question 1", elem_classes="level-tracker", visible=False)
        question_display = gr.Markdown("", elem_classes="question-header", visible=False)

        with gr.Row():
            b0 = gr.Button("", elem_classes="option", visible=False)
            b1 = gr.Button("", elem_classes="option", visible=False)
        with gr.Row():
            b2 = gr.Button("", elem_classes="option", visible=False)
            b3 = gr.Button("", elem_classes="option", visible=False)

        feedback = gr.Markdown()
        summary = gr.HTML(visible=False)
        play_again_btn = gr.Button("🔁 Play Again", visible=False)
        end_game_btn = gr.Button("⏹️ End Game", visible=False)

        # Logic connections
        start_btn.click(
            fn=show_mode_selection,
            inputs=None,
            outputs=[
                start_btn,
                mode_selector,
                proceed_btn,
                level_text,
                question_display,
                b0,
                b1,
                b2,
                b3,
                feedback,
                summary,
                play_again_btn,
                end_game_btn,
            ],
        )

        proceed_btn.click(
            fn=start_game,
            inputs=mode_selector,
            outputs=[
                mode_selector,
                proceed_btn,
                level_text,
                question_display,
                b0,
                b1,
                b2,
                b3,
                feedback,
                summary,
                play_again_btn,
                end_game_btn,
            ],
        )

        for btn in [b0, b1, b2, b3]:
            btn.click(
                fn=answer_question,
                inputs=btn,
                outputs=[
                    feedback,
                    question_display,
                    b0,
                    b1,
                    b2,
                    b3,
                    summary,
                    level_text,
                    play_again_btn,
                    end_game_btn,
                ],
            )

        play_again_btn.click(
            fn=show_mode_selection,
            inputs=None,
            outputs=[
                start_btn,
                mode_selector,
                proceed_btn,
                level_text,
                question_display,
                b0,
                b1,
                b2,
                b3,
                feedback,
                summary,
                play_again_btn,
                end_game_btn,
            ],
        )

        end_game_btn.click(
            fn=end_game,
            inputs=None,
            outputs=[
                feedback,
                question_display,
                b0,
                b1,
                b2,
                b3,
                summary,
                level_text,
                play_again_btn,
                end_game_btn,
            ],
        )

demo.launch()