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.
_TRANSLATION_LABEL_TO_CODE = {
    f"{name} Translation": 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.

    `exclude_code` is the actual transcript language at this point (this
    picker only exists once a Transcript does), never a pre-detection
    guess. Without a key, the Marian options are checked for real against
    the Hub for this specific source language -- not a fixed guess -- so
    a source with no path to a given language (direct or English-pivot)
    simply won't offer it, and one with good coverage offers everything
    that's actually reachable (which can be more than four languages).
    """
    if has_key:
        pool = ANTHROPIC_TARGET_LANGUAGES
    else:
        pool = available_marian_targets(exclude_code) if exclude_code else {}
    return [f"{name} Translation" 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,
):
    if cached_transcript is None:
        raise gr.Error("Generate a transcript first.")

    translation_service = get_translation_service()
    section_updates = {}

    for label in _TRANSLATION_SECTION_ORDER:
        code = _TRANSLATION_LABEL_TO_CODE[label]
        if label not in (selected_languages or []):
            section_updates[label] = (False, "", None)
            continue

        if code in cached_translations:
            text = cached_translations[code]
        else:
            try:
                translation = translation_service.translate(cached_transcript, code, api_key=api_key)
                text = translation.text
                cached_translations[code] = text
            except TranslationError as exc:
                # Surface the failure for this language only; deliberately
                # not cached, so the next click retries it.
                text = f"\u26a0\ufe0f Translation failed: {exc}"

        file_path = None
        if not text.startswith("\u26a0\ufe0f"):
            tmp_dir = Path(tempfile.mkdtemp(prefix="echoscript_"))
            file_path = _write_text_file(text, tmp_dir, f"{code}.txt")
        section_updates[label] = (True, text, file_path)

    outputs = []
    for label in _TRANSLATION_SECTION_ORDER:
        visible, text, file_path = section_updates[label]
        outputs += [gr.update(visible=visible), text, file_path]
    outputs.append(cached_translations)
    return outputs


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-06-28 22:17 UTC &middot; fixed en-jap Helsinki-NLP naming exception</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:")
                        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],
    )

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

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