"""Math game logic: problem generation + adaptive difficulty.

No UI / model dependencies live here so it can be unit-tested in isolation.

A "problem" is a plain dict:
    {
        "skill":   str,            # which skill this exercises
        "display": str,            # big text/emoji shown on screen
        "spoken":  str,            # plain-English text read aloud by TTS
        "answer":  str,            # the correct choice (as a string)
        "choices": list[str],      # 4 shuffled multiple-choice options
    }

Game state is a plain dict so it can live in a Gradio gr.State:
    {"level": int, "score": int, "streak": int, "wrong_streak": int, "asked": int}
"""

import random

MIN_LEVEL = 1
MAX_LEVEL = 10
LEVEL_UP_STREAK = 3      # correct-in-a-row needed to level up
LEVEL_DOWN_STREAK = 2    # wrong-in-a-row before easing down

SHAPES = ["circle", "square", "triangle", "star", "heart", "diamond"]
SHAPE_EMOJI = {
    "circle": "⚪",
    "square": "🟦",
    "triangle": "🔺",
    "star": "⭐",
    "heart": "❤️",
    "diamond": "🔷",
}
SHAPE_SIDES = {"triangle": 3, "square": 4, "circle": 0}

# Friendly emoji used to render "count the objects" problems.
OBJECT_EMOJI = ["🍎", "🍌", "🐶", "🐱", "🐟", "🌟", "🚗", "🎈", "🍓", "🦋"]


def new_state():
    """Return a fresh game-state dict."""
    return {"level": 1, "score": 0, "streak": 0, "wrong_streak": 0, "asked": 0}


# --------------------------------------------------------------------------
# Difficulty / skill selection
# --------------------------------------------------------------------------

def _skills_for_level(level):
    """Which skills are unlocked at a given level (harder skills appear later)."""
    skills = ["add_sub", "count_compare", "shapes_patterns"]
    if level >= 4:
        skills.append("multiply")
    return skills


def update_state(state, correct):
    """Apply the adaptive rule after an answer. Mutates and returns ``state``.

    Level up after LEVEL_UP_STREAK correct in a row; level down after
    LEVEL_DOWN_STREAK wrong in a row. Streaks reset on the opposite outcome.
    """
    state["asked"] += 1
    if correct:
        state["score"] += 1
        state["streak"] += 1
        state["wrong_streak"] = 0
        if state["streak"] >= LEVEL_UP_STREAK and state["level"] < MAX_LEVEL:
            state["level"] += 1
            state["streak"] = 0
    else:
        state["wrong_streak"] += 1
        state["streak"] = 0
        if state["wrong_streak"] >= LEVEL_DOWN_STREAK and state["level"] > MIN_LEVEL:
            state["level"] -= 1
            state["wrong_streak"] = 0
    return state


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

def _make_choices(answer, distractors, n=4):
    """Build n unique shuffled string choices that always include ``answer``."""
    answer_str = str(answer)
    choices = [answer_str]
    for d in distractors:
        d = str(d)
        if d not in choices:
            choices.append(d)
        if len(choices) >= n:
            break
    # Pad if we still came up short (e.g. tiny numbers near 0): step outward.
    if len(choices) < n:
        try:
            base = int(float(answer))
            k = 1
            while len(choices) < n and k < 50:
                cand = str(base + k)
                if cand not in choices:
                    choices.append(cand)
                k += 1
        except (ValueError, TypeError):
            pass
    random.shuffle(choices)
    return choices


def _numeric_distractors(answer, spread=3, allow_negative=False):
    """Plausible near-miss numbers around ``answer``."""
    candidates = set()
    attempts = 0
    while len(candidates) < 6 and attempts < 40:
        attempts += 1
        delta = random.randint(-spread, spread)
        val = answer + delta
        if val == answer:
            continue
        if val < 0 and not allow_negative:
            continue
        candidates.add(val)
    return list(candidates)


# --------------------------------------------------------------------------
# Skill generators -- each returns a problem dict
# --------------------------------------------------------------------------

def _gen_add_sub(level):
    if level <= 1:
        a, b = random.randint(0, 5), random.randint(0, 5)
        op = "+"
    elif level == 2:
        a, b = random.randint(0, 10), random.randint(0, 10)
        op = "+"
    elif level == 3:
        a = random.randint(2, 10)
        b = random.randint(0, a)            # subtraction within 10, no negatives
        op = "-"
    elif level == 4:
        op = random.choice(["+", "-"])
        if op == "+":
            a, b = random.randint(0, 20), random.randint(0, 20 - 0)
            b = random.randint(0, 20 - a) if a <= 20 else 0
        else:
            a = random.randint(2, 20)
            b = random.randint(0, a)
    elif level == 5:
        # missing addend:  a + ? = total
        a = random.randint(1, 15)
        b = random.randint(1, 20 - a) if a < 20 else 1
        total = a + b
        return {
            "skill": "add_sub",
            "display": f"{a} + ? = {total}",
            "spoken": f"{a} plus what makes {total}?",
            "answer": str(b),
            "choices": _make_choices(b, _numeric_distractors(b, 3)),
        }
    elif level <= 7:
        op = random.choice(["+", "-"])
        if op == "+":
            a = random.randint(10, 50)
            b = random.randint(1, 99 - a)
        else:
            a = random.randint(10, 99)
            b = random.randint(1, a)
    else:  # level 8-10
        op = random.choice(["+", "-"])
        if op == "+":
            a = random.randint(20, 80)
            b = random.randint(10, 99 - a) if a < 99 else 10
        else:
            a = random.randint(30, 99)
            b = random.randint(10, a)

    answer = a + b if op == "+" else a - b
    spread = 3 if answer < 20 else max(3, answer // 10)
    return {
        "skill": "add_sub",
        "display": f"{a} {op} {b} = ?",
        "spoken": f"What is {a} {'plus' if op == '+' else 'minus'} {b}?",
        "answer": str(answer),
        "choices": _make_choices(answer, _numeric_distractors(answer, spread)),
    }


def _gen_count_compare(level):
    if level <= 1:
        n = random.randint(1, 5)
        emoji = random.choice(OBJECT_EMOJI)
        return {
            "skill": "count_compare",
            "display": emoji * n,
            "spoken": "How many do you see?",
            "answer": str(n),
            "choices": _make_choices(n, _numeric_distractors(n, 2)),
        }
    if level == 2:
        n = random.randint(3, 10)
        emoji = random.choice(OBJECT_EMOJI)
        return {
            "skill": "count_compare",
            "display": emoji * n,
            "spoken": "How many do you see?",
            "answer": str(n),
            "choices": _make_choices(n, _numeric_distractors(n, 2)),
        }
    if level == 3:
        a, b = random.sample(range(0, 11), 2)
        bigger = max(a, b)
        return {
            "skill": "count_compare",
            "display": f"{a}   or   {b}",
            "spoken": f"Which number is bigger, {a} or {b}?",
            "answer": str(bigger),
            "choices": _make_choices(
                bigger, [min(a, b)] + _numeric_distractors(bigger, 3)
            ),
        }
    # level 4+: missing number in a sequence (skip-counting at higher levels)
    if level <= 4:
        step = 1
    elif level <= 6:
        step = random.choice([1, 2, 5])
    else:
        step = random.choice([2, 5, 10])
    start = random.randint(0, 10) * step + random.randint(0, step)
    seq = [start + step * i for i in range(4)]
    hide = random.randint(1, 2)            # hide an interior term
    answer = seq[hide]
    shown = [str(x) if i != hide else "?" for i, x in enumerate(seq)]
    return {
        "skill": "count_compare",
        "display": "  ".join(shown),
        "spoken": "What number is missing?",
        "answer": str(answer),
        "choices": _make_choices(answer, _numeric_distractors(answer, max(2, step))),
    }


def _gen_multiply(level):
    if level <= 4:
        groups = random.randint(2, 3)
        per = random.randint(2, 3)
        emoji = random.choice(OBJECT_EMOJI)
        display = "   ".join([emoji * per] * groups)
        answer = groups * per
        return {
            "skill": "multiply",
            "display": display,
            "spoken": f"{groups} groups of {per}. How many in total?",
            "answer": str(answer),
            "choices": _make_choices(answer, _numeric_distractors(answer, 3)),
        }
    if level == 5:
        table = random.choice([2, 10])
        other = random.randint(1, 10)
    elif level <= 7:
        table = random.choice([2, 5, 10])
        other = random.randint(1, 10)
    else:
        table = random.randint(2, 9)
        other = random.randint(2, 9)
    answer = table * other
    return {
        "skill": "multiply",
        "display": f"{table} × {other} = ?",
        "spoken": f"What is {table} times {other}?",
        "answer": str(answer),
        "choices": _make_choices(answer, _numeric_distractors(answer, max(3, table))),
    }


def _gen_shapes_patterns(level):
    if level <= 1:
        shape = random.choice(SHAPES)
        distract = random.sample([s for s in SHAPES if s != shape], 3)
        return {
            "skill": "shapes_patterns",
            "display": SHAPE_EMOJI[shape],
            "spoken": "What shape is this?",
            "answer": shape,
            "choices": _make_choices(shape, distract),
        }
    if level == 2:
        target = random.choice(SHAPES)
        return {
            "skill": "shapes_patterns",
            "display": "🔎",
            "spoken": f"Which one is the {target}?",
            "answer": SHAPE_EMOJI[target],
            "choices": _make_choices(
                SHAPE_EMOJI[target],
                [SHAPE_EMOJI[s] for s in SHAPES if s != target],
            ),
        }
    if level <= 4:
        # ABAB pattern -> what comes next?
        a, b = random.sample(SHAPES, 2)
        ea, eb = SHAPE_EMOJI[a], SHAPE_EMOJI[b]
        seq = [ea, eb, ea, eb]
        nxt = ea          # next after ...ea, eb, ea, eb is ea
        return {
            "skill": "shapes_patterns",
            "display": " ".join(seq) + " ?",
            "spoken": "What comes next in the pattern?",
            "answer": nxt,
            "choices": _make_choices(nxt, [eb] + [SHAPE_EMOJI[s] for s in SHAPES
                                                  if s not in (a, b)]),
        }
    if level <= 6:
        # AABB pattern
        a, b = random.sample(SHAPES, 2)
        ea, eb = SHAPE_EMOJI[a], SHAPE_EMOJI[b]
        seq = [ea, ea, eb, eb, ea, ea]
        nxt = eb
        return {
            "skill": "shapes_patterns",
            "display": " ".join(seq) + " ?",
            "spoken": "What comes next in the pattern?",
            "answer": nxt,
            "choices": _make_choices(nxt, [ea] + [SHAPE_EMOJI[s] for s in SHAPES
                                                  if s not in (a, b)]),
        }
    # level 7+: count the sides of a shape
    shape = random.choice([s for s in SHAPE_SIDES])
    answer = SHAPE_SIDES[shape]
    return {
        "skill": "shapes_patterns",
        "display": SHAPE_EMOJI[shape],
        "spoken": f"How many sides does a {shape} have?",
        "answer": str(answer),
        "choices": _make_choices(answer, [0, 3, 4, 5, 6]),
    }


_GENERATORS = {
    "add_sub": _gen_add_sub,
    "count_compare": _gen_count_compare,
    "multiply": _gen_multiply,
    "shapes_patterns": _gen_shapes_patterns,
}


def generate_problem(level, skill=None):
    """Return a problem dict for ``level`` (1-10), optionally forcing a skill."""
    level = max(MIN_LEVEL, min(MAX_LEVEL, int(level)))
    if skill is None:
        skill = random.choice(_skills_for_level(level))
    problem = _GENERATORS[skill](level)
    # Safety: guarantee the answer is among the (deduplicated) choices.
    if problem["answer"] not in problem["choices"]:
        problem["choices"][0] = problem["answer"]
        random.shuffle(problem["choices"])
    return problem