"""
🦁 NumZoo — Math practice with cute AI-generated animal rewards!
"""

import base64
import io
import random

import gradio as gr

from math_engine import generate_question, LEVEL_NAMES, level_up_message
from image_generator import generate_reward_image

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

ANSWERS_PER_LEVEL = 5

ANIMAL_EMOJIS = ["🐶", "🐱", "🐰", "🦊", "🐼", "🐨", "🦁", "🐯", "🐸", "🐧", "🦋", "🦄"]
PLACE_EMOJIS  = ["🌊", "🏔️", "🌸", "🌈", "🌙", "⭐", "🌴", "🏡", "🌺", "🍄"]

MATH_LEVEL_NAMES = {1: "Additions", 2: "Subtractions", 3: "Multiplications", 4: "Mix"}

# ---------------------------------------------------------------------------
# Game helpers
# ---------------------------------------------------------------------------

def _math_level(state: dict) -> int:
    return min(state.get("level", 1), 4)

def _status_text(state: dict) -> str:
    if not state:
        return ""
    diff    = MATH_LEVEL_NAMES.get(_math_level(state), "")
    correct = state.get("correct_this_level", 0)
    return (f"👤 {state['name']}  |  ⭐ {state['score']}  |  "
            f"📚 {diff}  |  ✅ {correct}/{ANSWERS_PER_LEVEL}")

def _safe_question(state: dict) -> str:
    return f"## {state.get('question', '')} = ?"

def _img_to_data_url(pil_image) -> str:
    buf = io.BytesIO()
    pil_image.save(buf, format="JPEG", quality=80)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()

def _picker_level_label(level: int) -> str:
    diff = MATH_LEVEL_NAMES.get(min(level, 4), "Mix")
    return f"## 🌟 Level {level} — {diff}"

# ---------------------------------------------------------------------------
# Collection helpers — server-side render, no JS polling
# ---------------------------------------------------------------------------

def _add_locked(items: list, level_id: int) -> list:
    """Append a locked placeholder for level_id if not already in the list."""
    id_str = str(level_id)
    if any(x["id"] == id_str for x in items):
        return items
    return items + [{"id": id_str, "status": "locked", "src": None}]

def _unlock_item(items: list, level_id: int, src: str) -> list:
    """Replace the locked entry for level_id with an unlocked image."""
    id_str = str(level_id)
    updated, found = [], False
    for item in items:
        if item["id"] == id_str:
            updated.append({"id": id_str, "status": "unlocked", "src": src})
            found = True
        else:
            updated.append(item)
    if not found:
        updated.append({"id": id_str, "status": "unlocked", "src": src})
    return updated

def _render_collection(items: list) -> str:
    """Render the collection as HTML, newest first. Empty string when no items."""
    if not items:
        return ""
    cards = []
    for item in reversed(items):
        if item["status"] == "locked":
            cards.append(
                '<div class="nz-card nz-card-locked">'
                '<div class="nz-card-pulse"></div>'
                '<div class="nz-card-icon">🔒</div>'
                '</div>'
            )
        elif item.get("src"):
            src = item["src"]
            cards.append(
                f'<div class="nz-card">'
                f'<img src="{src}" '
                f'onclick="var m=document.getElementById(\'nz-modal\'),'
                f'i=document.getElementById(\'nz-modal-img\');'
                f'i.src=this.src;m.style.display=\'flex\';">'
                f'</div>'
            )
    if not cards:
        return ""
    return (
        '<div id="nz-coll-wrap">'
        '<div class="nz-coll-header">🖼️ Collection</div>'
        '<div id="nz-coll-grid">'
        + "".join(cards)
        + '</div></div>'
    )

# ---------------------------------------------------------------------------
# Step 1 — Name entry
# ---------------------------------------------------------------------------

def enter_name(player_name: str, state: dict):
    try:
        name  = player_name.strip() or "Player"
        state = {"name": name, "level": 1, "score": 0, "streak": 0}
        return (
            state,
            gr.update(visible=False),   # welcome
            gr.update(visible=True),    # emoji picker
            gr.update(visible=False),   # game
            gr.update(visible=False),   # picker_level_md
            random.sample(ANIMAL_EMOJIS, 3),
            random.sample(PLACE_EMOJIS, 3),
        )
    except Exception as e:
        print(f"enter_name error: {e}")
        return (state, gr.update(visible=True), gr.update(visible=False),
                gr.update(visible=False), gr.update(visible=False), [], [])

# ---------------------------------------------------------------------------
# Step 2 — Emoji selection → start level
# ---------------------------------------------------------------------------

def start_level(animal_sel: list, place_sel: list, state: dict, coll_items: list):
    try:
        if not animal_sel or not place_sel:
            return (state, gr.update(visible=True), gr.update(visible=False),
                    "⚠️ Pick at least one animal and one place!", "", "", gr.update(), gr.update())

        level = state.get("level", 1)
        state.update({
            "selected_animals":   animal_sel[:3],
            "selected_places":    place_sel[:3],
            "correct_this_level": 0,
            "generate_now":       False,
        })

        math_lv          = _math_level(state)
        question, answer = generate_question(math_lv)
        state["question"] = question
        state["answer"]   = answer

        new_coll = _add_locked(coll_items, level)
        return (state,
                gr.update(visible=False),           # emoji_panel
                gr.update(visible=True),            # game_panel
                "",                                 # picker_error
                _status_text(state),                # status_md
                _safe_question(state),              # question_md
                new_coll,                           # coll_items state
                _render_collection(new_coll))       # collection_html

    except Exception as e:
        print(f"start_level error: {e}")
        return (state, gr.update(visible=True), gr.update(visible=False),
                "⚠️ Something went wrong, try again.", "", "", gr.update(), gr.update())

# ---------------------------------------------------------------------------
# Background pre-generation — fires after "Let's go!" so the image is ready
# when the user finishes the level.
# ---------------------------------------------------------------------------

def pregenerate_image(state: dict):
    level = state.get("level", 1)
    try:
        animals = state.get("selected_animals", [random.choice(ANIMAL_EMOJIS)])
        places  = state.get("selected_places",  [random.choice(PLACE_EMOJIS)])

        print(f"[pregenerate] level={level} | animals={animals} | places={places}")
        result, prompt = generate_reward_image(animals, places)
        print(f"[pregenerate] level={level} done | prompt={prompt!r}")

        if result is not None:
            data_url = _img_to_data_url(result)
            state["pending_reward_id"] = level
            return state, result, data_url

        state.pop("pending_reward_id", None)
        return state, None, ""

    except Exception as e:
        import traceback
        print(f"[pregenerate] ❌ {e}\n{traceback.format_exc()}")
        state.pop("pending_reward_id", None)
        return state, None, ""

# ---------------------------------------------------------------------------
# Step 3 — Answer checking
# ---------------------------------------------------------------------------

def check_answer(user_input: str, state: dict, pre_image, pre_data_url: str, coll_items: list):
    _no_change_coll = (gr.update(), gr.update())   # (coll_items, collection_html) no-ops

    def _empty():
        return (state, _status_text(state), _safe_question(state), "",
                "", gr.update(visible=False), "", gr.update(visible=False), "",
                gr.update(), gr.update(), gr.update(), gr.update())

    try:
        if not state or not state.get("question"):
            return _empty()

        try:
            user_answer = int(user_input.strip())
        except (ValueError, AttributeError):
            return (state, _status_text(state), _safe_question(state), "",
                    "⚠️ Numbers only!", gr.update(visible=False), "", gr.update(visible=False), "",
                    gr.update(), gr.update(), gr.update(), gr.update())

        correct = (user_answer == state["answer"])

        if correct:
            state["score"]             += 1
            state["streak"]            += 1
            state["correct_this_level"] = state.get("correct_this_level", 0) + 1

            level         = state.get("level", 1)
            correct_count = state["correct_this_level"]
            remaining     = ANSWERS_PER_LEVEL - correct_count
            streak_fire   = "🔥" * min(state["streak"], 5)

            if correct_count >= ANSWERS_PER_LEVEL:
                # ── Level complete! ───────────────────────────────────────
                state["generate_now"] = False
                pending_id = state.get("pending_reward_id")
                has_pre    = (pre_image is not None
                              and bool(pre_data_url)
                              and pending_id == level)

                if has_pre:
                    state.pop("pending_reward_id", None)
                    new_coll = _unlock_item(coll_items, level, pre_data_url)
                    return (state, _status_text(state), _safe_question(state), "",
                            f"🎉 Level {level} complete! {streak_fire}",
                            gr.update(visible=True),                    # reward_panel
                            "",                                         # loader cleared
                            gr.update(visible=True, value=pre_image),   # image shown
                            "",
                            None,                                       # clear hidden_image
                            "",                                         # clear hidden_data_url
                            new_coll,
                            _render_collection(new_coll))
                else:
                    state["generate_now"] = True
                    return (state, _status_text(state), _safe_question(state), "",
                            f"🎉 Level {level} complete! {streak_fire}",
                            gr.update(visible=True),    # reward_panel
                            LOADER_HTML,                # show loader
                            gr.update(visible=False),   # image hidden
                            "",
                            gr.update(), gr.update(),
                            gr.update(), gr.update())   # collection unchanged until image ready
            else:
                # ── Keep going ────────────────────────────────────────────
                math_lv          = _math_level(state)
                question, answer = generate_question(math_lv)
                state["question"] = question
                state["answer"]   = answer
                return (state, _status_text(state), _safe_question(state), "",
                        f"✅ {remaining} to go! {streak_fire}",
                        gr.update(visible=False), "", gr.update(visible=False), "",
                        gr.update(), gr.update(), gr.update(), gr.update())

        else:
            state["streak"]  = 0
            math_lv          = _math_level(state)
            question, answer = generate_question(math_lv)
            state["question"] = question
            state["answer"]   = answer
            return (state, _status_text(state), _safe_question(state), "",
                    f"❌ Answer: **{state['answer']}**",
                    gr.update(visible=False), "", gr.update(visible=False), "",
                    gr.update(), gr.update(), gr.update(), gr.update())

    except Exception as e:
        import traceback
        print(f"check_answer error: {e}\n{traceback.format_exc()}")
        state["generate_now"] = False
        math_lv          = _math_level(state)
        question, answer = generate_question(math_lv)
        state["question"] = question
        state["answer"]   = answer
        return (state, _status_text(state), _safe_question(state), "",
                "⚠️ Something went wrong!",
                gr.update(visible=False), "", gr.update(visible=False), "",
                gr.update(), gr.update(), gr.update(), gr.update())

# ---------------------------------------------------------------------------
# On-demand generation — fallback when pre-image wasn't ready at level end
# ---------------------------------------------------------------------------

def generate_on_demand(state: dict, coll_items: list):
    """Generate reward image on demand. No-op if generate_now is False."""
    if not state.get("generate_now"):
        return state, "", gr.update(), "", gr.update(), gr.update()

    state["generate_now"] = False
    level = state.get("level", 1)

    try:
        animals = state.get("selected_animals", [random.choice(ANIMAL_EMOJIS)])
        places  = state.get("selected_places",  [random.choice(PLACE_EMOJIS)])

        print(f"[on_demand] level={level}")
        result, prompt = generate_reward_image(animals, places)
        print(f"[on_demand] level={level} done")

        if result is not None:
            data_url = _img_to_data_url(result)
            new_coll = _unlock_item(coll_items, level, data_url)
            return (state, "",
                    gr.update(visible=True, value=result),
                    "",
                    new_coll,
                    _render_collection(new_coll))

        return state, "", gr.update(visible=False), "⚠️ Could not generate image — try again later!", gr.update(), gr.update()

    except Exception as e:
        import traceback
        print(f"[on_demand] ❌ {e}\n{traceback.format_exc()}")
        return state, "", gr.update(visible=False), f"⚠️ Error: {e}", gr.update(), gr.update()

# ---------------------------------------------------------------------------
# Next level — back to picker with incremented level
# ---------------------------------------------------------------------------

def next_level(state: dict):
    try:
        new_level = state.get("level", 1) + 1
        state["level"] = new_level
        state["correct_this_level"] = 0
        state.pop("pending_reward_id", None)

        label = _picker_level_label(new_level)
        return (state,
                gr.update(visible=False),               # game_panel
                gr.update(visible=True),                # emoji_panel
                gr.update(value=label, visible=True),   # picker_level_md
                random.sample(ANIMAL_EMOJIS, 2),
                random.sample(PLACE_EMOJIS, 2),
                None,                                   # clear hidden_image
                "")                                     # clear hidden_data_url
    except Exception as e:
        print(f"next_level error: {e}")
        return (state, gr.update(), gr.update(), gr.update(),
                gr.update(), gr.update(), None, "")

# ---------------------------------------------------------------------------
# Restart
# ---------------------------------------------------------------------------

def restart(state: dict):
    return ({}, gr.update(visible=True), gr.update(visible=False),
            gr.update(visible=False), gr.update(), "", None, "", [], "")

# ---------------------------------------------------------------------------
# Gradio UI
# ---------------------------------------------------------------------------

CSS = """
.gradio-container { max-width: 560px !important; margin: 0 auto !important; }
.gradio-group > .form { padding: 20px 24px !important; }
.prose h3 { margin-top: 14px !important; margin-bottom: 8px !important; padding-left: 4px !important; }
#title        { text-align: center; font-size: 2.2em; margin-bottom: 0.1em; }
#subtitle     { text-align: center; color: #888; margin-bottom: 1em; font-style: italic; }
#question-box { text-align: center; font-size: 2.6em; font-weight: bold; padding: 0.5em; }
#feedback-box { text-align: center; font-size: 1.3em; min-height: 2em; }
#status-box   { text-align: center; padding: 0.4em; border-radius: 8px; }
#reward-img   { border-radius: 16px; }
#picker-level { text-align: center; padding: 0.3em 0 0.6em; }
.answer-input input { font-size: 2em !important; text-align: center !important; }
/* Emoji picker grid */
.emoji-group .wrap {
  display: grid !important; gap: 8px !important;
  justify-content: center !important; justify-items: center !important;
}
#animal-picker .wrap { grid-template-columns: repeat(6, 52px) !important; }
#place-picker  .wrap { grid-template-columns: repeat(5, 52px) !important; }
.emoji-group label {
  width: 52px !important; height: 52px !important;
  display: flex !important; align-items: center !important; justify-content: center !important;
  font-size: 1.8em !important; cursor: pointer !important;
  border-radius: 12px !important; border: 2px solid transparent !important;
  transition: all 0.15s !important; user-select: none !important;
}
.emoji-group label:has(input:checked) {
  background: #e9d5ff !important; border-color: #7c3aed !important;
}
.emoji-group input[type="checkbox"] { display: none !important; }
/* Collection */
@keyframes numzoo-pulse { 0%,100% { opacity:.4; } 50% { opacity:.9; } }
#nz-coll-wrap   { margin-top: 16px; }
.nz-coll-header { font-weight: bold; font-size: 1.05em; margin: 0 0 8px 4px; }
#nz-coll-grid   { display: flex; flex-wrap: wrap; gap: 8px; justify-content: center; }
.nz-card        { width: 80px; height: 80px; border-radius: 12px; overflow: hidden;
                  flex-shrink: 0; position: relative; }
.nz-card-locked { background: #1e1e2e; border: 2px dashed #555; }
.nz-card-pulse  { position: absolute; inset: 0;
                  background: linear-gradient(135deg,#2d2d44,#111128);
                  animation: numzoo-pulse 2s ease-in-out infinite; }
.nz-card-icon   { width: 100%; height: 100%; display: flex; align-items: center;
                  justify-content: center; font-size: 2em; z-index: 1; position: relative; }
.nz-card img    { width: 100%; height: 100%; object-fit: cover; cursor: pointer; display: block; }
"""

LOADER_HTML = """
<div style="text-align:center; padding:2em;">
  <div style="font-size:3em; animation:spin 1s linear infinite; display:inline-block;">🎨</div>
  <div style="margin-top:0.5em; color:#888; font-size:1.1em;">Generating your reward…</div>
</div>
<style>@keyframes spin{0%{transform:rotate(0deg)}100%{transform:rotate(360deg)}}</style>
"""

ZOOM_MODAL_HTML = """
<div id="nz-modal" onclick="this.style.display='none'"
     style="display:none; position:fixed; inset:0; background:rgba(0,0,0,0.88);
            z-index:9999; align-items:center; justify-content:center; cursor:pointer;">
  <img id="nz-modal-img"
       style="max-width:92vw; max-height:92vh; border-radius:20px; box-shadow:0 8px 40px #0006;">
</div>
"""

with gr.Blocks(title="🦁 NumZoo") as demo:

    state           = gr.State({})
    coll_items      = gr.State([])
    hidden_image    = gr.Image(visible=False, label="", type="pil")
    hidden_data_url = gr.Textbox(visible=False, value="")

    gr.Markdown("# 🦁 NumZoo", elem_id="title")
    gr.Markdown("*Do maths. Win cute animals!*", elem_id="subtitle")

    gr.HTML("""
<script>
  document.addEventListener('DOMContentLoaded', function() {
    var saved = localStorage.getItem('numzoo_name');
    if (saved) {
      setTimeout(function() {
        var input = document.querySelector('input[placeholder="Your name…"]');
        if (input) {
          Object.getOwnPropertyDescriptor(HTMLInputElement.prototype,'value').set.call(input, saved);
          input.dispatchEvent(new Event('input', {bubbles:true}));
        }
      }, 800);
    }
  });
  document.addEventListener('input', function(e) {
    if (e.target.placeholder === 'Your name…') localStorage.setItem('numzoo_name', e.target.value);
  });
</script>
""")

    # ── Welcome ────────────────────────────────────────────────────────────
    with gr.Group(visible=True) as welcome_panel:
        gr.Markdown("### Who are you?")
        player_name = gr.Textbox(placeholder="Your name…", label=" ", show_label=False, max_lines=1)
        next_btn    = gr.Button("Next ➡️", variant="primary", size="lg")

    # ── Emoji picker ───────────────────────────────────────────────────────
    with gr.Group(visible=False) as emoji_panel:
        picker_level_md = gr.Markdown("", visible=False, elem_id="picker-level")
        gr.Markdown("### Pick your animals 🐾")
        animal_picker   = gr.CheckboxGroup(choices=ANIMAL_EMOJIS, label="",
                                           show_label=False, elem_classes=["emoji-group"],
                                           elem_id="animal-picker")
        gr.Markdown("### Pick your places 🌍")
        place_picker    = gr.CheckboxGroup(choices=PLACE_EMOJIS, label="",
                                           show_label=False, elem_classes=["emoji-group"],
                                           elem_id="place-picker")
        picker_error    = gr.Markdown("")
        go_btn          = gr.Button("Let's go! 🚀", variant="primary", size="lg")

    # ── Game ───────────────────────────────────────────────────────────────
    with gr.Group(visible=False) as game_panel:
        status_md    = gr.Markdown("", elem_id="status-box", container=True)
        question_md  = gr.Markdown("", elem_id="question-box")
        answer_input = gr.Textbox(placeholder="?", label="Your answer",
                                  max_lines=1, elem_classes=["answer-input"])
        check_btn    = gr.Button("✔️ Check!", variant="primary")
        feedback_md  = gr.Markdown("", elem_id="feedback-box")

        with gr.Group(visible=False) as reward_panel:
            gr.Markdown("### 🎁 Your reward!")
            loader_html  = gr.HTML("")
            reward_image = gr.Image(label="", show_label=False,
                                    elem_id="reward-img", height=400, visible=False)
            reward_error = gr.Markdown("")
            next_level_btn = gr.Button("Next Level ➡️", variant="primary", size="lg")

    # ── Collection (server-side rendered) ─────────────────────────────────
    collection_html = gr.HTML("", elem_id="nz-collection-html")
    gr.HTML(ZOOM_MODAL_HTML)

    restart_btn = gr.Button("🔄 Restart", variant="secondary", size="sm")

    # ── Wiring ─────────────────────────────────────────────────────────────

    enter_name_outputs = [state, welcome_panel, emoji_panel, game_panel,
                          picker_level_md, animal_picker, place_picker]
    next_btn.click(enter_name, [player_name, state], enter_name_outputs)
    player_name.submit(enter_name, [player_name, state], enter_name_outputs)

    start_level_outputs = [state, emoji_panel, game_panel, picker_error,
                           status_md, question_md, coll_items, collection_html]
    go_btn.click(
        start_level, [animal_picker, place_picker, state, coll_items], start_level_outputs
    ).then(
        pregenerate_image, [state], [state, hidden_image, hidden_data_url]
    )

    check_outputs = [
        state, status_md, question_md, answer_input,
        feedback_md, reward_panel, loader_html, reward_image, reward_error,
        hidden_image, hidden_data_url, coll_items, collection_html,
    ]
    ondemand_outputs = [state, loader_html, reward_image, reward_error,
                        coll_items, collection_html]

    for trigger in [check_btn.click, answer_input.submit]:
        trigger(
            check_answer,
            [answer_input, state, hidden_image, hidden_data_url, coll_items],
            check_outputs,
        ).then(
            generate_on_demand, [state, coll_items], ondemand_outputs
        )

    next_level_outputs = [state, game_panel, emoji_panel, picker_level_md,
                          animal_picker, place_picker, hidden_image, hidden_data_url]
    next_level_btn.click(next_level, [state], next_level_outputs)

    restart_btn.click(
        restart, [state],
        [state, welcome_panel, emoji_panel, game_panel,
         player_name, picker_error, hidden_image, hidden_data_url, coll_items, collection_html]
    )


if __name__ == "__main__":
    demo.launch(
        css=CSS,
        theme=gr.themes.Soft(primary_hue="purple", secondary_hue="pink"),
    )