src/ui/mfa_compute.py

"""MFA compute flows — qua_sdk transport, batch submission, the Animate-All
generator, the synchronous API variant, and the on-demand segment-WAV slicer.

The wire transport is qua_sdk's MfaRemoteSpace (timing.mfa_remote). Batches go
through the backend's transport with caller-built refs instead of
``backend.timestamps()`` because the legacy raw-results contract requires it:
refs may repeat within one batch (two Basmala takes must not share timings),
and the fused-combined split ref ("Isti'adha+Basmala+<span>") is not
expressible through ``build_ref``'s segment grammar.
"""
import os

import gradio as gr

from config import (MFA_SPACE_URL, MFA_MODEL_ID, MFA_TIMEOUT, MFA_PROGRESS_SEGMENT_RATE,
                    MFA_SPLIT_PADDING)
from src.ui.mfa_lookups import (
    _build_crossword_groups, _build_enriched_json, _build_mfa_ref,
    _build_timestamp_lookups, _extend_word_timestamps,
)
from src.ui.mfa_render import _ts_progress_bar_html, inject_timestamps_into_html


def _backend():
    from qua_sdk.components.timing.runtimes.remote_space import MfaRemoteSpace
    return MfaRemoteSpace(space_url=MFA_SPACE_URL)


def _timing_params(padding):
    from qua_sdk.components.timing.runtimes.mfa_params import MfaParams
    return MfaParams(padding=padding, timeout_s=MFA_TIMEOUT,
                     model_id=MFA_MODEL_ID or None)


def _mfa_align_batch(refs, audio_paths, padding="forward", on_started=None):
    """Submit one align_batch; block until the raw results list arrives.

    InfraError (Space paused / timeout) surfaces as gr.Error with the same
    user-facing text the legacy client raised; component failures become
    RuntimeError like the legacy SSE error paths.
    """
    from qua_sdk.errors import ComponentError, InfraError

    backend = _backend()
    params = _timing_params(padding)
    try:
        event_id, headers = backend._upload_and_submit(refs, audio_paths, params)
        return backend._wait_result(event_id, headers, params, on_started)
    except InfraError as e:
        raise gr.Error(e.user_message) from e
    except ComponentError as e:
        detail = f" ({e.internal_detail})" if e.internal_detail else ""
        raise RuntimeError(f"{e.user_message}{detail}") from e



def _mfa_align_batch_streaming(refs, audio_paths, padding="forward"):
    """Generator: yields ("started", None) when the MFA Space exits its queue
    and starts processing, then ("result", results_list) when alignment
    finishes. Lets callers swap a "Preparing alignment" indicator for a real
    progress bar at the moment work actually begins.

    The blocking SDK wait runs in a worker thread; its on_started callback is
    bridged through a queue so this stays a plain generator.
    """
    import queue
    import threading

    from qua_sdk.errors import ComponentError, InfraError

    backend = _backend()
    params = _timing_params(padding)
    try:
        event_id, headers = backend._upload_and_submit(refs, audio_paths, params)
    except InfraError as e:
        raise gr.Error(e.user_message) from e

    q = queue.Queue()

    def _wait():
        try:
            results = backend._wait_result(
                event_id, headers, params, lambda: q.put(("started", None)))
            q.put(("result", results))
        except BaseException as e:  # noqa: BLE001 — re-raised on the caller thread
            q.put(("error", e))

    threading.Thread(target=_wait, name="mfa_wait_result", daemon=True).start()

    while True:
        kind, data = q.get()
        if kind == "error":
            if isinstance(data, InfraError):
                raise gr.Error(data.user_message) from data
            if isinstance(data, ComponentError):
                detail = f" ({data.internal_detail})" if data.internal_detail else ""
                raise RuntimeError(f"{data.user_message}{detail}") from data
            raise data
        yield kind, data
        if kind == "result":
            return



def mfa_split_timestamps(audio_int16, sample_rate, mfa_refs,
                         padding=MFA_SPLIT_PADDING):
    """Call MFA to get word timestamps for splitting segments.

    Args:
        audio_int16: List of int16 audio arrays (one per segment to split).
        sample_rate: Audio sample rate.
        mfa_refs: List of MFA ref strings (one per segment).
        padding: Gap-padding strategy ("forward", "symmetric", "none").

    Returns:
        List of results (one per segment), each a list of
        {location, start, end} dicts, or None on failure for that segment.
    """
    import tempfile
    import wave

    if not mfa_refs or not audio_int16:
        return [None] * len(mfa_refs)

    # Write segment audio to temp WAV files
    audio_paths = []
    for audio in audio_int16:
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        with wave.open(tmp.name, "wb") as wf:
            wf.setnchannels(1)
            wf.setsampwidth(2)
            wf.setframerate(sample_rate)
            wf.writeframes(audio.tobytes())
        audio_paths.append(tmp.name)

    try:
        results = _mfa_align_batch(mfa_refs, audio_paths, padding=padding)
        print(f"[MFA_SPLIT] Got {len(results)} results from MFA API")

        out = []
        for result in results:
            if result.get("status") != "ok":
                print(f"[MFA_SPLIT] Segment failed: ref={result.get('ref')} error={result.get('error')}")
                out.append(None)
            else:
                out.append(result.get("words", []))
        return out

    except Exception as e:
        print(f"[MFA_SPLIT] MFA call failed: {e}")
        return [None] * len(mfa_refs)
    finally:
        import os as _os
        for p in audio_paths:
            try:
                _os.unlink(p)
            except OSError:
                pass


def _ensure_segment_wavs(segments, segment_dir):
    """Write individual segment WAVs from full.wav on demand (for MFA).

    Segments are sliced from the full recording using soundfile's
    frame-level random access — no need to load the entire file.
    """
    if not segment_dir:
        return
    full_path = os.path.join(segment_dir, "full.wav")
    if not os.path.exists(full_path):
        return
    import soundfile as sf
    info = sf.info(full_path)
    sr = info.samplerate
    written = 0
    for seg in segments:
        idx = seg.get("segment", 0) - 1
        wav_path = os.path.join(segment_dir, f"seg_{idx}.wav")
        if os.path.exists(wav_path):
            continue
        start_frame = int(seg.get("time_from", 0) * sr)
        stop_frame = int(seg.get("time_to", 0) * sr)
        audio_slice, _ = sf.read(full_path, start=start_frame, stop=stop_frame, dtype='int16')
        sf.write(wav_path, audio_slice, sr, format='WAV', subtype='PCM_16')
        written += 1
    if written:
        print(f"[MFA] Wrote {written} segment WAVs on demand from full.wav")


def _build_mfa_refs(segments, segment_dir, require_existing=True):
    """Build MFA refs and audio paths from segments.

    Returns (refs, audio_paths, seg_to_result_idx).
    """
    refs = []
    audio_paths = []
    seg_to_result_idx = {}

    for seg in segments:
        seg_idx = seg.get("segment", 0) - 1
        mfa_ref = _build_mfa_ref(seg)
        if mfa_ref is None:
            continue

        audio_path = os.path.join(segment_dir, f"seg_{seg_idx}.wav") if segment_dir else None
        if not audio_path or (require_existing and not os.path.exists(audio_path)):
            continue

        seg_to_result_idx[seg_idx] = len(refs)
        refs.append(mfa_ref)
        audio_paths.append(audio_path)

    return refs, audio_paths, seg_to_result_idx


def compute_mfa_timestamps_api(segments, segment_dir, granularity="words"):
    """Run MFA forced alignment and return enriched segments (no UI/HTML).

    Args:
        segments: List of segment dicts (same format as alignment response).
        segment_dir: Path to directory containing per-segment WAV files.
        granularity: "words" or "words+chars".

    Returns:
        Dict with "segments" key containing enriched segment data.
    """
    if not granularity or granularity not in ("words", "words+chars"):
        granularity = "words"

    # Write individual segment WAVs on demand (sliced from full.wav)
    _ensure_segment_wavs(segments, segment_dir)

    refs, audio_paths, seg_to_result_idx = _build_mfa_refs(segments, segment_dir)
    if not refs:
        return {"segments": segments}

    results = _mfa_align_batch(refs, audio_paths)

    word_ts, letter_ts, _ = _build_timestamp_lookups(results)
    _build_crossword_groups(results, letter_ts)
    _extend_word_timestamps(word_ts, segments, seg_to_result_idx, results, segment_dir)
    return _build_enriched_json(segments, results, seg_to_result_idx,
                                word_ts, letter_ts, granularity, minimal=True)


def compute_mfa_timestamps(
    current_html,
    json_output,
    segment_dir,
    cached_log_row=None,
):
    """Animate All handler: MFA-align only uncomputed segments, then signal the mega card.

    Generator that yields (output_html, animate_all_btn, edit_patch, progress_bar, json_output)
    tuples. Skips segments whose SegmentInfo.words is already populated (from prior per-card
    or batch MFA). Progress counter reflects the uncomputed subset only.

    Preload/CDN mode: segments with pre-baked ``.words`` participate without per-segment
    WAVs because ``current_html`` arrives already stamped (see ``_stamp_prebuilt_words``).
    """
    import json as _json_mod
    import time as _time_mod
    import traceback

    # Signal-only patch: tells the JS observer "MFA is done, start the mega
    # card." The JS handler polls .segments-container for data-start before
    # building the mega card, so it doesn't matter whether Gradio's @html flush
    # for output_html lands before or after the patch arrives.
    def _make_start_patch():
        return _json_mod.dumps({
            "status": "start_megacard",
            "nonce": _time_mod.time(),
        })

    # json_output is now List[SegmentInfo] from gr.State (not a JSON dict)
    segments_state = json_output if isinstance(json_output, list) else []
    if not segments_state:
        yield current_html, gr.update(), gr.update(), gr.update(), gr.update()
        return

    # Preload (no segment_dir, no per-segment WAVs): current_html arrives
    # already stamped with data-start via _stamp_prebuilt_words at preload-load
    # time. Skip the entire MFA path — emit the patch unchanged so the JS
    # observer fires without forcing a Gradio output_html replacement (which
    # would clobber the mega card the observer is about to append).
    if not segment_dir:
        yield (
            current_html,
            gr.update(visible=True, interactive=True, variant="primary"),
            gr.update(value=_make_start_patch()),
            gr.update(visible=False),
            segments_state,
        )
        return

    # Re-render HTML from SegmentInfo to pick up any inline edits and any
    # prior MFA timestamps already attached to .words.
    from src.ui.segments import render_segments
    full_audio_url = f"/gradio_api/file={segment_dir}/full.wav"
    current_html = render_segments(
        segments_state, full_audio_url=full_audio_url, segment_dir=str(segment_dir)
    )

    if not current_html or '<span class="word"' not in current_html:
        yield current_html, gr.update(), gr.update(), gr.update(), gr.update()
        return

    # Convert to dicts at the MFA boundary (MFA internals expect dict-based segments)
    segment_dicts = [seg.to_json_dict() for seg in segments_state]

    all_refs, all_audio_paths, all_seg_to_result_idx = _build_mfa_refs(
        segment_dicts, segment_dir, require_existing=False
    )

    # Partition animatable segments into "already computed" (synthesize a
    # result from SegmentInfo.words) and "needs MFA" (batch-align). Both
    # paths feed a unified results list so inject_timestamps_into_html can
    # repaint every animatable segment in one pass after re-rendering.
    refs, audio_paths = [], []
    seg_to_result_idx = {}
    prebuilt_slots = {}
    new_batch_slots = []
    for seg_idx, old_ri in all_seg_to_result_idx.items():
        combined_idx = len(seg_to_result_idx)
        seg_to_result_idx[seg_idx] = combined_idx
        if 0 <= seg_idx < len(segments_state) and segments_state[seg_idx].words:
            prebuilt_slots[combined_idx] = {
                "status": "ok",
                "ref": all_refs[old_ri],
                "words": segments_state[seg_idx].words,
            }
        else:
            refs.append(all_refs[old_ri])
            audio_paths.append(all_audio_paths[old_ri])
            new_batch_slots.append(combined_idx)

    if not refs:
        # Everything is already timestamped — DOM still has the data-start from
        # the prior MFA run, so no per-segment text patch is needed. Empty
        # patch list just signals "start the mega card."
        yield (
            gr.update(),
            gr.update(visible=True, interactive=True, variant="primary"),
            gr.update(value=_make_start_patch()),
            gr.update(visible=False),
            segments_state,
        )
        return

    # Yield 1: hide the button so the progress bar occupies the ts-row slot.
    # Bar shows static "Preparing Alignment..." — held until the MFA Space
    # signals it has exited its queue and started actual work.
    total_segments = len(refs)
    static_bar = _ts_progress_bar_html(total_segments, MFA_PROGRESS_SEGMENT_RATE, animated=False)
    yield (
        gr.update(),
        gr.update(visible=False),
        gr.update(),
        gr.update(value=static_bar, visible=True),
        gr.update(),
    )

    # Write individual segment WAVs on demand (sliced from full.wav) after the
    # UI has flushed the preparing state, so long disk work is not invisible.
    _ensure_segment_wavs(segment_dicts, segment_dir)

    # Upload, then stream events from the MFA Space (upload happens on the
    # generator's first pull — the bar holds "Preparing..." through it). The
    # "started" event arrives the moment MFA exits its queue and actually
    # begins work — that's when we swap the static bar for the animated
    # counter. The "result" event delivers the final batch.
    batch_results = None
    try:
        for kind, data in _mfa_align_batch_streaming(refs, audio_paths):
            if kind == "started":
                animated_bar = _ts_progress_bar_html(
                    total_segments, MFA_PROGRESS_SEGMENT_RATE, animated=True)
                yield (
                    gr.update(),
                    gr.update(),
                    gr.update(),
                    gr.update(value=animated_bar),
                    gr.update(),
                )
            elif kind == "result":
                batch_results = data
    except Exception as e:
        traceback.print_exc()
        yield (
            gr.update(),
            gr.update(visible=True, interactive=True, variant="primary"),
            gr.update(),
            gr.update(visible=False),
            gr.update(),
        )
        raise

    if batch_results is None:
        raise RuntimeError("MFA align_batch did not return a result")

    # Splice synthesized "prior" results and fresh MFA results into a single
    # list aligned with seg_to_result_idx.
    total_slots = len(seg_to_result_idx)
    results = [None] * total_slots
    for combined_idx, synth in prebuilt_slots.items():
        results[combined_idx] = synth
    for i, combined_idx in enumerate(new_batch_slots):
        results[combined_idx] = batch_results[i] if i < len(batch_results) else {"status": "failed"}

    html, enriched_json = inject_timestamps_into_html(
        current_html, segment_dicts, results, seg_to_result_idx, segment_dir
    )

    # V3 note: word/char timestamps are no longer logged to the main dataset.
    # The offline `extract_timestamps.py` + Inspector flows are the authoritative
    # timestamp producers. A separate `quran-aligner-timestamps` dataset may be
    # added in v3.1 if the Space MFA path shows enough traffic to warrant it.

    # Copy MFA word/letter data back onto SegmentInfo objects
    enriched_segs = enriched_json.get("segments", []) if enriched_json else []
    for seg in segments_state:
        idx = seg.segment_number - 1
        if 0 <= idx < len(enriched_segs) and "words" in enriched_segs[idx]:
            seg.words = enriched_segs[idx]["words"]

    # Final yield: emit the freshly stamped html so Gradio @html-flushes
    # .segments-container with timestamped cards. The mega card lives in the
    # right column (sibling of the .output-html wrapper, not inside it), so it
    # survives the flush. The JS observer polls for data-start before
    # triggering the synth click — the timing of the flush vs the patch
    # doesn't matter.
    # Yield a NEW list reference for c.output_json. segments_state was mutated
    # in place (seg.words populated above), so the previous c.output_json
    # already pointed to this same list — without a fresh reference,
    # gr.State.change short-circuits on identity and the autosave hook on
    # c.output_json.change never fires. Result: pipeline_state.pkl on the
    # bucket stays without .words, so reloading + clicking Animate All again
    # re-runs MFA instead of finding everything pre-baked.
    yield (
        html,
        gr.update(visible=True, interactive=True, variant="primary"),
        gr.update(value=_make_start_patch()),
        gr.update(visible=False),
        list(segments_state),
    )