app.py

"""EchoScript v1.0 UI.

Two-stage workflow, matching the frozen v1.0 product spec exactly:

    Upload Audio -> Select Audio Window (optional)
                  -> Detect Language & Generate Canonical Transcript
                  -> Preview Transcript + available translation languages
                  -> Generate Translations -> Copy / Download

These are two distinct user actions, not one combined click:

1. "Generate Transcript" -- audio, time window, and an optional source-
   language hint go in; a canonical Transcript comes out (detected
   language, confidence, duration, word count, full text). This is the
   only step that touches the audio. Once it's done, the translation-
   language picker appears, scoped to the language that was *actually*
   detected (or forced) -- never a pre-detection guess.
2. "Generate Translations" -- pick which languages to translate the
   transcript into (the list depends on whether an Anthropic API key is
   present: more languages with a key, the offline-safe set without one)
   and click again. Always derived from the cached Transcript, never from
   the audio.

Caching: a `gr.State` holds the last Transcript plus the exact (audio,
window, source-language) signature that produced it -- clicking "Generate
Transcript" again with that signature unchanged reuses it instead of
re-running Whisper. A second `gr.State` dict caches each Translation by
language code, scoped to the current transcript, so adding one more
language to "Generate Translations" doesn't redo the others, and
deselecting a language doesn't drop it from the cache (just hides it).

Services are instantiated lazily (on first use) rather than at import
time, so the app can start up without needing model weights on disk yet.
"""

from __future__ import annotations

import tempfile
from pathlib import Path
from typing import Optional

import gradio as gr

from models.transcript import Transcript
from services.audio import AudioError, extract_window, resolve_window, validate_extension
from services.subtitles import generate_srt, generate_vtt
from services.transcription import SUPPORTED_LANGUAGES, TranscriptionService
from services.translation import (
    ANTHROPIC_TARGET_LANGUAGES,
    TranslationError,
    TranslationService,
    available_marian_targets,
)

# ---------------------------------------------------------------------------
# Lazy service singletons -- safe to share across requests/users; see
# services/translation.py for why TranslationService holds no key state.
# ---------------------------------------------------------------------------

_transcription_service: Optional[TranscriptionService] = None
_translation_service: Optional[TranslationService] = None


def get_transcription_service() -> TranscriptionService:
    global _transcription_service
    if _transcription_service is None:
        _transcription_service = TranscriptionService(
            model_size="base",
            device="cpu",
            compute_type="int8",
            download_root="/tmp/whisper_models",
        )
    return _transcription_service


def get_translation_service() -> TranslationService:
    global _translation_service
    if _translation_service is None:
        _translation_service = TranslationService()
    return _translation_service


# ---------------------------------------------------------------------------
# UI <-> service-layer vocabulary
# ---------------------------------------------------------------------------

# "Source Language" dropdown: display name -> ISO 639-1 code (None = auto).
_NAME_TO_CODE = {name: code for code, name in SUPPORTED_LANGUAGES.items()}
SOURCE_LANGUAGE_CHOICES = ["Auto Detect"] + list(SUPPORTED_LANGUAGES.values())

# Label -> ISO 639-1 code, built from the full (Anthropic) superset so it
# resolves correctly regardless of which list is currently offered.
# Labels are just language names (e.g. "French", not "French Translation").
_TRANSLATION_LABEL_TO_CODE = {
    name: code for code, name in ANTHROPIC_TARGET_LANGUAGES.items()
}

# Sections are pre-built for every language in the full superset (hidden by
# default) so that toggling the API key never needs to add/remove
# components -- only which ones are visible changes.
_TRANSLATION_SECTION_ORDER = list(_TRANSLATION_LABEL_TO_CODE.keys())


def _compute_translation_choices(has_key: bool, exclude_code: Optional[str]) -> list[str]:
    """The translate-to picker, cascaded from key presence + source language."""
    if has_key:
        pool = ANTHROPIC_TARGET_LANGUAGES
    else:
        pool = available_marian_targets(exclude_code) if exclude_code else {}
    return [name for code, name in pool.items() if code != exclude_code]


def _on_api_key_change(api_key: str, cached_transcript: Optional[Transcript], current_value: list[str]):
    """Re-cascade the translate-to picker live as the API key field changes.

    No-op until a transcript exists -- the picker isn't shown before then,
    so there's nothing yet to cascade.
    """
    if cached_transcript is None:
        return gr.update()
    has_key = bool((api_key or "").strip())
    choices = _compute_translation_choices(has_key, cached_transcript.language)
    filtered_value = [v for v in (current_value or []) if v in choices]
    return gr.update(choices=choices, value=filtered_value)


def _format_duration(seconds: float) -> str:
    seconds = max(0, int(round(seconds)))
    hours, remainder = divmod(seconds, 3600)
    minutes, secs = divmod(remainder, 60)
    return f"{hours:02}:{minutes:02}:{secs:02}"


def _write_text_file(text: str, tmp_dir: Path, filename: str) -> str:
    path = tmp_dir / filename
    path.write_text(text, encoding="utf-8")
    return str(path)


# ---------------------------------------------------------------------------
# Stage 1: Generate Transcript
# ---------------------------------------------------------------------------

def generate_transcript(
    audio_path: Optional[str],
    start_value: str,
    end_value: str,
    source_language_label: str,
    api_key: str,
    cached_transcript: Optional[Transcript],
    cached_signature,
    cached_translations: dict,
):
    if not audio_path:
        raise gr.Error("Please upload an audio file first.")

    try:
        validate_extension(audio_path)
        start, end = resolve_window(start_value, end_value)
    except AudioError as exc:
        raise gr.Error(str(exc)) from exc

    source_code = _NAME_TO_CODE.get(source_language_label)  # None = auto-detect
    signature = (audio_path, start, end, source_code)
    regenerated = not (cached_transcript is not None and cached_signature == signature)

    if regenerated:
        working_path = audio_path
        if start is not None or end is not None:
            try:
                working_path = extract_window(audio_path, start, end)
            except AudioError as exc:
                raise gr.Error(str(exc)) from exc

        # The only step that touches the audio. Detection happens here too
        # when no source language is forced.
        transcript = get_transcription_service().transcribe(
            working_path,
            source_filename=Path(audio_path).name,
            language=source_code,
            window_start=start,
            window_end=end,
        )
        cached_signature = signature
        cached_translations = {}  # old translations were derived from a different transcript
    else:
        # Same audio, window, and source-language hint as last time --
        # skip Whisper entirely and reuse the cached Transcript.
        transcript = cached_transcript

    tmp_dir = Path(tempfile.mkdtemp(prefix="echoscript_"))

    detected_language = SUPPORTED_LANGUAGES.get(transcript.language, transcript.language)
    dashboard_md = (
        f"### \u2713 {detected_language} detected\n\n"
        f"**Confidence:** {transcript.language_probability:.0%}&nbsp;&nbsp;&nbsp;"
        f"**Duration:** {_format_duration(transcript.duration)}&nbsp;&nbsp;&nbsp;"
        f"**Words:** {transcript.word_count:,}"
    )

    transcript_text = transcript.text
    transcript_file = _write_text_file(transcript_text, tmp_dir, "transcript.txt")

    # Now that the actual language is known -- the whole point of doing
    # this as its own step -- compute the translate-to picker against it,
    # never against a pre-detection guess.
    has_key = bool((api_key or "").strip())
    translation_choices = _compute_translation_choices(has_key, transcript.language)
    default_targets = ["English Translation"] if "English Translation" in translation_choices else []
    translate_choices_update = gr.update(choices=translation_choices, value=default_targets)

    srt_path = _write_text_file(generate_srt(transcript.segments), tmp_dir, "transcript.srt")
    vtt_path = _write_text_file(generate_vtt(transcript.segments), tmp_dir, "transcript.vtt")

    # Per-language result sections: only reset (hide + clear) if the
    # transcript actually changed. If it was reused, leave whatever
    # translations are already showing exactly as they are.
    section_outputs = []
    if regenerated:
        for _ in _TRANSLATION_SECTION_ORDER:
            section_outputs += [gr.update(visible=False), "", None]
    else:
        for _ in _TRANSLATION_SECTION_ORDER:
            section_outputs += [gr.update(), gr.update(), gr.update()]

    outputs = [
        dashboard_md,
        transcript_text,
        transcript_file,
        translate_choices_update,
        gr.update(visible=True),   # reveal the "choose languages" picker
        gr.update(visible=False),  # hide the "generate a transcript first" placeholder
    ]
    outputs += section_outputs
    outputs += [srt_path, vtt_path, transcript, cached_signature, cached_translations]
    return outputs


# ---------------------------------------------------------------------------
# Stage 2: Generate Translations
# ---------------------------------------------------------------------------

def generate_translations(
    selected_languages: list[str],
    api_key: str,
    cached_transcript: Optional[Transcript],
    cached_translations: dict,
):
    """Generator: yield cached results immediately, then compute only new ones.

    This ensures the loading spinner only appears on sections that are
    actually being translated. Languages already in the cache are yielded
    instantly in the first pass; only genuinely new languages trigger
    model/API calls in the second pass. Gradio generators allow partial
    yields, so the UI updates progressively rather than waiting for the
    slowest language.
    """
    if cached_transcript is None:
        raise gr.Error("Generate a transcript first.")

    selected = set(selected_languages or [])
    translation_service = get_translation_service()
    tmp_dir = Path(tempfile.mkdtemp(prefix="echoscript_"))

    def _make_outputs(section_states: dict) -> list:
        """Build the flat output list from a dict of label -> (visible, text, file)."""
        result = []
        for label in _TRANSLATION_SECTION_ORDER:
            state = section_states.get(label)
            if state is None:
                # No decision yet for this label -- emit a no-op so Gradio
                # doesn't touch it (preserves whatever is already shown).
                result += [gr.update(), gr.update(), gr.update()]
            else:
                visible, text, file_path = state
                result += [gr.update(visible=visible), text if text is not None else gr.update(), file_path]
        result.append(cached_translations)
        return result

    # ------------------------------------------------------------------
    # Pass 1: Resolve every section immediately from the cache or by
    # hiding unselected ones. Sections that need a real translation show
    # a "⏳ Translating..." placeholder so the user sees all boxes right
    # away rather than having to wait for each one to appear.
    # ------------------------------------------------------------------
    section_states: dict[str, Optional[tuple]] = {}
    needs_translation: list[str] = []

    for label in _TRANSLATION_SECTION_ORDER:
        code = _TRANSLATION_LABEL_TO_CODE[label]
        if label not in selected:
            section_states[label] = (False, "", None)
        elif code in cached_translations:
            text = cached_translations[code]
            file_path = _write_text_file(text, tmp_dir, f"{code}_cached.txt")
            section_states[label] = (True, text, file_path)
        else:
            # Show the box immediately with a placeholder; fill it in pass 2.
            section_states[label] = (True, "⏳ Translating...", None)
            needs_translation.append(label)

    # Yield immediately so cached/placeholder results appear at once.
    yield _make_outputs(section_states)

    # ------------------------------------------------------------------
    # Pass 2: Translate only the languages that aren't cached yet,
    # yielding after each one completes.
    # ------------------------------------------------------------------
    for label in needs_translation:
        code = _TRANSLATION_LABEL_TO_CODE[label]
        try:
            translation = translation_service.translate(cached_transcript, code, api_key=api_key)
            text = translation.text
            cached_translations[code] = text
            file_path = _write_text_file(text, tmp_dir, f"{code}.txt")
            section_states[label] = (True, text, file_path)
        except TranslationError as exc:
            # Surface the failure for this language only. Deliberately not
            # cached, so the next click will retry.
            text = f"\u26a0\ufe0f Translation failed: {exc}"
            section_states[label] = (True, text, None)

        # Yield after each language so the UI updates progressively.
        yield _make_outputs(section_states)


def reset_session_state():
    """Clear cached transcript/translations and everything on screen."""
    ui_reset = [
        "Upload an audio file and click **Generate Transcript** to begin.",
        "",
        None,
        gr.update(choices=[], value=[]),
        gr.update(visible=False),
        gr.update(visible=True),
    ]
    for _ in _TRANSLATION_SECTION_ORDER:
        ui_reset += [gr.update(visible=False), "", None]
    ui_reset += [None, None, None, None, {}]
    return ui_reset


# ---------------------------------------------------------------------------
# UI layout
# ---------------------------------------------------------------------------

with gr.Blocks(title="EchoScript") as demo:

    transcript_state = gr.State(value=None)
    signature_state = gr.State(value=None)
    translations_state = gr.State(value={})

    gr.Markdown(
        """
# EchoScript

**Upload Audio &rarr; Select Audio Window &rarr; Detect Language & Generate Transcript
&rarr; Preview & Choose Languages &rarr; Generate Translations &rarr; Copy / Download**

<sub>build: 2026-07-02 01:21 UTC &middot; fixed fa/pt/tr model names &middot; short labels &middot; Select All</sub>
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### Upload Audio")
            audio_input = gr.Audio(
                label="Drop audio file here or click to browse",
                sources=["upload"],
                type="filepath",
            )
            gr.Markdown("Supported: mp3 &middot; wav &middot; m4a &middot; flac")

            gr.Markdown("### Processing Window")
            with gr.Row():
                start_input = gr.Textbox(label="Start Time (optional)", placeholder="HH:MM:SS")
                end_input = gr.Textbox(label="End Time (optional)", placeholder="HH:MM:SS")
            gr.Markdown("Leave blank: entire file")

            gr.Markdown("### Processing Options")
            language_input = gr.Dropdown(
                choices=SOURCE_LANGUAGE_CHOICES,
                value="Auto Detect",
                label="Source Language",
                info="A hint for transcription, not a guess at translation targets.",
            )

            api_key_input = gr.Textbox(
                label="Anthropic API Key (optional)",
                type="password",
                placeholder="sk-ant-...",
                info=(
                    "Provide your own key to translate into many more languages via Claude. "
                    "Without one, translation uses local offline models (English, German, "
                    "Persian, Spanish only). Used for this session only -- never stored."
                ),
            )

            generate_transcript_button = gr.Button("Generate Transcript", variant="primary")
            reset_button = gr.Button("Reset (clear cache)", size="sm")

        with gr.Column(scale=2):
            gr.Markdown("### Results Dashboard")
            dashboard_output = gr.Markdown("Upload an audio file and click **Generate Transcript** to begin.")

            with gr.Tabs():
                with gr.Tab("Transcript"):
                    transcript_box = gr.Textbox(
                        label="Transcript",
                        lines=16,
                        interactive=True,
                        buttons=["copy"],
                    )
                    transcript_download = gr.DownloadButton("Download TXT")

                with gr.Tab("Translations"):
                    translations_placeholder = gr.Markdown(
                        "Generate a transcript first to see the languages available to "
                        "translate it into."
                    )
                    with gr.Group(visible=False) as translations_picker_group:
                        gr.Markdown("Translate to:")
                        with gr.Row():
                            select_all_btn = gr.Button("Select All", size="sm")
                            select_none_btn = gr.Button("Deselect All", size="sm")
                        translate_choices_input = gr.CheckboxGroup(choices=[], value=[], label=None)
                        generate_translations_button = gr.Button("Generate Translations", variant="primary")

                    translation_groups = {}
                    translation_boxes = {}
                    translation_downloads = {}
                    for label in _TRANSLATION_SECTION_ORDER:
                        short_name = label.replace(" Translation", "")
                        with gr.Group(visible=False) as group:
                            box = gr.Textbox(
                                label=short_name,
                                lines=10,
                                interactive=True,
                                buttons=["copy"],
                            )
                            download = gr.DownloadButton("Download TXT")
                        translation_groups[label] = group
                        translation_boxes[label] = box
                        translation_downloads[label] = download

                with gr.Tab("Subtitles"):
                    gr.Markdown(
                        "Subtitles are generated from the transcript "
                        "(source language) as soon as it's ready -- no "
                        "translation needed."
                    )
                    with gr.Row():
                        srt_download = gr.DownloadButton("Download SRT")
                        vtt_download = gr.DownloadButton("Download VTT")

    # Outputs shared by Stage 1 (Generate Transcript) and Reset.
    transcript_stage_outputs = [
        dashboard_output,
        transcript_box,
        transcript_download,
        translate_choices_input,
        translations_picker_group,
        translations_placeholder,
    ]
    for label in _TRANSLATION_SECTION_ORDER:
        transcript_stage_outputs += [
            translation_groups[label],
            translation_boxes[label],
            translation_downloads[label],
        ]
    transcript_stage_outputs += [
        srt_download,
        vtt_download,
        transcript_state,
        signature_state,
        translations_state,
    ]

    generate_transcript_button.click(
        fn=generate_transcript,
        inputs=[
            audio_input,
            start_input,
            end_input,
            language_input,
            api_key_input,
            transcript_state,
            signature_state,
            translations_state,
        ],
        outputs=transcript_stage_outputs,
    )

    # Outputs for Stage 2 (Generate Translations): just the per-language
    # sections plus the translation cache.
    translation_stage_outputs = []
    for label in _TRANSLATION_SECTION_ORDER:
        translation_stage_outputs += [
            translation_groups[label],
            translation_boxes[label],
            translation_downloads[label],
        ]
    translation_stage_outputs += [translations_state]

    generate_translations_button.click(
        fn=generate_translations,
        inputs=[translate_choices_input, api_key_input, transcript_state, translations_state],
        outputs=translation_stage_outputs,
    )

    # Cascade the translate-to picker live as the key changes, once a
    # transcript exists (no-op before that -- the picker isn't shown yet).
    api_key_input.input(
        fn=_on_api_key_change,
        inputs=[api_key_input, transcript_state, translate_choices_input],
        outputs=[translate_choices_input],
    )
    api_key_input.change(
        fn=_on_api_key_change,
        inputs=[api_key_input, transcript_state, translate_choices_input],
        outputs=[translate_choices_input],
    )

    select_all_btn.click(
        fn=lambda choices: gr.update(value=choices),
        inputs=[translate_choices_input],
        outputs=[translate_choices_input],
    )
    select_none_btn.click(
        fn=lambda: gr.update(value=[]),
        outputs=[translate_choices_input],
    )

    reset_button.click(fn=reset_session_state, outputs=transcript_stage_outputs)

if __name__ == "__main__":
    demo.launch()