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

Two things make repeated clicks cheap:

1. Transcript caching. A `gr.State` holds the last Transcript together with
   the exact (audio path, time window, source language) signature that
   produced it. If "Generate Outputs" is clicked again with that signature
   unchanged -- e.g. only the Outputs checkboxes changed -- transcription
   is skipped entirely and the cached Transcript is reused.
2. Per-language translation caching. A second `gr.State` dict caches each
   Translation by language code, scoped to the current transcript
   signature. Selecting an additional language only translates that new
   language; languages already translated (even if briefly deselected and
   reselected) are reused rather than recomputed.

Available translation languages are cascaded from two things:

- Whether an Anthropic API key is present for this session (see
  services/translation.py for the offline-vs-Claude tradeoff). 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.
- The source language, explicit or detected: translating a language into
  itself isn't offered. If "Auto Detect" is used, the list is re-filtered
  once the actual language is known, after transcription.
"""

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. Per-request caching of
# results (not the services themselves) lives in gr.State, below.
# ---------------------------------------------------------------------------

_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 or source language never needs to
# add/remove components -- only which ones are visible changes.
_TRANSLATION_SECTION_ORDER = list(_TRANSLATION_LABEL_TO_CODE.keys())

DEFAULT_OUTPUTS = ["Transcript", "English Translation"]


def _compute_output_choices(has_key: bool, exclude_code: Optional[str]) -> list[str]:
    """The Outputs checkbox list, cascaded from key presence + source language.

    `exclude_code` removes "translate into the language it's already in"
    from the list -- it's the explicit source language if one was chosen,
    or the detected language once transcription has run.
    """
    pool = ANTHROPIC_TARGET_LANGUAGES if has_key else MARIAN_TARGET_LANGUAGES
    return ["Transcript"] + [
        f"{name} Translation" for code, name in pool.items() if code != exclude_code
    ]


def _on_key_or_source_change(api_key: str, source_language_label: str, current_value: list[str]):
    """Re-cascade the Outputs choices when the API key or source language changes."""
    has_key = bool((api_key or "").strip())
    exclude_code = _NAME_TO_CODE.get(source_language_label)  # None when "Auto Detect"
    choices = _compute_output_choices(has_key, exclude_code)
    filtered_value = [v for v in (current_value or []) if v in choices] or ["Transcript"]
    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)


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

    if cached_transcript is not None and cached_signature == signature:
        # Same file, same window, same forced source language as last time
        # -- the canonical Transcript hasn't changed, so skip Whisper
        # entirely and reuse it.
        transcript = cached_transcript
    else:
        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

        # Step 1: Audio -> canonical Transcript. This is the only step
        # that touches the audio; everything below works off `transcript`.
        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

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

    # Now that the actual language is known (important when "Auto Detect"
    # was used), re-cascade the Outputs choices to drop "translate into
    # the language it's already in" and drop that selection if it was
    # only chosen by default/before detection.
    has_key = bool((api_key or "").strip())
    updated_choices = _compute_output_choices(has_key, transcript.language)
    effective_outputs = [v for v in selected_outputs if v in updated_choices] or ["Transcript"]
    outputs_update = gr.update(choices=updated_choices, value=effective_outputs)

    # --- Transcript section (always computed -- it's the source of truth
    # -- but only exposed if the user kept "Transcript" checked) ---
    transcript_text = transcript.text
    transcript_file = _write_text_file(transcript_text, tmp_dir, "transcript.txt")
    transcript_visible = "Transcript" in effective_outputs

    # --- Translations: reuse anything already cached for this transcript;
    # only compute the languages newly selected since the last click. ---
    translation_service = get_translation_service()
    section_updates = {}  # label -> (visible, text, file_path)

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

        if code in cached_translations:
            text = cached_translations[code]
        else:
            try:
                translation = translation_service.translate(transcript, code, api_key=api_key)
                text = translation.text
                cached_translations[code] = text
            except TranslationError as exc:
                # Surface the failure in that section rather than aborting
                # every other output that already succeeded. Deliberately
                # not cached, so it's retried on the next click.
                text = f"\u26a0\ufe0f Translation failed: {exc}"

        file_path = None
        if not text.startswith("\u26a0\ufe0f"):
            file_path = _write_text_file(text, tmp_dir, f"{code}.txt")
        section_updates[label] = (True, 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_visible),
        outputs_update,
    ]

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

    outputs += [srt_path, vtt_path, transcript, cached_signature, cached_translations]

    return outputs


def reset_session_state():
    """Clear cached transcript/translations and the visible results."""
    cleared = [None, None, {}]
    ui_reset = [
        "Upload an audio file and click **Generate Outputs** to begin.",
        "",
        None,
        gr.update(visible=True),
        gr.update(choices=_compute_output_choices(has_key=False, exclude_code=None), value=DEFAULT_OUTPUTS),
    ]
    for _ in _TRANSLATION_SECTION_ORDER:
        ui_reset += [gr.update(visible=False), "", None]
    ui_reset += [None, None]
    return ui_reset + cleared


# ---------------------------------------------------------------------------
# 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; 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=_compute_output_choices(has_key=False, exclude_code=None),
                value=DEFAULT_OUTPUTS,
                label="Outputs",
            )

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

            process_button = gr.Button("Generate Outputs", 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 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")

                with gr.Tab("Translations"):
                    gr.Markdown(
                        "Every language selected in **Outputs** appears below at "
                        "once -- nothing is hidden behind a tab you have to click "
                        "through."
                    )
                    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), 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,
        outputs_input,
    ]
    for label in _TRANSLATION_SECTION_ORDER:
        click_outputs += [translation_groups[label], translation_boxes[label], translation_downloads[label]]
    click_outputs += [
        srt_download,
        vtt_download,
        transcript_state,
        signature_state,
        translations_state,
    ]

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

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

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