dataset/templates.py

"""Control-flow templates for Python, JavaScript, C++ and C.

Each template builds the source code and its Mermaid flowchart from the same
state, so the `<linemap>` is correct by construction regardless of language
(C++ snippets carry `#include` lines that shift line numbers — handled
automatically because line numbers are read back from the builder).

C/C++ specifics handled here:
  * typed signatures and return types (via Lang); C collections gain a `_len`
    companion parameter and use index loops, C++ uses std::vector + ranged-for.
  * every return statement in a function is type-consistent (so it compiles).
  * no calls to undeclared helpers (C/C++ would reject them); loops self-advance.
  * C++ emits `#include <vector>`/`<string>`; C stays header-free.

A template signature is ``fn(rng, lang: Lang) -> Example``. TEMPLATES records the
languages each one supports.
"""

from __future__ import annotations

from engine import Example, Mermaid, assemble_output
from pools import (
    CLASSIFY_TRIPLES, COMMENTS, DOCSTRINGS, NUM1, NUM_T, RET, STATUS_PAIRS,
    COLLECTIONS, FLAGS, ITEMVARS, KEYS, Lang, SCALARS, STRINGS, TARGETS,
    camel, fn_name,
)

THRESHOLDS = [1, 2, 3, 5, 10, 18, 50, 60, 90, 100]
TF = [("True", "False"), ("Yes", "No")]
LOOP_EACH = ["Each item", "Each value", "Next element", "For each"]
LOOP_DONE = ["Exhausted", "All processed", "Loop done", "No more items"]

CFAM = ("cpp", "c")

# (function name, base-case condition, base return, recursive return)
RECURSIONS = [
    ("factorial", "{n} <= 1", "1", "{n} * {fn}({n} - 1)"),
    ("fibonacci", "{n} <= 1", "{n}", "{fn}({n} - 1) + {fn}({n} - 2)"),
    ("sum_to", "{n} == 0", "0", "{n} + {fn}({n} - 1)"),
    ("power_two", "{n} == 0", "1", "2 * {fn}({n} - 1)"),
    ("count_down", "{n} <= 0", "0", "1 + {fn}({n} - 1)"),
    ("digit_count", "{n} == 0", "0", "1 + {fn}({n} - 1)"),
]
RVARS = ["n", "k", "depth", "num", "level", "count"]


def _preamble(cb, L: Lang) -> None:
    """C++ snippets need their STL headers to be valid (and syntax-checkable)."""
    if L.name == "cpp":
        cb.add("#include <vector>")
        cb.add("#include <string>")


def simple_cond(rng, L: Lang, allow=("num", "numt", "coll", "null", "flag")):
    """Return (label, rendered_code, param_spec) for a single-variable condition.

    The param_spec ((kind, name)) lets the caller type the function signature.
    """
    kinds = [k for k in allow if not (k == "null" and L.name in CFAM)]
    kind = rng.choice(kinds)
    if kind == "num":
        c = rng.choice(NUM1); v = rng.choice(SCALARS)
        return c.label, c.code(L.name, x=v), ("int", v)
    if kind == "numt":
        c = rng.choice(NUM_T); v = rng.choice(SCALARS); t = rng.choice(THRESHOLDS)
        return c.label, c.code(L.name, x=v, t=t), ("int", v)
    if kind == "coll":
        v = rng.choice(COLLECTIONS)
        if rng.random() < 0.5:
            return "Collection is empty?", L.cond_empty(v), ("coll", v)
        return "Collection has items?", L.cond_nonempty(v), ("coll", v)
    if kind == "null":
        v = rng.choice(SCALARS + STRINGS)
        if rng.random() < 0.5:
            return "Value is missing?", L.cond_null(v), ("str", v)
        return "Value is present?", L.cond_notnull(v), ("str", v)
    v = rng.choice(FLAGS)
    if rng.random() < 0.5:
        return "Flag is set?", L.cond_flag(v), ("bool", v)
    return "Input is invalid?", L.cond_notflag(v), ("bool", v)


def _finalize(rng, L: Lang, name: str, cb, m: Mermaid) -> Example:
    return Example(
        language=L.name,
        template=name,
        code=cb.render_numbered(),
        output=assemble_output(rng, m),
        source=cb.source(),
        n_nodes=len(m.nodes),
    )


# --------------------------------------------------------------------------- #
# Templates
# --------------------------------------------------------------------------- #


def t_guard_return(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, rng.choice(["validate", "math", "process"]), L.name)
    label, code, pspec = simple_cond(rng, L, allow=("null", "coll", "flag", "num"))
    cfam = L.name in CFAM
    ln_def = cb.add(L.def_header(fn, [pspec], ret="bool" if cfam else "int"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_if = cb.add(L.if_header(code), 1)
    if cfam:
        bad_code, bad_label = L.ret(L.false()), "Return false"
        good_code, good_label = L.ret(L.true()), "Return true"
    else:
        bad = RET[rng.choice(["none", "neg1", "empty_list", "false", "zero"])]
        good = RET[rng.choice(["result", "value", "total", "true"])]
        bad_code, bad_label = bad.code(L.name), bad.label
        good_code, good_label = good.code(L.name), good.label
    ln_bad = cb.add(bad_code, 2)
    L.close(cb, 1)
    cb.maybe_comment(rng, COMMENTS, 1)
    ln_good = cb.add(good_code, 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add(label, "decision", ln_if)
    c = m.add(bad_label, "rect", ln_bad)
    d = m.add(good_label, "rect", ln_good)
    t, f = rng.choice(TF)
    m.edge(a, b)
    m.edge(b, c, t)
    m.edge(b, d, f)
    return _finalize(rng, L, "guard_return", cb, m)


def t_if_else(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, rng.choice(["validate", "string", "game"]), L.name)
    var = rng.choice(SCALARS)
    pos_lit, neg_lit, pos_lbl, neg_lbl = rng.choice(STATUS_PAIRS)
    cond = rng.choice(NUM_T)
    thr = rng.choice(THRESHOLDS)
    code = cond.code(L.name, x=var, t=thr)

    ln_def = cb.add(L.def_header(fn, [("int", var)], ret="str"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_if = cb.add(L.if_header(code), 1)
    ln_pos = cb.add(L.ret(pos_lit), 2)
    cb.add(L.else_header(), 1)
    ln_neg = cb.add(L.ret(neg_lit), 2)
    L.close(cb, 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add(cond.label, "decision", ln_if)
    c = m.add(pos_lbl, "rect", ln_pos)
    d = m.add(neg_lbl, "rect", ln_neg)
    t, f = rng.choice(TF)
    m.edge(a, b)
    m.edge(b, c, t)
    m.edge(b, d, f)
    return _finalize(rng, L, "if_else", cb, m)


def t_if_elif_else(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, rng.choice(["validate", "math"]), L.name)
    var = rng.choice(SCALARS)
    hi_lit, mid_lit, lo_lit, hi_lbl, mid_lbl, lo_lbl = rng.choice(CLASSIFY_TRIPLES)
    t1, t2 = sorted(rng.sample(THRESHOLDS, 2), reverse=True)

    ln_def = cb.add(L.def_header(fn, [("int", var)], ret="str"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_if = cb.add(L.if_header(f"{var} >= {t1}"), 1)
    ln_hi = cb.add(L.ret(hi_lit), 2)
    ln_elif = cb.add(L.elif_header(f"{var} >= {t2}"), 1)
    ln_mid = cb.add(L.ret(mid_lit), 2)
    cb.add(L.else_header(), 1)
    ln_lo = cb.add(L.ret(lo_lit), 2)
    L.close(cb, 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Value meets upper threshold?", "decision", ln_if)
    c = m.add(hi_lbl, "rect", ln_hi)
    d = m.add("Value meets middle threshold?", "decision", ln_elif)
    e = m.add(mid_lbl, "rect", ln_mid)
    f = m.add(lo_lbl, "rect", ln_lo)
    m.edge(a, b)
    m.edge(b, c, "True")
    m.edge(b, d, "False")
    m.edge(d, e, "True")
    m.edge(d, f, "False")
    return _finalize(rng, L, "if_elif_else", cb, m)


def t_multi_guard(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "validate", L.name)
    p1 = rng.choice(SCALARS)
    p2 = rng.choice(COLLECTIONS)
    numc = rng.choice(NUM1)
    guards = [("Collection is empty?", L.cond_empty(p2)),
              (numc.label, numc.code(L.name, x=p1))]
    if L.name not in CFAM and rng.random() < 0.6:
        guards.insert(rng.randint(0, len(guards)), ("Value is missing?", L.cond_null(p1)))

    ln_def = cb.add(L.def_header(fn, [("int", p1), ("coll", p2)], ret="bool"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    prev = a
    for label, code in guards:
        ln_if = cb.add(L.if_header(code), 1)
        ln_bad = cb.add(L.ret(L.false()), 2)
        L.close(cb, 1)
        dec = m.add(label, "decision", ln_if)
        ret = m.add("Return false", "rect", ln_bad)
        m.edge(prev, dec) if prev is a else m.edge(prev, dec, "False")
        m.edge(dec, ret, "True")
        prev = dec
    ln_good = cb.add(L.ret(L.true()), 1)
    L.close(cb, 0)
    g = m.add("Return true", "rect", ln_good)
    m.edge(prev, g, "False")
    return _finalize(rng, L, "multi_guard", cb, m)


def t_for_accumulate(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "math", L.name)
    coll = rng.choice(COLLECTIONS)
    item = rng.choice(ITEMVARS)
    acc = rng.choice(["total", "count"])
    inner = rng.choice(NUM1 + NUM_T)
    thr = rng.choice(THRESHOLDS)

    ln_def = cb.add(L.def_header(fn, [("coll", coll)], ret="int"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_init = cb.add(L.declare(acc, "0"), 1)
    cb.maybe_comment(rng, COMMENTS, 1)
    ln_for, elem = L.foreach(cb, 1, item, coll)
    ln_if = cb.add(L.if_header(inner.code(L.name, x=elem, t=thr)), 2)
    add_amt = elem if acc == "total" else "1"
    ln_add = cb.add(L.aug(acc, "+", add_amt), 3)
    L.close(cb, 2)
    L.close(cb, 1)
    ln_ret = cb.add(L.ret(acc), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add(f"Set {acc} to zero", "rect", ln_init)
    c = m.add("Iterate over collection", "rect", ln_for)
    d = m.add(inner.label, "decision", ln_if)
    e = m.add(f"Update the {acc}", "rect", ln_add)
    g = m.add(f"Return the {acc}", "rect", ln_ret)
    m.edge(a, b)
    m.edge(b, c)
    m.edge(c, d, rng.choice(LOOP_EACH))
    m.edge(c, g, rng.choice(LOOP_DONE))
    m.edge(d, e, "True")
    m.edge(d, c, "False")
    m.edge(e, c, loop=True)
    return _finalize(rng, L, "for_accumulate", cb, m)


def t_for_search(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "search", L.name)
    coll = rng.choice(COLLECTIONS)
    target = rng.choice(TARGETS)
    item = rng.choice(ITEMVARS)

    ln_def = cb.add(L.def_header(fn, [("coll", coll), ("int", target)], ret="int"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
        ln_for = cb.add(f"for i, {item} in enumerate({coll}):", 1)
        ln_if = cb.add(f"if {item} == {target}:", 2)
        ln_hit = cb.add("return i", 3)
        ln_miss = cb.add("return -1", 1)
    else:
        ln_for = L.forrange(cb, 1, "i", L.length(coll))
        ln_if = cb.add(f"if ({coll}[i] {L.eq()} {target}) {{", 2)
        ln_hit = cb.add(L.ret("i"), 3)
        L.close(cb, 2)
        L.close(cb, 1)
        ln_miss = cb.add(L.ret("-1"), 1)
        L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Iterate over collection", "rect", ln_for)
    c = m.add("Item matches target?", "decision", ln_if)
    d = m.add("Return the index", "rect", ln_hit)
    e = m.add("Return not found", "rect", ln_miss)
    m.edge(a, b)
    m.edge(b, c, rng.choice(LOOP_EACH))
    m.edge(b, e, rng.choice(LOOP_DONE))
    m.edge(c, d, "Match")
    m.edge(c, b, "No match", loop=True)
    return _finalize(rng, L, "for_search", cb, m)


def t_while_counter(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "math", L.name)
    n = rng.choice(["n", "count", "remaining"])
    steps = rng.choice(["steps", "iterations", "passes"])

    ln_def = cb.add(L.def_header(fn, [("int", n)], ret="int"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_init = cb.add(L.declare(steps, "0"), 1)
    ln_while = cb.add(L.while_header(f"{n} > 0"), 1)
    ln_dec = cb.add(L.aug(n, "-", "1"), 2)
    ln_step = cb.add(L.aug(steps, "+", "1"), 2)
    L.close(cb, 1)
    ln_ret = cb.add(L.ret(steps), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add(f"Set {steps} to zero", "rect", ln_init)
    c = m.add("Counter still positive?", "decision", ln_while)
    d = m.add("Decrement and tally", "rect", f"{ln_dec}-{ln_step}")
    e = m.add(f"Return the {steps}", "rect", ln_ret)
    m.edge(a, b)
    m.edge(b, c)
    m.edge(c, d, "True")
    m.edge(d, c, loop=True)
    m.edge(c, e, "False")
    return _finalize(rng, L, "while_counter", cb, m)


def t_while_scan(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "search", L.name)
    coll = rng.choice(COLLECTIONS)

    ln_def = cb.add(L.def_header(fn, [("coll", coll)], ret="int"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
        ln_init = cb.add("idx = 0", 1)
        ln_while = cb.add(f"while idx < len({coll}):", 1)
        ln_if = cb.add(f"if {coll}[idx] <= 0:", 2)
        ln_adv = cb.add("idx += 1", 3)
        ln_cont = cb.add("continue", 3)
        ln_hit = cb.add(f"return {coll}[idx]", 2)
        ln_miss = cb.add("return -1", 1)
    else:
        ln_init = cb.add(L.declare("idx", "0"), 1)
        ln_while = cb.add(f"while (idx < {L.length(coll)}) {{", 1)
        ln_if = cb.add(f"if ({coll}[idx] <= 0) {{", 2)
        ln_adv = cb.add(L.aug("idx", "+", "1"), 3)
        ln_cont = cb.add(L.continue_(), 3)
        L.close(cb, 2)
        ln_hit = cb.add(L.ret(f"{coll}[idx]"), 2)
        L.close(cb, 1)
        ln_miss = cb.add(L.ret("-1"), 1)
        L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Set idx to zero", "rect", ln_init)
    c = m.add("Index in bounds?", "decision", ln_while)
    d = m.add("Value is not positive?", "decision", ln_if)
    e = m.add("Advance and continue", "rect", f"{ln_adv}-{ln_cont}")
    f = m.add("Return current value", "rect", ln_hit)
    g = m.add("Return not found", "rect", ln_miss)
    m.edge(a, b)
    m.edge(b, c)
    m.edge(c, d, "In bounds")
    m.edge(c, g, "Out of bounds")
    m.edge(d, e, "Yes")
    m.edge(d, f, "No")
    m.edge(e, c, loop=True)
    return _finalize(rng, L, "while_scan", cb, m)


def t_nested_for(rng, L: Lang) -> Example:
    cb = L.builder()
    fn = fn_name(rng, "process", L.name)
    matrix = rng.choice(["matrix", "grid", "rows", "table"])
    row = rng.choice(["row", "line", "group"])
    val = rng.choice(["value", "cell", "entry"])
    acc = "count"
    inner = rng.choice(NUM1)

    ln_def = cb.add(L.def_header(fn, [matrix]))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_init = cb.add(L.declare(acc, "0"), 1)
    ln_outer, _ = L.foreach(cb, 1, row, matrix)
    ln_inner, elem = L.foreach(cb, 2, val, row, ivar="j")
    ln_if = cb.add(L.if_header(inner.code(L.name, x=elem)), 3)
    ln_add = cb.add(L.aug(acc, "+", "1"), 4)
    L.close(cb, 3)
    L.close(cb, 2)
    L.close(cb, 1)
    ln_ret = cb.add(L.ret(acc), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add(f"Set {acc} to zero", "rect", ln_init)
    c = m.add("Iterate over rows", "rect", ln_outer)
    d = m.add("Iterate over values", "rect", ln_inner)
    e = m.add(inner.label, "decision", ln_if)
    f = m.add(f"Increment {acc}", "rect", ln_add)
    g = m.add(f"Return the {acc}", "rect", ln_ret)
    m.edge(a, b)
    m.edge(b, c)
    m.edge(c, d, "Each row")
    m.edge(c, g, "Rows done")
    m.edge(d, e, "Each value")
    m.edge(d, c, "Row done")
    m.edge(e, f, "True")
    m.edge(e, d, "False")
    m.edge(f, d, loop=True)
    return _finalize(rng, L, "nested_for", cb, m)


def t_nested_if(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "validate", L.name)
    var = rng.choice(SCALARS)
    thr = rng.choice(THRESHOLDS)

    ln_def = cb.add(L.def_header(fn, [("int", var)], ret="str"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_out = cb.add(L.if_header(f"{var} > 0"), 1)
    ln_in = cb.add(L.if_header(f"{var} > {thr}"), 2)
    ln_big = cb.add(L.ret('"large"'), 3)
    L.close(cb, 2)
    ln_small = cb.add(L.ret('"small"'), 2)
    L.close(cb, 1)
    ln_neg = cb.add(L.ret('"nonpositive"'), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Number is positive?", "decision", ln_out)
    c = m.add("Value above threshold?", "decision", ln_in)
    d = m.add("Return large band", "rect", ln_big)
    e = m.add("Return small band", "rect", ln_small)
    f = m.add("Return nonpositive band", "rect", ln_neg)
    m.edge(a, b)
    m.edge(b, c, "True")
    m.edge(b, f, "False")
    m.edge(c, d, "True")
    m.edge(c, e, "False")
    return _finalize(rng, L, "nested_if", cb, m)


def t_try_except(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "string", L.name)
    text = rng.choice(STRINGS)
    with_finally = (not L.name == "cpp") and rng.random() < 0.45

    if not with_finally:
        ln_def = cb.add(L.def_header(fn, [("str", text)], ret="int"))
        if L.py:
            ln_try = cb.add("try:", 1)
            ln_ok = cb.add(f"return int({text})", 2)
            cb.add("except ValueError:", 1)
            ln_bad = cb.add("return 0", 2)
        elif L.name == "javascript":
            ln_try = cb.add("try {", 1)
            ln_ok = cb.add(f"return parseInt({text});", 2)
            cb.add("} catch (err) {", 1)
            ln_bad = cb.add("return 0;", 2)
            L.close(cb, 1)
            L.close(cb, 0)
        else:  # cpp
            ln_try = cb.add("try {", 1)
            ln_ok = cb.add(f"return std::stoi({text});", 2)
            cb.add("} catch (...) {", 1)
            ln_bad = cb.add("return 0;", 2)
            L.close(cb, 1)
            L.close(cb, 0)
        m = Mermaid()
        a = m.add(f"Start: {fn}", "rect", ln_def)
        b = m.add("Conversion succeeds?", "decision", ln_try)
        c = m.add("Return the value", "rect", ln_ok)
        d = m.add("Return zero", "rect", ln_bad)
        m.edge(a, b)
        m.edge(b, c, "Success")
        m.edge(b, d, "Failure")
        return _finalize(rng, L, "try_except", cb, m)

    # python / javascript with finally
    path = rng.choice(["path", "source", "handle"])
    if L.py:
        ln_def = cb.add(L.def_header(fn, [path]))
        ln_init = cb.add("data = None", 1)
        ln_try = cb.add("try:", 1)
        ln_read = cb.add(f"data = read({path})", 2)
        cb.add("except OSError:", 1)
        ln_fail = cb.add('data = ""', 2)
        cb.add("finally:", 1)
        ln_clean = cb.add("cleanup()", 2)
        ln_ret = cb.add("return data", 1)
    else:
        ln_def = cb.add(L.def_header(fn, [path]))
        ln_init = cb.add("let data = null;", 1)
        ln_try = cb.add("try {", 1)
        ln_read = cb.add(f"data = read({path});", 2)
        cb.add("} catch (err) {", 1)
        ln_fail = cb.add('data = "";', 2)
        cb.add("} finally {", 1)
        ln_clean = cb.add("cleanup();", 2)
        L.close(cb, 1)
        ln_ret = cb.add("return data;", 1)
        L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Set data to empty", "rect", ln_init)
    c = m.add("Read succeeds?", "decision", ln_try)
    d = m.add("Store the contents", "rect", ln_read)
    e = m.add("Use empty data", "rect", ln_fail)
    f = m.add("Run cleanup", "rect", ln_clean)
    g = m.add("Return the value", "rect", ln_ret)
    m.edge(a, b)
    m.edge(b, c)
    m.edge(c, d, "Success")
    m.edge(c, e, "Failure")
    m.edge(d, f)
    m.edge(e, f)
    m.edge(f, g)
    return _finalize(rng, L, "try_except", cb, m)


def t_switch_match(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "string", L.name)
    var = rng.choice(["kind", "mode", "code", "status"])
    ret_lits = ['"first"', '"second"', '"third"']
    ret_lbls = ["Return first label", "Return second label", "Return third label"]
    case_words = ["First case", "Second case", "Third case"]
    n_cases = rng.choice([2, 3])

    ln_def = cb.add(L.def_header(fn, [("int", var)], ret="str"))
    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    if L.py:
        ln_switch = cb.add(f"match {var}:", 1)
    else:
        ln_switch = cb.add(f"switch ({var}) {{", 1)
    b = m.add("Which case applies?", "decision", ln_switch)
    m.edge(a, b)

    nodes = []
    for i in range(n_cases):
        cb.add(f"case {i}:", 2)
        ln_ret = cb.add(L.ret(ret_lits[i]), 3)
        nodes.append(m.add(ret_lbls[i], "rect", ln_ret))
    cb.add("case _:" if L.py else "default:", 2)
    ln_def_ret = cb.add(L.ret('"other"'), 3)
    if not L.py:
        L.close(cb, 1)
        L.close(cb, 0)

    d = m.add("Return default label", "rect", ln_def_ret)
    for i, node in enumerate(nodes):
        m.edge(b, node, case_words[i])
    m.edge(b, d, "Default")
    return _finalize(rng, L, "switch_match", cb, m)


def t_recursion(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    name, base_cond, base_ret, rec = rng.choice(RECURSIONS)
    fn = camel(name) if L.name in ("javascript", "cpp") else name
    n = rng.choice(RVARS)

    ln_def = cb.add(L.def_header(fn, [("int", n)], ret="int"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    cb.maybe_comment(rng, COMMENTS, 1, p=0.2)
    ln_if = cb.add(L.if_header(base_cond.format(n=n)), 1)
    ln_base = cb.add(L.ret(base_ret.format(n=n)), 2)
    L.close(cb, 1)
    ln_rec = cb.add(L.ret(rec.format(fn=fn, n=n)), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Base case reached?", "decision", ln_if)
    c = m.add("Return base value", "rect", ln_base)
    d = m.add("Recurse and combine", "rect", ln_rec)
    m.edge(a, b)
    m.edge(b, c, "True")
    m.edge(b, d, "False")
    return _finalize(rng, L, "recursion", cb, m)


def t_clamp(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "math", L.name)
    var = rng.choice(["value", "amount", "score", "x"])

    ln_def = cb.add(L.def_header(fn, [("int", var), ("int", "low"), ("int", "high")], ret="int"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_if1 = cb.add(L.if_header(f"{var} < low"), 1)
    ln_low = cb.add(L.ret("low"), 2)
    L.close(cb, 1)
    ln_if2 = cb.add(L.if_header(f"{var} > high"), 1)
    ln_high = cb.add(L.ret("high"), 2)
    L.close(cb, 1)
    ln_val = cb.add(L.ret(var), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Below the minimum?", "decision", ln_if1)
    c = m.add("Return the minimum", "rect", ln_low)
    d = m.add("Above the maximum?", "decision", ln_if2)
    e = m.add("Return the maximum", "rect", ln_high)
    f = m.add("Return the value", "rect", ln_val)
    m.edge(a, b)
    m.edge(b, c, "True")
    m.edge(b, d, "False")
    m.edge(d, e, "True")
    m.edge(d, f, "False")
    return _finalize(rng, L, "clamp", cb, m)


def t_dict_lookup(rng, L: Lang) -> Example:
    cb = L.builder()
    fn = fn_name(rng, "search", L.name)
    table = rng.choice(["table", "mapping", "lookup", "registry"])
    key = rng.choice(KEYS)

    if L.py:
        ln_def = cb.add(L.def_header(fn, [table, key]))
        ln_if = cb.add(f"if {key} in {table}:", 1)
        ln_hit = cb.add(f"return {table}[{key}]", 2)
        ln_miss = cb.add("return None", 1)
    else:
        ln_def = cb.add(L.def_header(fn, [table, key]))
        ln_if = cb.add(f"if ({key} in {table}) {{", 1)
        ln_hit = cb.add(f"return {table}[{key}];", 2)
        L.close(cb, 1)
        ln_miss = cb.add("return null;", 1)
        L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Key exists in map?", "decision", ln_if)
    c = m.add("Return the mapped value", "rect", ln_hit)
    d = m.add("Return nothing", "rect", ln_miss)
    m.edge(a, b)
    m.edge(b, c, "Found")
    m.edge(b, d, "Missing")
    return _finalize(rng, L, "dict_lookup", cb, m)


def t_ternary(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "string", L.name)
    var = rng.choice(SCALARS)
    thr = rng.choice(THRESHOLDS)
    hi, lo = '"high"', '"low"'

    ln_def = cb.add(L.def_header(fn, [("int", var)], ret="str"))
    if L.py:
        ln_assign = cb.add(f"result = {hi} if {var} > {thr} else {lo}", 1)
        ln_ret = cb.add("return result", 1)
    else:
        typ = {"javascript": "const", "cpp": "std::string", "c": "const char*"}[L.name]
        ln_assign = cb.add(f"{typ} result = {var} > {thr} ? {hi} : {lo};", 1)
        ln_ret = cb.add("return result;", 1)
        L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Value above threshold?", "decision", ln_assign)
    c = m.add("Use high label", "rect", ln_assign)
    d = m.add("Use low label", "rect", ln_assign)
    e = m.add("Return the result", "rect", ln_ret)
    m.edge(a, b)
    m.edge(b, c, "True")
    m.edge(b, d, "False")
    m.edge(c, e)
    m.edge(d, e)
    return _finalize(rng, L, "ternary", cb, m)


def t_bool_and(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "validate", L.name)
    age = rng.choice(["age", "score", "level"])
    flag = rng.choice(["member", "active", "verified"])
    thr = rng.choice(THRESHOLDS)

    ln_def = cb.add(L.def_header(fn, [("int", age), ("bool", flag)], ret="bool"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    ln_if = cb.add(L.if_header(f"{age} >= {thr} {L.and_op()} {flag}"), 1)
    ln_true = cb.add(L.ret(L.true()), 2)
    L.close(cb, 1)
    ln_false = cb.add(L.ret(L.false()), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Value meets threshold?", "decision", ln_if)
    c = m.add("Flag is set?", "decision", ln_if)
    d = m.add("Return true", "rect", ln_true)
    e = m.add("Return false", "rect", ln_false)
    m.edge(a, b)
    m.edge(b, c, "True")
    m.edge(b, e, "False")
    m.edge(c, d, "True")
    m.edge(c, e, "False")
    return _finalize(rng, L, "bool_and", cb, m)


def t_state_machine(rng, L: Lang) -> Example:
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "game", L.name)
    var = "state"
    transitions = [('"idle"', '"running"', "Currently idle?", "Return running state"),
                   ('"running"', '"paused"', "Currently running?", "Return paused state"),
                   ('"paused"', '"running"', "Currently paused?", "Return running state")]

    ln_def = cb.add(L.def_header(fn, [("str", var)], ret="str"))
    if L.py:
        cb.maybe_docstring(rng, DOCSTRINGS)
    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    prev = a
    first = True
    for lit, nxt, qlabel, rlabel in transitions:
        cond = f"{var} {L.eq()} {lit}"
        ln_if = cb.add(L.if_header(cond) if first else L.elif_header(cond), 1)
        first = False
        ln_ret = cb.add(L.ret(nxt), 2)
        dec = m.add(qlabel, "decision", ln_if)
        ret = m.add(rlabel, "rect", ln_ret)
        m.edge(prev, dec) if prev is a else m.edge(prev, dec, "No")
        m.edge(dec, ret, "Yes")
        prev = dec
    cb.add(L.else_header(), 1)
    ln_default = cb.add(L.ret('"idle"'), 2)
    L.close(cb, 1)
    L.close(cb, 0)
    g = m.add("Return idle state", "rect", ln_default)
    m.edge(prev, g, "No")
    return _finalize(rng, L, "state_machine", cb, m)


def t_for_else_py(rng, L: Lang) -> Example:
    cb = L.builder()
    fn = fn_name(rng, "search", L.name)
    coll = rng.choice(COLLECTIONS)
    target = rng.choice(TARGETS)
    item = rng.choice(ITEMVARS)

    ln_def = cb.add(L.def_header(fn, [coll, target]))
    ln_for = cb.add(f"for {item} in {coll}:", 1)
    ln_if = cb.add(f"if {item} == {target}:", 2)
    ln_break = cb.add("break", 3)
    cb.add("else:", 1)
    ln_missing = cb.add('return "missing"', 2)
    ln_found = cb.add('return "found"', 1)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Iterate over collection", "rect", ln_for)
    c = m.add("Item matches target?", "decision", ln_if)
    d = m.add("Stop searching", "rect", ln_break)
    e = m.add("Return missing", "rect", ln_missing)
    f = m.add("Return found", "rect", ln_found)
    m.edge(a, b)
    m.edge(b, c, "Each item")
    m.edge(b, e, "Exhausted")
    m.edge(c, d, "Match")
    m.edge(c, b, "No match", loop=True)
    m.edge(d, f)
    return _finalize(rng, L, "for_else", cb, m)


def t_for_found_js(rng, L: Lang) -> Example:
    cb = L.builder()
    fn = fn_name(rng, "search", L.name)
    coll = rng.choice(COLLECTIONS)
    target = rng.choice(TARGETS)
    item = rng.choice(ITEMVARS)

    ln_def = cb.add(L.def_header(fn, [coll, target]))
    ln_init = cb.add("let found = false;", 1)
    ln_for, _ = L.foreach(cb, 1, item, coll)
    ln_if = cb.add(f"if ({item} === {target}) {{", 2)
    ln_set = cb.add("found = true;", 3)
    ln_break = cb.add("break;", 3)
    L.close(cb, 2)
    L.close(cb, 1)
    ln_check = cb.add("if (found) {", 1)
    ln_found = cb.add('return "found";', 2)
    L.close(cb, 1)
    ln_missing = cb.add('return "missing";', 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Set found to false", "rect", ln_init)
    c = m.add("Iterate over collection", "rect", ln_for)
    d = m.add("Item matches target?", "decision", ln_if)
    e = m.add("Mark found and stop", "rect", f"{ln_set}-{ln_break}")
    f = m.add("Was a match found?", "decision", ln_check)
    g = m.add("Return found", "rect", ln_found)
    h = m.add("Return missing", "rect", ln_missing)
    m.edge(a, b)
    m.edge(b, c)
    m.edge(c, d, "Each item")
    m.edge(c, f, "Exhausted")
    m.edge(d, e, "Match")
    m.edge(d, c, "No match", loop=True)
    m.edge(e, f)
    m.edge(f, g, "True")
    m.edge(f, h, "False")
    return _finalize(rng, L, "for_found", cb, m)


def t_do_while(rng, L: Lang) -> Example:
    """do/while — JavaScript, C++ and C (Python has no do/while)."""
    cb = L.builder()
    _preamble(cb, L)
    fn = fn_name(rng, "math", L.name)
    acc = rng.choice(["total", "sum", "acc", "running"])
    cnt = rng.choice(["n", "x", "value", "step"])
    limit = rng.choice(["limit", "cap", "budget", "target"])

    ln_def = cb.add(L.def_header(fn, [("int", limit)], ret="int"))
    cb.maybe_comment(rng, COMMENTS, 1, p=0.25)
    ln_init1 = cb.add(L.declare(acc, "0"), 1)
    ln_init2 = cb.add(L.declare(cnt, "0"), 1)
    ln_do = cb.add("do {", 1)
    ln_draw = cb.add(L.aug(cnt, "+", "1"), 2)
    ln_add = cb.add(L.aug(acc, "+", cnt), 2)
    ln_while = cb.add(f"}} while ({acc} < {limit});", 1)
    ln_ret = cb.add(L.ret(acc), 1)
    L.close(cb, 0)

    m = Mermaid()
    a = m.add(f"Start: {fn}", "rect", ln_def)
    b = m.add("Initialize totals", "rect", f"{ln_init1}-{ln_init2}")
    c = m.add("Advance and add", "rect", f"{ln_draw}-{ln_add}")
    d = m.add("Total below limit?", "decision", ln_while)
    e = m.add("Return the total", "rect", ln_ret)
    m.edge(a, b)
    m.edge(b, c)
    m.edge(c, d)
    m.edge(d, c, "Yes", loop=True)
    m.edge(d, e, "No")
    return _finalize(rng, L, "do_while", cb, m)


# --------------------------------------------------------------------------- #
# Error node (unparseable input -> isolated error node, per constraint 3)
# --------------------------------------------------------------------------- #

_BROKEN = [
    ["<<<>>> not valid @@@", "??? )("],
    ["def (:", "    return"],
    ["function {{{", "  >> <<"],
    ["%%%% garbage input %%%%"],
    ["}{ ]] [[ ))((", "=== ;;;"],
    ["lorem ipsum dolor sit", "amet consectetur"],
    ["#### ?? ## !!", "$$ ^^ &&"],
    ["class :: ;;", "    @@@ ..."],
    ["import import import", "from from"],
    ["][ }{ )( ><", "~~ `` ''"],
    ["....,,,,;;;;", "----++++"],
    ["the quick brown fox", "jumps over"],
]


def t_error_node(rng, L: Lang) -> Example:
    lines = rng.choice(_BROKEN)
    cb = L.builder()
    for ln in lines:
        cb.add(ln)
    lead = rng.choice([
        "The provided text is not parseable source code; no control structures can be extracted.",
        "Input does not form valid source; the parser cannot recover any execution flow.",
        "No recognizable syntax is present, so no control-flow graph can be built.",
    ])
    thinking = ("<thinking>\n"
                f"1. {lead}\n"
                "2. Per protocol, emit a single isolated error node and no source map.\n"
                "</thinking>")
    output = thinking + "\ngraph TD\n    A[Error: unable to parse source]"
    return Example(language=L.name, template="error_node",
                   code=cb.render_numbered(), output=output,
                   source=cb.source(), n_nodes=1)


# --------------------------------------------------------------------------- #
# Registry
# --------------------------------------------------------------------------- #

ALL4 = ("python", "javascript", "cpp", "c")
PJC = ("python", "javascript", "cpp")
PJ = ("python", "javascript")
PY = ("python",)
JS = ("javascript",)
JCC = ("javascript", "cpp", "c")

TEMPLATES = [
    (t_guard_return, ALL4, 1.0),
    (t_if_else, ALL4, 1.0),
    (t_if_elif_else, ALL4, 1.0),
    (t_multi_guard, ALL4, 1.0),
    (t_for_accumulate, ALL4, 1.2),
    (t_for_search, ALL4, 1.1),
    (t_while_counter, ALL4, 1.0),
    (t_while_scan, ALL4, 1.0),
    (t_nested_for, PJ, 0.9),
    (t_nested_if, ALL4, 1.0),
    (t_try_except, PJC, 1.0),
    (t_switch_match, ALL4, 1.0),
    (t_recursion, ALL4, 1.1),
    (t_clamp, ALL4, 0.9),
    (t_dict_lookup, PJ, 0.8),
    (t_ternary, ALL4, 0.9),
    (t_bool_and, ALL4, 0.9),
    (t_state_machine, PJC, 1.0),
    (t_for_else_py, PY, 0.8),
    (t_for_found_js, JS, 0.7),
    (t_do_while, JCC, 1.0),
    (t_error_node, ALL4, 0.3),
]