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

import base64
import io
import json
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   # correct answers needed to complete a level and earn a reward

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

# Math difficulty caps at 4 (Mix), but level number shown to user keeps increasing
MATH_LEVEL_NAMES = {1: "Additions", 2: "Subtractions", 3: "Multiplications", 4: "Mix"}

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _math_level(state: dict) -> int:
    """Return math difficulty (1–4), capped, from game level."""
    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 _coll(*actions) -> str:
    """Encode collection actions into the collection_trigger HTML component."""
    payload = json.dumps({
        "ts": random.random(),
        "actions": [{"action": a, "id": str(i), "src": s} for a, i, s in actions],
    })
    return f"<div id='numzoo-coll-data' data-payload='{payload}'></div>"

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

# ---------------------------------------------------------------------------
# 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 (hidden for level 1)
            random.sample(ANIMAL_EMOJIS, 2),
            random.sample(PLACE_EMOJIS, 2),
        )
    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):
    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!", "", "", "")

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

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

        # Add locked placeholder for this level's reward immediately
        return (state, gr.update(visible=False), gr.update(visible=True),
                "", _status_text(state), _safe_question(state),
                _coll(("add-locked", level, "")))

    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.", "", "", "")

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

def pregenerate_image(state: dict):
    """Pre-generate the reward image for the current level."""
    level = state.get("level", 1)
    try:
        animals = state.get("selected_animals", [random.choice(ANIMAL_EMOJIS)])
        places  = state.get("selected_places",  [random.choice(PLACE_EMOJIS)])
        streak  = state.get("streak", 0)

        print(f"[pregenerate] level={level} | animals={animals} | places={places}")
        result, prompt = generate_reward_image(streak, 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):
    _empty = (state, _status_text(state), _safe_question(state), "",
              "", gr.update(visible=False), "", gr.update(visible=False), "",
              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(), "")

        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)
                    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
                            _coll(("unlock", level, pre_data_url)))
                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(),
                            "")
            else:
                # ── Keep going ───────────────────────────────────────────────
                feedback = f"✅ {remaining} to go! {streak_fire}"
                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), "",
                        feedback,
                        gr.update(visible=False), "", gr.update(visible=False), "",
                        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(), "")

    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(), "")

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

def generate_on_demand(state: dict):
    """Generate reward image on demand. No-op if generate_now is False."""
    if not state.get("generate_now"):
        return state, "", 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)])
        streak  = state.get("streak", 0)

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

        if result is not None:
            data_url = _img_to_data_url(result)
            return (state, "",
                    gr.update(visible=True, value=result),
                    "",
                    _coll(("unlock", level, data_url)))

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

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

# ---------------------------------------------------------------------------
# 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),                         # reset animal picker
                random.sample(PLACE_EMOJIS, 2),                          # reset place picker
                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 locked pulse */
@keyframes numzoo-pulse { 0%,100% { opacity:.4; } 50% { opacity:.9; } }
"""

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>
"""

COLLECTION_HTML = """
<div id="numzoo-collection-wrap" style="margin-top:16px;">
  <div id="numzoo-collection-header"
       style="display:none; font-weight:bold; font-size:1.05em; margin:0 0 8px 4px;">🖼️ Collection</div>
  <div id="numzoo-collection"
       style="display:flex; flex-wrap:wrap; gap:8px; justify-content:center;"></div>
</div>

<!-- Zoom modal -->
<div id="numzoo-modal"
     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="numzoo-modal-img"
       style="max-width:92vw; max-height:92vh; border-radius:20px; box-shadow:0 8px 40px #0006;">
</div>

<script>
(function() {
  var KEY = 'numzoo_collection_v2', MAX = 30;

  function getCol() { try { return JSON.parse(localStorage.getItem(KEY)||'[]'); } catch(e) { return []; } }
  function saveCol(c) { localStorage.setItem(KEY, JSON.stringify(c.slice(-MAX))); }

  window.numzooAddLocked = function(id) {
    id = String(id);
    var col = getCol();
    if (!col.find(function(x){return x.id===id;})) {
      col.push({id:id, status:'locked', src:null});
      saveCol(col);
    }
    renderCollection();
  };

  window.numzooUnlock = function(id, src) {
    id = String(id);
    var col = getCol(), item = col.find(function(x){return x.id===id;});
    if (item) { item.status='unlocked'; item.src=src; }
    else { col.push({id:id, status:'unlocked', src:src}); }
    saveCol(col);
    renderCollection();
  };

  window.numzooMarkFailed = function(id) {
    id = String(id);
    saveCol(getCol().filter(function(x){return x.id!==id;}));
    renderCollection();
  };

  function renderCollection() {
    var col = getCol(), wrap = document.getElementById('numzoo-collection'),
        hdr  = document.getElementById('numzoo-collection-header');
    if (!wrap) return;
    wrap.innerHTML = '';
    if (!col.length) { hdr.style.display='none'; return; }
    hdr.style.display = 'block';
    col.slice().reverse().forEach(function(item) {
      var el = document.createElement('div');
      el.style.cssText = 'position:relative;width:80px;height:80px;border-radius:12px;overflow:hidden;flex-shrink:0;';
      if (item.status === 'locked') {
        el.style.background = '#1e1e2e';
        el.style.border = '2px dashed #555';
        el.innerHTML =
          '<div style="width:100%;height:100%;display:flex;align-items:center;justify-content:center;font-size:2em;z-index:1;position:relative;">🔒</div>' +
          '<div style="position:absolute;inset:0;background:linear-gradient(135deg,#2d2d44,#111128);animation:numzoo-pulse 2s ease-in-out infinite;"></div>';
      } else if (item.src) {
        var img = document.createElement('img');
        img.src = item.src;
        img.style.cssText = 'width:100%;height:100%;object-fit:cover;cursor:pointer;display:block;';
        img.onclick = function() {
          document.getElementById('numzoo-modal-img').src = item.src;
          document.getElementById('numzoo-modal').style.display = 'flex';
        };
        el.appendChild(img);
      }
      wrap.appendChild(el);
    });
  }

  document.getElementById('numzoo-modal').onclick = function() { this.style.display='none'; };

  // Poll the collection_trigger component for action payloads
  var lastTs = -1;
  setInterval(function() {
    var el = document.getElementById('numzoo-coll-data');
    if (!el) return;
    try {
      var data = JSON.parse(el.dataset.payload||'{}');
      if (!data.ts || data.ts === lastTs) return;
      lastTs = data.ts;
      (data.actions||[]).forEach(function(act) {
        if (act.action === 'add-locked') window.numzooAddLocked(act.id);
        else if (act.action === 'unlock')    window.numzooUnlock(act.id, act.src);
        else if (act.action === 'failed')    window.numzooMarkFailed(act.id);
      });
    } catch(e) {}
  }, 200);

  setTimeout(renderCollection, 800);
})();
</script>
"""

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

    state           = 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 ─────────────────────────────────────────────────────────
    collection_trigger = gr.HTML("")
    gr.HTML(COLLECTION_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)

    # "Let's go!" → start level → immediately pre-generate reward image
    go_btn.click(
        start_level, [animal_picker, place_picker, state],
        [state, emoji_panel, game_panel, picker_error, status_md, question_md, collection_trigger]
    ).then(
        pregenerate_image, [state],
        [state, hidden_image, hidden_data_url, collection_trigger]
    )

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

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

    # "Next Level" → increment level → back to picker with fresh picks
    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])


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