src/ui/mfa_render.py

"""MFA presentation — regex HTML timestamp injection + the Animate-All
progress bar. Everything between "raw MFA results" and "timestamped HTML"."""
from src.ui.mfa_lookups import (
    _build_crossword_groups, _build_enriched_json, _build_timestamp_lookups,
    _extend_word_timestamps,
)


def _ts_progress_bar_html(total_segments, rate, animated=True):
    """Return HTML for a progress bar showing Segment x/N.

    When *animated* is False, shows static "Preparing Alignment..." at 0%.
    When True, a single requestAnimationFrame loop drives BOTH the fill
    (`transform:scaleX(0→1)`) AND the counter text from one shared
    elapsed-time source — so they cannot drift, and `transform:scaleX`
    can never overflow the container's bounding box (no width-animation
    jitter). The script is injected via an `<img onerror>` trick because
    Gradio innerHTML doesn't execute `<script>` tags.
    """
    import random
    duration = total_segments * rate
    uid = f"tspb{random.randint(0, 999999)}"

    if animated:
        # One RAF loop owns both the fill scale and the counter text.
        counter_js = f'''<img src="data:," style="display:none"
            onerror="(function(){{
                var T={total_segments},DUR={duration};
                var fill=document.getElementById('{uid}-fill');
                var text=document.getElementById('{uid}-text');
                if(!fill||!text)return;
                var t0=performance.now();
                function tick(now){{
                    var e=(now-t0)/1000,pct=Math.min(e/DUR,1);
                    fill.style.transform='scaleX('+pct+')';
                    if(pct>=1){{
                        text.textContent='Almost Done...';
                    }}else{{
                        var seg=Math.min(Math.floor(e/(DUR/T))+1,T);
                        text.textContent='Segment '+seg+'/'+T;
                    }}
                    if(pct<1)requestAnimationFrame(tick);
                }}
                requestAnimationFrame(tick);
            }})()" />'''
        initial_text = f"Segment 0/{total_segments}"
    else:
        counter_js = ""
        initial_text = "Preparing Alignment..."

    return f'''<div id="{uid}" style="
        position:relative; width:100%; height:40px;
        background:#e5e7eb; border-radius:8px; overflow:hidden;
        font-family:system-ui,sans-serif; font-size:14px;
        isolation:isolate;
    ">
        <div id="{uid}-fill" style="
            position:absolute; inset:0;
            background:linear-gradient(90deg,#3b82f6,#2563eb);
            transform:scaleX(0); transform-origin:left center;
            will-change:transform;
        "></div>
        <span id="{uid}-text" style="
            position:absolute; inset:0; display:flex;
            align-items:center; justify-content:center;
            color:#1f2937; font-weight:600; z-index:1;
            text-shadow:0 0 4px rgba(255,255,255,0.8);
        ">{initial_text}</span>
        {counter_js}
    </div>'''


def inject_timestamps_into_html(current_html, segments, results, seg_to_result_idx, segment_dir):
    """Inject word and char timestamps into rendered segment HTML.

    Builds lookups, cross-word groups, extends timestamps, then performs
    regex-based injection of data-start/data-end attributes into word and
    char spans. Reusable by both the main MFA flow and the Dev tab
    log-based flow.

    Returns (enriched_html, enriched_json).
    """
    import re
    import unicodedata

    # Build timestamp lookups
    word_timestamps, letter_timestamps, word_to_all_results = _build_timestamp_lookups(results)
    crossword_groups = _build_crossword_groups(results, letter_timestamps)
    _extend_word_timestamps(word_timestamps, segments, seg_to_result_idx, results, segment_dir)

    # Inject timestamps into word spans, using segment boundaries to determine result_idx
    seg_boundaries = []
    for m in re.finditer(r'data-segment-idx="(\d+)"', current_html):
        seg_boundaries.append((m.start(), int(m.group(1))))
    seg_boundaries.sort(key=lambda x: x[0])

    seg_offset_map = {}
    for seg in segments:
        idx = seg.get("segment", 0) - 1
        seg_offset_map[idx] = seg.get("time_from", 0)

    def _get_seg_idx_at_pos(pos):
        seg_idx = None
        for boundary_pos, idx in seg_boundaries:
            if boundary_pos > pos:
                break
            seg_idx = idx
        return seg_idx

    word_open_re = r'<span class="word"[^>]*>'

    def _inject_word_ts(m):
        orig = m.group(0)
        pos_m = re.search(r'data-pos="([^"]+)"', orig)
        if not pos_m:
            return orig
        pos = pos_m.group(1)
        line_m = re.search(r'data-line-idx="(\d+)"', orig)
        line_idx = int(line_m.group(1)) if line_m else None
        seg_idx = _get_seg_idx_at_pos(m.start())
        if seg_idx is None:
            return orig
        expected_result_idx = seg_to_result_idx.get(seg_idx)
        result_idx = None
        if pos and not pos.startswith("0:0:"):
            candidates = word_to_all_results.get(pos, [])
            if candidates:
                if len(candidates) == 1:
                    result_idx = candidates[0]
                elif expected_result_idx in candidates:
                    result_idx = expected_result_idx
                else:
                    result_idx = min(candidates, key=lambda r: abs(r - (expected_result_idx or 0)))
        if result_idx is None:
            result_idx = expected_result_idx
        if result_idx is None:
            return orig
        # Try line-keyed lookup first (merge-with-repetition), then fall back
        # to the location-only key for normal results.
        ts = None
        if line_idx is not None:
            ts = word_timestamps.get(f"{result_idx}:{pos}:L{line_idx}")
        if ts is None:
            ts = word_timestamps.get(f"{result_idx}:{pos}")
        if not ts:
            return orig
        seg_offset = seg_offset_map.get(seg_idx, 0)
        abs_start = ts[0] + seg_offset
        abs_end = ts[1] + seg_offset
        return orig[:-1] + f' data-result-idx="{result_idx}" data-start="{abs_start:.4f}" data-end="{abs_end:.4f}">'

    html = re.sub(word_open_re, _inject_word_ts, current_html)

    # Enable per-segment animate buttons
    html = re.sub(r'(<button class="animate-btn"[^>]*?)\s+disabled(?:="[^"]*")?', r'\1', html)

    # Create char spans for timestamped words that don't have them yet
    # (char spans are deferred from initial render to reduce HTML size)
    from src.ui.segments import split_into_char_groups, ZWSP, DAGGER_ALEF

    def _create_char_spans(m):
        word_open = m.group(1)
        inner = m.group(2)
        if '<span class="char"' in inner:
            return m.group(0)  # Already has char spans (attrs or not)
        chars = []
        for g in split_into_char_groups(inner):
            if g.startswith(DAGGER_ALEF):
                chars.append(f'<span class="char">{ZWSP}{g}</span>')
            else:
                chars.append(f'<span class="char">{g}</span>')
        return f'{word_open}{"".join(chars)}</span>'

    html = re.sub(
        r'(<span class="word"[^>]*data-start="[\d.]+"[^>]*>)(.*?)</span>',
        _create_char_spans,
        html,
    )

    # Stamp char spans with MFA letter timestamps
    def _stamp_chars_with_mfa(word_m):
        word_open = word_m.group(1)
        word_abs_start = float(word_m.group(2))
        inner = word_m.group(4)

        pos_m = re.search(r'data-pos="([^"]+)"', word_open)
        word_pos = pos_m.group(1) if pos_m else None
        line_m = re.search(r'data-line-idx="(\d+)"', word_open)
        word_line_idx = int(line_m.group(1)) if line_m else None

        result_idx_m = re.search(r'data-result-idx="(\d+)"', word_open)
        if result_idx_m:
            result_idx = int(result_idx_m.group(1))
        else:
            result_idx = None
            if word_pos and not word_pos.startswith("0:0:"):
                candidates = word_to_all_results.get(word_pos, [])
                if candidates:
                    if len(candidates) == 1:
                        result_idx = candidates[0]
                    else:
                        result_idx = candidates[0]

        if result_idx is None or not word_pos:
            key = None
        elif word_line_idx is not None:
            key = f"{result_idx}:{word_pos}:L{word_line_idx}"
        else:
            key = f"{result_idx}:{word_pos}"

        word_ts = word_timestamps.get(key) if key else None
        mfa_letters = letter_timestamps.get(key) if key else None
        # Fall back to location-only key if the line-keyed lookup misses
        # (e.g. legacy injected HTML without data-line-idx).
        if (word_ts is None or mfa_letters is None) and key and word_line_idx is not None:
            fallback_key = f"{result_idx}:{word_pos}"
            if word_ts is None:
                word_ts = word_timestamps.get(fallback_key)
            if mfa_letters is None:
                mfa_letters = letter_timestamps.get(fallback_key)
        if not mfa_letters or not word_ts:
            return word_m.group(0)

        word_rel_start = word_ts[0]

        char_matches = list(re.finditer(r'<span class="char">([^<]*)</span>', inner))
        if not char_matches:
            return word_m.group(0)

        mfa_chars = [l["char"] for l in mfa_letters]
        html_chars = [m.group(1).replace('\u0640', '') for m in char_matches]

        CHAR_EQUIVALENTS = {
            '\u0649': '\u064a',  # alef maqsura <-> ya
            '\u064a': '\u0649',
        }

        def _first_base(s):
            for c in unicodedata.normalize("NFD", s):
                if not unicodedata.category(c).startswith('M'):
                    return c
            return s[0] if s else ''

        def chars_match(mfa_c, html_c):
            if mfa_c == html_c or html_c in mfa_c or mfa_c in html_c:
                return True
            if CHAR_EQUIVALENTS.get(mfa_c) == html_c:
                return True
            mb, hb = _first_base(mfa_c), _first_base(html_c)
            if mb and hb and (mb == hb or CHAR_EQUIVALENTS.get(mb) == hb):
                return True
            return False

        mfa_idx = 0
        char_replacements = []
        stamped_html = set()
        for html_idx, cm in enumerate(char_matches):
            if html_idx in stamped_html:
                continue
            html_char = html_chars[html_idx]
            if mfa_idx < len(mfa_letters):
                mfa_char = mfa_chars[mfa_idx]
                if chars_match(mfa_char, html_char):
                    letter = mfa_letters[mfa_idx]
                    if letter["start"] is None or letter["end"] is None:
                        if chars_match(mfa_char, html_char) or len(html_char) >= len(mfa_char):
                            mfa_idx += 1
                        continue
                    abs_start = word_abs_start + (letter["start"] - word_rel_start)
                    abs_end = word_abs_start + (letter["end"] - word_rel_start)
                    crossword_gid = crossword_groups.get((key, mfa_idx), "")
                    final_group_id = crossword_gid or letter.get("group_id", "")
                    char_replacements.append((
                        cm.start(), cm.end(),
                        f'<span class="char" data-start="{abs_start:.4f}" data-end="{abs_end:.4f}" data-group-id="{final_group_id}">{cm.group(1)}</span>'
                    ))
                    mfa_nfd = unicodedata.normalize("NFD", letter["char"])
                    peek = html_idx + 1
                    while peek < len(char_matches):
                        peek_raw = char_matches[peek].group(1).replace('\u0640', '')
                        if not peek_raw or not all(unicodedata.category(c).startswith('M') for c in peek_raw):
                            break
                        if not any(c in mfa_nfd for c in peek_raw):
                            break
                        char_replacements.append((
                            char_matches[peek].start(), char_matches[peek].end(),
                            f'<span class="char" data-start="{abs_start:.4f}" data-end="{abs_end:.4f}" data-group-id="{final_group_id}">{char_matches[peek].group(1)}</span>'
                        ))
                        stamped_html.add(peek)
                        peek += 1
                    if chars_match(mfa_char, html_char) or len(html_char) >= len(mfa_char):
                        mfa_idx += 1

        stamped_inner = inner
        for start, end, replacement in reversed(char_replacements):
            stamped_inner = stamped_inner[:start] + replacement + stamped_inner[end:]

        return f'{word_open}{stamped_inner}</span>'

    html = re.sub(
        r'(<span class="word"(?:\s+data-pos="[^"]*")?(?:\s+data-result-idx="\d+")?\s+data-start="([\d.]+)"\s+data-end="([\d.]+)">)((?:<span class="char">.*?</span>)+)</span>',
        _stamp_chars_with_mfa,
        html,
    )

    # Build enriched JSON (words only for download)
    enriched_json = _build_enriched_json(
        segments, results, seg_to_result_idx,
        word_timestamps, letter_timestamps, "words",
    )

    return html, enriched_json