main.py

import subprocess
from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from time import perf_counter

import typer
import whisper
from rich.console import Console
from rich.panel import Panel
from rich.table import Table


class WhisperModel(str, Enum):
    TINY_EN = "tiny.en"
    TINY = "tiny"
    BASE_EN = "base.en"
    BASE = "base"
    SMALL_EN = "small.en"
    SMALL = "small"
    MEDIUM_EN = "medium.en"
    MEDIUM = "medium"
    LARGE = "large"
    TURBO = "turbo"


class WhisperTask(str, Enum):
    TRANSCRIBE = "transcribe"
    TRANSLATE = "translate"


@dataclass(frozen=True)
class RunConfig:
    output_dir: Path
    input_dir: Path
    extractor_script: Path
    language: str | None
    task: WhisperTask
    realtime: bool
    fast_decode: bool
    timestamps: bool
    model: WhisperModel


@dataclass(frozen=True)
class TranscriptionResult:
    file_name: str
    language: str
    output_file: Path
    elapsed_seconds: float


app = typer.Typer(add_completion=False)
console = Console()


def list_files_by_suffix(directory: Path, suffix: str) -> list[Path]:
    return sorted(
        path
        for path in directory.iterdir()
        if path.is_file() and path.suffix.lower() == suffix.lower()
    )


def format_timestamp(seconds: float) -> str:
    total_ms = int(seconds * 1000)
    hours, remainder = divmod(total_ms, 3_600_000)
    minutes, remainder = divmod(remainder, 60_000)
    secs, ms = divmod(remainder, 1000)
    if hours > 0:
        return f"{hours:02d}:{minutes:02d}:{secs:02d}.{ms:03d}"
    return f"{minutes:02d}:{secs:02d}.{ms:03d}"


def render_output_text(raw_result: dict, include_timestamps: bool) -> str:
    if not include_timestamps:
        return raw_result.get("text", "").strip()

    segments = raw_result.get("segments") or []
    lines: list[str] = []
    for segment in segments:
        start = format_timestamp(float(segment.get("start", 0.0)))
        end = format_timestamp(float(segment.get("end", 0.0)))
        text = str(segment.get("text", "")).strip()
        if text:
            lines.append(f"[{start} --> {end}] {text}")

    if lines:
        return "\n".join(lines)
    return raw_result.get("text", "").strip()


class TranscriptionService:
    def __init__(self, cfg: RunConfig, ui: Console) -> None:
        self.cfg = cfg
        self.ui = ui
        self.model = self._load_model()

    def _load_model(self) -> whisper.Whisper:
        with self.ui.status("[bold yellow]Loading model...[/bold yellow]", spinner="dots"):
            return whisper.load_model(self.cfg.model.value)

    def print_header(self, mode: str, source: Path | None = None) -> None:
        source_line = (
            f"Input: [green]{source}[/green]"
            if source
            else f"Input dir: [green]{self.cfg.input_dir}[/green]"
        )
        self.ui.print(
            Panel.fit(
                f"[bold cyan]Whisper Transcriber[/bold cyan]\n"
                f"Mode: [magenta]{mode}[/magenta]\n"
                f"{source_line}\n"
                f"Model: [magenta]{self.cfg.model.value}[/magenta]\n"
                f"Task: [magenta]{self.cfg.task.value}[/magenta]\n"
                f"Language: [magenta]{self.cfg.language or 'auto'}[/magenta]\n"
                f"Realtime: [magenta]{'on' if self.cfg.realtime else 'off'}[/magenta]\n"
                f"Fast decode: [magenta]{'on' if self.cfg.fast_decode else 'off'}[/magenta]\n"
                f"Timestamps: [magenta]{'on' if self.cfg.timestamps else 'off'}[/magenta]",
                border_style="cyan",
            )
        )

    def print_translation_warning_if_needed(self) -> None:
        if self.cfg.task == WhisperTask.TRANSLATE and self.cfg.model == WhisperModel.TURBO:
            self.ui.print(
                "[yellow]Warning:[/yellow] model 'turbo' is not ideal for translation. "
                "Prefer 'medium' or 'large'."
            )

    def transcribe_file(self, audio_file: Path) -> TranscriptionResult:
        start_time = perf_counter()
        if self.cfg.realtime:
            self.ui.print(
                f"[bold yellow]Streaming transcript for {audio_file.name}...[/bold yellow]"
            )
        else:
            self.ui.print(
                f"[bold yellow]Transcribing full audio: {audio_file.name}...[/bold yellow]"
            )

        transcribe_kwargs = {
            "language": self.cfg.language,
            "task": self.cfg.task.value,
            "verbose": self.cfg.realtime,
            # Avoid CPU FP16 warning/fallback path.
            "fp16": str(self.model.device) != "cpu",
        }
        if self.cfg.fast_decode:
            transcribe_kwargs.update(
                {
                    "beam_size": 1,
                    "best_of": 1,
                    "temperature": 0.0,
                }
            )

        raw_result = self.model.transcribe(str(audio_file), **transcribe_kwargs)
        transcript_text = render_output_text(
            raw_result, include_timestamps=self.cfg.timestamps
        )
        detected_language = raw_result.get("language", "unknown")

        output_file = self.cfg.output_dir / f"{audio_file.stem}.txt"
        output_file.write_text(f"{transcript_text}\n", encoding="utf-8")
        elapsed = perf_counter() - start_time

        details = Table.grid(padding=(0, 2))
        details.add_row("[bold]File[/bold]", f"[green]{audio_file.name}[/green]")
        details.add_row("[bold]Language[/bold]", f"[green]{detected_language}[/green]")
        details.add_row("[bold]Output[/bold]", f"[cyan]{output_file}[/cyan]")
        details.add_row("[bold]Elapsed[/bold]", f"[yellow]{elapsed:.2f}s[/yellow]")
        self.ui.print(Panel(details, title="Done", border_style="green"))
        self.ui.print(
            Panel(
                transcript_text,
                title=f"Transcript - {audio_file.name}",
                border_style="blue",
            )
        )

        return TranscriptionResult(
            file_name=audio_file.name,
            language=detected_language,
            output_file=output_file,
            elapsed_seconds=elapsed,
        )


class BatchTranscriber:
    def __init__(self, cfg: RunConfig, service: TranscriptionService, ui: Console) -> None:
        self.cfg = cfg
        self.service = service
        self.ui = ui

    def _convert_mov_files(self) -> None:
        mov_files = list_files_by_suffix(self.cfg.input_dir, ".mov")
        if not mov_files:
            self.ui.print("[yellow]No MOV files found for conversion.[/yellow]")
            return

        if not self.cfg.extractor_script.exists():
            self.ui.print(
                f"[red]Extractor script not found:[/red] {self.cfg.extractor_script}"
            )
            raise typer.Exit(code=1)

        self.ui.print(
            f"[cyan]Found {len(mov_files)} MOV file(s); converting to MP3 with {self.cfg.extractor_script}[/cyan]"
        )
        for mov_file in mov_files:
            mp3_file = mov_file.with_suffix(".mp3")
            if mp3_file.exists():
                self.ui.print(
                    f"[yellow]Skipping conversion:[/yellow] {mov_file.name} -> {mp3_file.name} (already exists)"
                )
                continue

            with self.ui.status(
                f"[bold yellow]Converting {mov_file.name} to MP3...[/bold yellow]",
                spinner="dots",
            ):
                try:
                    subprocess.run(
                        ["bash", str(self.cfg.extractor_script), str(mov_file)],
                        check=True,
                        capture_output=True,
                        text=True,
                    )
                except subprocess.CalledProcessError as exc:
                    self.ui.print(f"[red]Failed to convert {mov_file}[/red]")
                    if exc.stderr:
                        self.ui.print(exc.stderr.strip())
                    raise typer.Exit(code=1) from exc
            self.ui.print(f"[green]Converted:[/green] {mov_file.name} -> {mp3_file.name}")

    def run(self) -> None:
        self._convert_mov_files()
        mp3_files = list_files_by_suffix(self.cfg.input_dir, ".mp3")
        if not mp3_files:
            self.ui.print(f"[red]No MP3 files found in {self.cfg.input_dir}.[/red]")
            raise typer.Exit(code=1)

        results: list[TranscriptionResult] = []
        for mp3_file in mp3_files:
            self.ui.rule(f"[bold white]Transcribing {mp3_file.name}")
            results.append(self.service.transcribe_file(mp3_file))

        table = Table(title="Batch Results", border_style="green")
        table.add_column("File", style="cyan")
        table.add_column("Language", style="green")
        table.add_column("Output", style="blue")
        table.add_column("Elapsed", style="yellow", justify="right")
        for result in results:
            table.add_row(
                result.file_name,
                result.language,
                str(result.output_file),
                f"{result.elapsed_seconds:.2f}s",
            )
        self.ui.print(table)


@app.command()
def main(
    audio_file: Path | None = typer.Argument(
        None,
        help="Path to one MP3 file. If omitted, batch mode runs over data/input.",
    ),
    output_dir: Path = typer.Option(
        Path("data/output"),
        "--output-dir",
        help="Directory where transcript files are saved (default: data/output)",
    ),
    input_dir: Path = typer.Option(
        Path("data/input"),
        "--input-dir",
        help="Input directory used for batch mode (default: data/input)",
    ),
    extractor_script: Path = typer.Option(
        Path("scripts/extract_audio.sh"),
        "--extractor-script",
        help="Path to MOV-to-MP3 extractor script.",
    ),
    language: str | None = typer.Option(
        None,
        "--language",
        help="Optional spoken language (e.g. English, Japanese, sv).",
    ),
    task: WhisperTask = typer.Option(
        WhisperTask.TRANSCRIBE,
        "--task",
        case_sensitive=False,
        help="Use transcribe (default) or translate.",
    ),
    realtime: bool = typer.Option(
        False,
        "--realtime/--no-realtime",
        help="Stream transcript segments while processing.",
    ),
    fast_decode: bool = typer.Option(
        True,
        "--fast-decode/--no-fast-decode",
        help="Use faster decoding settings (beam_size=1, best_of=1, temperature=0).",
    ),
    timestamps: bool = typer.Option(
        False,
        "--timestamps/--no-timestamps",
        help="Include [start --> end] timestamps in saved output files.",
    ),
    model: WhisperModel = typer.Option(
        WhisperModel.TURBO,
        "--model",
        case_sensitive=True,
        help="Whisper model to use (default: turbo)",
    ),
) -> None:
    run_start = perf_counter()
    output_dir.mkdir(parents=True, exist_ok=True)
    input_dir.mkdir(parents=True, exist_ok=True)

    cfg = RunConfig(
        output_dir=output_dir,
        input_dir=input_dir,
        extractor_script=extractor_script,
        language=language,
        task=task,
        realtime=realtime,
        fast_decode=fast_decode,
        timestamps=timestamps,
        model=model,
    )

    service = TranscriptionService(cfg, console)
    service.print_translation_warning_if_needed()

    if audio_file is not None:
        if not audio_file.exists() or not audio_file.is_file():
            console.print(f"[red]Audio file not found:[/red] {audio_file}")
            raise typer.Exit(code=1)
        service.print_header("single", source=audio_file)
        service.transcribe_file(audio_file)
    else:
        service.print_header("batch")
        BatchTranscriber(cfg, service, console).run()

    total_seconds = perf_counter() - run_start
    console.print(
        Panel.fit(
            f"[bold green]Run complete[/bold green]\nTotal elapsed: [yellow]{total_seconds:.2f}s[/yellow]",
            border_style="green",
        )
    )


if __name__ == "__main__":
    app()