app.py

"""EchoScript v1.0 UI.

Implements the frozen v1.0 workflow:

    Upload Audio -> Generate Canonical Transcript -> Preview Results
                  -> Generate Outputs -> Copy / Download

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, and so
this module stays import-safe in environments without network access.

Translation languages on offer depend on whether the person supplies their
own Anthropic API key for that session:

- No key: translation runs on the local, offline MarianMT backend, so only
  the languages it can reliably reach are offered (services.translation.
  MARIAN_TARGET_LANGUAGES).
- Key supplied: translation runs through Claude, which has no missing-
  language-pair problem, so the full language list is offered (services.
  translation.ANTHROPIC_TARGET_LANGUAGES). The key is used only for the
  request(s) made during this session and is never written to disk,
  logged, or cached anywhere server-side.
"""

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,
    MARIAN_TARGET_LANGUAGES,
    TranslationError,
    TranslationService,
)

# ---------------------------------------------------------------------------
# Lazy service singletons
#
# Safe to share across requests/users: TranscriptionService holds no
# per-request state, and TranslationService resolves its backend (and
# takes the API key, if any) fresh on every translate() call rather than
# storing it -- see services/translation.py.
# ---------------------------------------------------------------------------

_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())

# "Outputs" checkbox labels, for each of the two language sets. The
# Anthropic set is a superset of the Marian one, so switching a key in/out
# only ever adds or removes options -- it never renames existing ones.
_MARIAN_OUTPUT_CHOICES = ["Transcript"] + [f"{name} Translation" for name in MARIAN_TARGET_LANGUAGES.values()]
_ANTHROPIC_OUTPUT_CHOICES = ["Transcript"] + [f"{name} Translation" for name in ANTHROPIC_TARGET_LANGUAGES.values()]
DEFAULT_OUTPUTS = ["Transcript", "English Translation"]

# 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()
}

# Tabs 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_TAB_ORDER = list(_TRANSLATION_LABEL_TO_CODE.keys())


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)


def _update_output_choices(api_key: str, current_value: list[str]):
    """Swap the Outputs checkbox list when the API key field changes.

    With a key: offer the full Anthropic language list. Without: restrict
    to what the offline MarianMT backend can reliably translate. Any
    currently-selected language that's no longer offered is dropped from
    the selection (this only happens when a key is removed).
    """
    has_key = bool((api_key or "").strip())
    choices = _ANTHROPIC_OUTPUT_CHOICES if has_key else _MARIAN_OUTPUT_CHOICES
    filtered_value = [v for v in (current_value or []) if v in choices] or ["Transcript"]
    return gr.update(choices=choices, value=filtered_value)


# ---------------------------------------------------------------------------
# Main processing callback
# ---------------------------------------------------------------------------

def process_audio(
    audio_path: Optional[str],
    start_value: str,
    end_value: str,
    source_language_label: str,
    selected_outputs: list[str],
    api_key: str,
):
    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

    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

    language_code = _NAME_TO_CODE.get(source_language_label)  # None = auto-detect

    # Step 1: Audio -> canonical Transcript. This is the only step that
    # touches the audio; everything below works off `transcript`.
    transcript: Transcript = get_transcription_service().transcribe(
        working_path,
        source_filename=Path(audio_path).name,
        language=language_code,
        window_start=start,
        window_end=end,
    )

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

    # --- Results dashboard ---
    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 tab (always computed -- it's the source of truth --
    # but only exposed as a tab if the user kept "Transcript" checked) ---
    transcript_text = transcript.text
    transcript_file = _write_text_file(transcript_text, tmp_dir, "transcript.txt")
    transcript_tab_visible = "Transcript" in selected_outputs

    # --- Translation tabs ---
    translation_service = get_translation_service()
    translation_updates = {}  # label -> (text, file_path)
    for label in _TRANSLATION_TAB_ORDER:
        if label not in selected_outputs:
            translation_updates[label] = (None, None)
            continue

        code = _TRANSLATION_LABEL_TO_CODE[label]
        try:
            translation = translation_service.translate(transcript, code, api_key=api_key)
            text = translation.text
            file_path = _write_text_file(text, tmp_dir, f"{code}.txt")
        except TranslationError as exc:
            # Surface the failure in that tab rather than aborting every
            # other output that already succeeded.
            text = f"\u26a0\ufe0f Translation failed: {exc}"
            file_path = None
        translation_updates[label] = (text, file_path)

    # --- Subtitles (always derived from the transcript, the canonical
    # source of truth -- never regenerated from audio) ---
    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")

    outputs = [dashboard_md, transcript_text, transcript_file, gr.update(visible=transcript_tab_visible)]

    for label in _TRANSLATION_TAB_ORDER:
        text, file_path = translation_updates[label]
        visible = text is not None
        outputs += [gr.update(visible=visible), text or "", file_path]

    outputs += [srt_path, vtt_path]

    return outputs


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

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

    gr.Markdown(
        """
# EchoScript

**Upload Audio &rarr; Generate Canonical Transcript &rarr; Preview Results &rarr; Generate Outputs &rarr; Copy / Download**
        """
    )

    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",
            )

            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."
                ),
            )

            outputs_input = gr.CheckboxGroup(
                choices=_MARIAN_OUTPUT_CHOICES,
                value=DEFAULT_OUTPUTS,
                label="Outputs",
            )

            api_key_input.change(
                fn=_update_output_choices,
                inputs=[api_key_input, outputs_input],
                outputs=[outputs_input],
            )

            process_button = gr.Button("Generate Outputs", variant="primary")

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

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

                translation_tabs = {}
                translation_boxes = {}
                translation_downloads = {}
                for label in _TRANSLATION_TAB_ORDER:
                    short_name = label.replace(" Translation", "")
                    with gr.Tab(short_name, visible=False) as tab:
                        box = gr.Textbox(
                            label=short_name,
                            lines=16,
                            interactive=True,
                            buttons=["copy"],
                        )
                        download = gr.DownloadButton("Download TXT")
                    translation_tabs[label] = tab
                    translation_boxes[label] = box
                    translation_downloads[label] = download

                with gr.Tab("Subtitles"):
                    gr.Markdown(
                        "Subtitles are generated from the transcript "
                        "(source language), so they stay in sync no matter "
                        "which translations are also generated."
                    )
                    with gr.Row():
                        srt_download = gr.DownloadButton("Download SRT")
                        vtt_download = gr.DownloadButton("Download VTT")

    # Build the flat outputs list in the exact order process_audio() returns.
    click_outputs = [dashboard_output, transcript_box, transcript_download, transcript_tab]
    for label in _TRANSLATION_TAB_ORDER:
        click_outputs += [translation_tabs[label], translation_boxes[label], translation_downloads[label]]
    click_outputs += [srt_download, vtt_download]

    process_button.click(
        fn=process_audio,
        inputs=[audio_input, start_input, end_input, language_input, outputs_input, api_key_input],
        outputs=click_outputs,
    )

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