# ------------------------------ Visualization (NEW) ------------------------------
# Diffusion-style consecutive output: only show the CURRENT output per frame.
# ------------------ Visualization (sanitized, masks stripped) ------------------
from PIL import Image, ImageDraw, ImageFont

import re
import unicodedata
from typing import Optional, List, Tuple, Annotated


def render_consecutive_trace_gif(
    trace: dict,
    tokenizer,
    out_path: str = "decode_trace.gif",
    font_size: int = 30,
    line_spacing: int = 12,
    frame_ms: int = 250,
    final_ms: int = 5000,  # final clean frame duration (ms)
    max_width: int = 1400,
    max_height: int = 3000,
    margin: int = 32,
    title_color=(80, 80, 80),
    text_color=(0, 0, 0),  # base black
    mask_color=(150, 150, 150),
    sub_nonmask_color=(200, 0, 0),  # persistent red
    ins_color=(0, 0, 200),  # persistent blue
    del_strike_color=(120, 120, 120),
    events_color=(30, 30, 30),
    box_color=(120, 120, 120),
    bg_color=(255, 255, 255),
):
    """
    Persistent coloring keyed by token *instance* (not token id):
      - Inserted tokens -> BLUE across frames (until deleted/substituted again).
      - Substitution nonmask→nonmask -> RED across frames (until deleted/substituted again).
      - Substitution mask→nonmask -> stays BLACK (no extra color).
    Adds a final clean frame (5s) with no events box.
    """
    from PIL import Image, ImageDraw, ImageFont
    import unicodedata

    # ---------- font ----------
    try:
        font = ImageFont.truetype(
            "assets/JetBrainsMono-VariableFont_wght.ttf", font_size
        )
    except Exception:
        print(f"fail to load target font")
        font = ImageFont.load_default()

    # ---------- helpers ----------
    def _sanitize_token(s: str) -> str:
        vis_mask_token = "[m]"
        s = unicodedata.normalize("NFKC", s)
        s = s.replace("Ċ", "\n").replace("▁", " ").replace("Ġ", " ")
        s = s.replace("\t", "    ")
        s = s.replace("\u00a0", " ").replace("\u2007", " ").replace("\u202f", " ")

        # replace mask variants
        if "mdm_mask" in s.lower():
            s = re.sub(r"<[\|]?\s*mdm_mask\s*[\|]?>", "[m]", s, flags=re.IGNORECASE)
            s = s.replace("mdm_mask", "[m]")
        if "mask" in s.lower():
            s = re.sub(r"<[\|]?\s*mask\s*[\|]?>", "[m]", s, flags=re.IGNORECASE)
            s = s.replace("mask", "[m]")

        # replace <|...|> format tokens with bracketed form
        s = re.sub(r"<\|\s*(.*?)\s*\|>", r"[\1]", s)
        return s

    def _tok_str(tok_id: int) -> str:
        try:
            s = tokenizer.decode([int(tok_id)], skip_special_tokens=False)
            s = s if s.strip() else f"<{int(tok_id)}>"
        except Exception:
            s = f"<{int(tok_id)}>"
        return s.replace("\n", "\\n")

    TOKEN_RE = re.compile(r"\s+|\S+")

    def _wrap_text(draw: ImageDraw.ImageDraw, text: str, width_px: int) -> List[str]:
        if text == "":
            return [""]
        lines: List[str] = []
        for para in text.split("\n"):
            tokens = TOKEN_RE.findall(para)
            cur = ""
            for tok in tokens:
                candidate = cur + tok
                if draw.textlength(candidate, font=font) <= width_px:
                    cur = candidate
                else:
                    if cur:
                        lines.append(cur)
                        cur = tok
                        while (
                            draw.textlength(cur, font=font) > width_px and len(cur) > 0
                        ):
                            lo, hi, fit = 1, len(cur), 1
                            while lo <= hi:
                                mid = (lo + hi) // 2
                                if draw.textlength(cur[:mid], font=font) <= width_px:
                                    fit, lo = mid, mid + 1
                                else:
                                    hi = mid - 1
                            lines.append(cur[:fit])
                            cur = cur[fit:]
                    else:
                        t = tok
                        while draw.textlength(t, font=font) > width_px and len(t) > 0:
                            lo, hi, fit = 1, len(t), 1
                            while lo <= hi:
                                mid = (lo + hi) // 2
                                if draw.textlength(t[:mid], font=font) <= width_px:
                                    fit, lo = mid, mid + 1
                                else:
                                    hi = mid - 1
                            lines.append(t[:fit])
                            t = t[fit:]
                        cur = t
            lines.append(cur)
        return lines or [""]

    tmp_img = Image.new("RGB", (10, 10), bg_color)
    tmp_draw = ImageDraw.Draw(tmp_img)
    text_width_budget = max_width - 2 * margin

    # mask detection
    MASK_IDS = set()
    if getattr(tokenizer, "mask_token_id", None) is not None:
        MASK_IDS.add(int(tokenizer.mask_token_id))
    MASK_STRINGS = set()
    mt = getattr(tokenizer, "mask_token", None)
    if mt is not None:
        MASK_STRINGS.add(str(mt))
    MASK_STRINGS.add("<mdm_mask>")

    def _is_mask_token(tok_id: int, tok_str_exact: str) -> bool:
        return (int(tok_id) in MASK_IDS) or (tok_str_exact in MASK_STRINGS)

    def _wrap_tokens_with_index(tokens, deleted_flags):
        lines, cur, cur_w = [], [], 0
        for i, tok in enumerate(tokens):
            t = _sanitize_token(tok)
            parts = t.split("\n")
            for j, seg in enumerate(parts):
                seg_rest = seg
                while seg_rest:
                    w = tmp_draw.textlength(seg_rest, font=font)
                    if cur_w + w <= text_width_budget or not cur:
                        cur.append((seg_rest, i, deleted_flags[i]))
                        cur_w += w
                        seg_rest = ""
                    else:
                        lines.append(cur)
                        cur, cur_w = [], 0
                if j != len(parts) - 1:
                    lines.append(cur)
                    cur, cur_w = [], 0
        if cur:
            lines.append(cur)
        return lines

    def _draw_dashed_rectangle(
        draw, xy, dash=8, gap=6, width=2, outline=(120, 120, 120)
    ):
        x0, y0, x1, y1 = xy
        x = x0
        while x < x1:
            x2 = min(x + dash, x1)
            draw.line([(x, y0), (x2, y0)], fill=outline, width=width)
            draw.line([(x, y1), (x2, y1)], fill=outline, width=width)
            x += dash + gap
        y = y0
        while y < y1:
            y2 = min(y + dash, y1)
            draw.line([(x0, y), (x0, y2)], fill=outline, width=width)
            draw.line([(x1, y), (x1, y2)], fill=outline, width=width)
            y += dash + gap

    def _ops_lines_for_step(st: dict):
        if st is None:
            return ["(no events)"]
        lines = []
        x_before = st["x_before_ids"]
        choose_del = st["choose_del"]
        choose_sub = st["choose_sub"]
        sub_samples = st["sub_samples"]
        ins_samples = st["ins_samples"]
        T = len(x_before)
        for i in range(T):
            if choose_del[i]:
                lines.append(f"DEL@{i}:{_tok_str(int(x_before[i]))}")
            elif choose_sub[i]:
                lines.append(
                    f"SUB@{i}:{_tok_str(int(x_before[i]))}->{_tok_str(int(sub_samples[i]))}"
                )
            if ins_samples[i] is not None:
                lines.append(f"INS@{i}->{i+1}:{_tok_str(int(ins_samples[i]))}")
        if not lines:
            lines.append("(no events)")
        return lines

    # ---- Instance-id machinery ----
    next_instance_id = 0

    def _new_inst():
        nonlocal next_instance_id
        val = next_instance_id
        next_instance_id += 1
        return val

    # Current sequence at the *start* (ids + instance_ids)
    curr_ids = list(trace["init"]["x_ids"])
    curr_inst = [_new_inst() for _ in curr_ids]

    # Persistent color by instance_id: {"blue", "red"}
    color_by_inst = {}

    # ---------- PASS 1: measure required heights per frame ----------
    measurement_payload = []

    for step_idx, st in enumerate([None] + trace["steps"]):
        # build augmented view
        if st is None:
            aug_ids = list(curr_ids)
            deleted_flags = [False] * len(aug_ids)
        else:
            x_before = st["x_before_ids"]
            choose_del = st["choose_del"]
            after_ids = st["x_after_ids"]
            deleted_positions = [i for i, d in enumerate(choose_del) if d]

            aug_ids = list(after_ids)
            deleted_flags = [False] * len(after_ids)
            for i in sorted(deleted_positions, reverse=True):
                aug_ids.insert(i, x_before[i])
                deleted_flags.insert(i, True)

        tokens = tokenizer.convert_ids_to_tokens(aug_ids)
        wrapped_lines = _wrap_tokens_with_index(tokens, deleted_flags)

        # estimate ops lines for this step
        if st:
            ops_text = "  • " + "  • ".join(_ops_lines_for_step(st))
        else:
            ops_text = "(no events)"
        ops_lines = _wrap_text(tmp_draw, ops_text, text_width_budget)

        # compute height needed
        body_h = len(wrapped_lines) * (font_size + line_spacing)
        ops_h = len(ops_lines) * (font_size + line_spacing) + font_size  # + 20
        required_h = margin + (font_size + line_spacing) + body_h + 20

        measurement_payload.append(
            {
                "step_idx": step_idx,
                "st": st,
                "aug_ids": aug_ids,
                "tokens": tokens,
                "deleted_flags": deleted_flags,
                "wrapped_lines": wrapped_lines,
                "ops_lines": ops_lines,
                "required_h": required_h,
            }
        )

    # Measure clean final frame (no events)
    final_text_ids = (
        trace["steps"][-1]["x_after_ids"] if trace["steps"] else trace["init"]["x_ids"]
    )
    final_tokens = tokenizer.convert_ids_to_tokens(final_text_ids)
    wrapped_clean = _wrap_tokens_with_index(final_tokens, [False] * len(final_tokens))
    clean_body_h = len(wrapped_clean) * (font_size + line_spacing)
    clean_required_h = margin + (font_size + line_spacing) + clean_body_h

    # Pick a single uniform canvas height
    max_required_h = max(
        [p["required_h"] for p in measurement_payload] + [clean_required_h]
    )  # + 20
    H = min(max_required_h, max_height)
    W = max_width

    # For each frame we need an augmented view (with deleted placeholders) to draw
    frames = []

    # Iterate steps; for step_idx==0 we still draw "initial state"
    steps_with_initial = [None] + trace["steps"]

    for step_idx, st in enumerate(steps_with_initial):
        if st is None:
            # initial frame: augmented is just current tokens
            aug_ids = list(curr_ids)
            aug_inst = list(curr_inst)
            aug_deleted = [False] * len(aug_ids)
            ops_lines = ["(no events)"]
            title = "initial state"
        else:
            title = f"t = {st['t']:.3f}"
            x_before = list(st["x_before_ids"])
            choose_del = list(st["choose_del"])
            choose_sub = list(st["choose_sub"])
            sub_samples = list(st["sub_samples"])
            ins_samples = list(st["ins_samples"])
            assert (
                len(x_before) == len(curr_ids) == len(curr_inst)
            ), "trace 'x_before' must match current sequence."

            # Build augmented (drawn) and next (state-after) in one pass
            aug_ids, aug_inst, aug_deleted = [], [], []
            next_ids, next_inst = [], []

            for i in range(len(x_before)):
                before_id = int(curr_ids[i])
                before_inst = curr_inst[i]

                if choose_del[i]:
                    # show deleted placeholder (strike-through)
                    aug_ids.append(before_id)
                    aug_inst.append(None)
                    aug_deleted.append(True)
                    # remove from next; also clear any persistent color
                    color_by_inst.pop(before_inst, None)
                else:
                    if choose_sub[i]:
                        after_id = int(sub_samples[i])
                        # in augmented we show the *after* token at same instance
                        aug_ids.append(after_id)
                        aug_inst.append(before_inst)
                        aug_deleted.append(False)
                        next_ids.append(after_id)
                        next_inst.append(before_inst)

                        # update persistence by source type
                        if int(before_id) in MASK_IDS:
                            # mask → nonmask: no extra color (ensure cleared)
                            color_by_inst.pop(before_inst, None)
                        else:
                            # nonmask → nonmask: mark RED
                            color_by_inst[before_inst] = "red"
                    else:
                        # keep
                        aug_ids.append(before_id)
                        aug_inst.append(before_inst)
                        aug_deleted.append(False)
                        next_ids.append(before_id)
                        next_inst.append(before_inst)

                # insertion AFTER position i
                if ins_samples[i] is not None:
                    ins_id = int(ins_samples[i])
                    ins_inst = _new_inst()
                    aug_ids.append(ins_id)
                    aug_inst.append(ins_inst)
                    aug_deleted.append(False)
                    next_ids.append(ins_id)
                    next_inst.append(ins_inst)
                    # mark persistent BLUE for this *instance only*
                    color_by_inst[ins_inst] = "blue"

            # commit next state
            curr_ids, curr_inst = next_ids, next_inst
            ops_text = "  • " + "  • ".join(_ops_lines_for_step(st))
            ops_lines = _wrap_text(tmp_draw, ops_text, text_width_budget)

        # ----- render this frame -----
        tokens = tokenizer.convert_ids_to_tokens(aug_ids)
        wrapped_lines = _wrap_tokens_with_index(tokens, aug_deleted)

        img = Image.new("RGB", (W, H), bg_color)
        draw = ImageDraw.Draw(img)

        y = margin
        draw.text((margin, y), title, fill=title_color, font=font)
        y += font_size + line_spacing

        for line in wrapped_lines:
            x = margin
            for seg_text, tok_idx, is_deleted in line:
                tok_id = int(aug_ids[tok_idx])
                tok_str_exact = tokens[tok_idx]
                inst = aug_inst[tok_idx]

                if is_deleted:
                    # strike deleted — grey masks slightly different if desired
                    strike_color = (
                        mask_color
                        if _is_mask_token(tok_id, tok_str_exact)
                        else del_strike_color
                    )
                    strike = "".join(ch + "\u0336" for ch in seg_text)
                    draw.text((x, y), strike, fill=strike_color, font=font)
                    x += tmp_draw.textlength(strike, font=font)
                else:
                    # choose color by *instance*
                    color = text_color
                    if inst is not None and inst in color_by_inst:
                        color = (
                            ins_color
                            if color_by_inst[inst] == "blue"
                            else sub_nonmask_color
                        )
                    elif _is_mask_token(tok_id, tok_str_exact):
                        color = mask_color
                    draw.text((x, y), seg_text, fill=color, font=font)
                    x += tmp_draw.textlength(seg_text, font=font)
            y += font_size + line_spacing

        # draw events box for all but the extra final-clean frame we'll add later
        # if step_idx != len(steps_with_initial) - 1:
        #     y += 20
        #     x0, y0 = margin, y
        #     x1 = max_width - margin
        #     box_h = len(ops_lines) * (font_size + line_spacing) + font_size + 20
        #     y1 = y0 + box_h
        #     _draw_dashed_rectangle(draw, (x0, y0, x1, y1), outline=box_color)
        #     draw.text((x0 + 10, y0 + 10), "events", fill=events_color, font=font)
        #     yy = y0 + font_size + 20
        #     for l in ops_lines:
        #         draw.text((x0 + 10, yy), l, fill=events_color, font=font)
        #         yy += font_size + line_spacing
        # y += 10
        frames.append(img)

    # ----- extra final clean frame (no events box), 5s -----
    final_ids = list(curr_ids)
    final_inst = list(curr_inst)
    final_tokens = tokenizer.convert_ids_to_tokens(final_ids)

    # wrap without deleted flags
    def _wrap_clean(tokens):
        lines, cur, cur_w = [], [], 0
        for i, tok in enumerate(tokens):
            t = _sanitize_token(tok)
            parts = t.split("\n")
            for j, seg in enumerate(parts):
                seg_rest = seg
                while seg_rest:
                    w = tmp_draw.textlength(seg_rest, font=font)
                    if cur_w + w <= text_width_budget or not cur:
                        cur.append((seg_rest, i))
                        cur_w += w
                        seg_rest = ""
                    else:
                        lines.append(cur)
                        cur, cur_w = [], 0
                if j != len(parts) - 1:
                    lines.append(cur)
                    cur, cur_w = [], 0
        if cur:
            lines.append(cur)
        return lines

    wrapped_clean = _wrap_clean(final_tokens)

    clean_img = Image.new("RGB", (W, H), bg_color)
    draw = ImageDraw.Draw(clean_img)
    draw.text((margin, margin), "final text", fill=title_color, font=font)
    y = margin + font_size + line_spacing
    for line in wrapped_clean:
        x = margin
        for seg_text, tok_idx in line:
            tok_id = int(final_ids[tok_idx])
            tok_str_exact = final_tokens[tok_idx]
            inst = final_inst[tok_idx]
            color = text_color
            if inst in color_by_inst:
                color = (
                    ins_color if color_by_inst[inst] == "blue" else sub_nonmask_color
                )
            elif _is_mask_token(tok_id, tok_str_exact):
                color = mask_color
            draw.text((x, y), seg_text, fill=color, font=font)
            x += tmp_draw.textlength(seg_text, font=font)
        y += font_size + line_spacing
    frames.append(clean_img)

    # save GIF
    durations = [frame_ms] * (len(frames) - 1) + [final_ms]
    frames[0].save(
        out_path,
        save_all=True,
        append_images=frames[1:],
        duration=durations,
        loop=0,
        disposal=2,
        optimize=True,
    )
    return out_path
