main.py

from __future__ import annotations

import json
import os
import sys
from pathlib import Path
from typing import Optional

import rich_click as click  # noqa: F401 — patches click for rich help output
import typer
from dotenv import load_dotenv
from rich.console import Console
from typing_extensions import Annotated

from transcriber.download import download_video, sanitize_filename
from transcriber.embed import embed_subtitles, extract_audio
from transcriber.subtitles import Cue, Word, build_cues, format_srt
from transcriber.transcribe import transcribe_audio

app = typer.Typer(help="Download a YouTube video and generate embedded subtitles via ElevenLabs.")
console = Console()


@app.command()
def main(
    url: Annotated[Optional[str], typer.Argument(help="YouTube URL to download and transcribe")] = None,
    input_video: Annotated[
        Optional[Path],
        typer.Option("--input-video", "-i", help="Path to an already-downloaded video file (skips download)"),
    ] = None,
    output_dir: Annotated[
        Optional[Path],
        typer.Option("--output-dir", "-o", help="Directory for output files (default: $OUTPUT_DIR, else current directory)"),
    ] = None,
    language: Annotated[
        Optional[str],
        typer.Option("--language", "-l", help="ISO 639-3 language code, e.g. 'eng', 'jpn' (overrides $VIDEO_LANGUAGE)"),
    ] = None,
    resolution: Annotated[
        int,
        typer.Option("--resolution", "-r", help="Preferred video resolution (smallest side, in pixels)"),
    ] = 720,
    max_length: Annotated[
        int,
        typer.Option("--max-length", help="Max characters per subtitle cue"),
    ] = 100,
    max_pause: Annotated[
        float,
        typer.Option("--max-pause", help="Silence gap in seconds that forces a new subtitle cue"),
    ] = 3.0,
    translate_to: Annotated[
        Optional[str],
        typer.Option("--translate-to", help="ISO 639-3 target language for a translated subtitle track, e.g. 'eng' (overrides $TRANSLATE_TO). Requires $TRANSLATION_MODEL."),
    ] = None,
    translation_context: Annotated[
        int,
        typer.Option("--translation-context", help="Number of surrounding cues used as context when translating"),
    ] = 3,
) -> None:
    if not url and not input_video:
        console.print("[red]Error:[/red] Provide either a YouTube URL or [bold]--input-video[/bold].")
        raise typer.Exit(1)

    load_dotenv()

    api_key = os.getenv("ELEVENLABS_API_KEY")
    if not api_key:
        console.print("[red]Error:[/red] ELEVENLABS_API_KEY is not set")
        raise typer.Exit(1)

    if not language:
        language = os.getenv("VIDEO_LANGUAGE") or None

    if not translate_to:
        translate_to = os.getenv("TRANSLATE_TO") or None

    if output_dir is None:
        output_dir = Path(os.getenv("OUTPUT_DIR") or ".")

    output_dir = output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    # ── Resolve source video and output stem ─────────────────────────────────
    try:
        if input_video:
            source_video = input_video.resolve()
            if not source_video.exists():
                console.print(f"[red]Error:[/red] File not found: {source_video}")
                raise typer.Exit(1)
            output_stem = source_video.stem
            console.print(f"Using local file: [cyan]{source_video.name}[/cyan]")
        else:
            console.print(f"Downloading: [cyan]{url}[/cyan]")
            source_video, title = download_video(url, output_dir, resolution, console)
            output_stem = sanitize_filename(title)
            console.print(f"Downloaded:  [cyan]{source_video.name}[/cyan]")
    except Exception as e:
        console.print(f"[red]Download error:[/red] {e}")
        raise typer.Exit(1)

    json_path = output_dir / f"{output_stem}.json"
    srt_path = output_dir / f"{output_stem}.srt"
    mkv_path = output_dir / f"{output_stem}.mkv"

    if mkv_path.resolve() == source_video.resolve():
        mkv_path = output_dir / f"{output_stem}_subtitled.mkv"

    # ── Transcription (skip if cached JSON exists) ────────────────────────────
    words: list[Word]
    if json_path.exists():
        console.print(f"Using cached transcription: [cyan]{json_path.name}[/cyan]")
        data = json.loads(json_path.read_text())
        words = [Word.from_dict(w) for w in data.get("words", [])]
    else:
        mp3_path = output_dir / f"{output_stem}.mp3"
        try:
            with console.status("[bold cyan]Extracting audio...[/bold cyan]"):
                extract_audio(source_video, mp3_path)

            with console.status("[bold cyan]Transcribing with ElevenLabs (this may take a while)...[/bold cyan]"):
                response = transcribe_audio(mp3_path, api_key, language)

            json_path.write_text(response.model_dump_json(indent=2))
            console.print(f"Saved transcription: [cyan]{json_path.name}[/cyan]")
            words = [Word.from_api(w) for w in response.words]
        except Exception as e:
            console.print(f"[red]Transcription error:[/red] {e}")
            raise typer.Exit(1)
        finally:
            if mp3_path.exists():
                mp3_path.unlink()

    if not words:
        console.print("[red]Error:[/red] No words in transcription — nothing to subtitle.")
        raise typer.Exit(1)

    # ── Build subtitle cues ───────────────────────────────────────────────────
    with console.status("[bold cyan]Building subtitles...[/bold cyan]"):
        cues = build_cues(words, max_chars=max_length, max_pause=max_pause)
    console.print(f"Generated [cyan]{len(cues)}[/cyan] subtitle cues")

    # ── Translate (skip if no target language) ────────────────────────────────
    translated_srt_path: Optional[Path] = None
    if translate_to:
        translation_model = os.getenv("TRANSLATION_MODEL")
        if not translation_model:
            console.print("[red]Error:[/red] TRANSLATION_MODEL is not set (required for translation)")
            raise typer.Exit(1)

        translation_cache_path = output_dir / f"{output_stem}.{translate_to}.json"
        translated_srt_path = output_dir / f"{output_stem}.{translate_to}.srt"

        translated_cues: list[Cue]
        if translation_cache_path.exists():
            console.print(f"Using cached translation: [cyan]{translation_cache_path.name}[/cyan]")
            raw = json.loads(translation_cache_path.read_text())
            translated_cues = [Cue.from_dict(d) for d in raw]
        else:
            try:
                from transcriber.translate import translate_cues
                with console.status(f"[bold cyan]Translating subtitles to {translate_to}...[/bold cyan]"):
                    translated_cues = translate_cues(
                        cues, translate_to, translation_model, translation_context, language
                    )
            except Exception as e:
                console.print(f"[red]Translation error:[/red] {e}")
                raise typer.Exit(1)

            translation_cache_path.write_text(
                json.dumps(
                    [{"start": c.start, "end": c.end, "text": c.text} for c in translated_cues],
                    indent=2,
                    ensure_ascii=False,
                ),
                encoding="utf-8",
            )
            console.print(f"Saved translation cache: [cyan]{translation_cache_path.name}[/cyan]")

        translated_srt_path.write_text(format_srt(translated_cues), encoding="utf-8")

    # ── Write original SRT ────────────────────────────────────────────────────
    srt_path.write_text(format_srt(cues), encoding="utf-8")

    # ── Embed into MKV ────────────────────────────────────────────────────────
    try:
        with console.status("[bold cyan]Embedding subtitles into MKV...[/bold cyan]"):
            embed_subtitles(
                source_video, srt_path, mkv_path, language,
                translated_srt_path=translated_srt_path,
                target_language=translate_to,
            )
    except Exception as e:
        console.print(f"[red]Embedding error:[/red] {e}")
        raise typer.Exit(1)
    finally:
        if srt_path.exists():
            srt_path.unlink()
        if translated_srt_path and translated_srt_path.exists():
            translated_srt_path.unlink()

    console.print(f"\n[bold green]Done![/bold green]  [cyan]{mkv_path}[/cyan]")


if __name__ == "__main__":
    app()