#!/usr/bin/env python3
"""
Transcribe audio recordings using Whisper models via faster-whisper.

For Chromebook (CPU-only): Uses 'tiny' or 'base' models which run reasonably fast.
For GPU (3090): Can use 'large-v3' for best quality.
"""

import argparse
import json
import os
import sys
from datetime import timedelta
from pathlib import Path


def format_timestamp(seconds: float) -> str:
    """Convert seconds to HH:MM:SS.mmm format."""
    td = timedelta(seconds=seconds)
    hours, remainder = divmod(td.seconds, 3600)
    minutes, secs = divmod(remainder, 60)
    return f"{hours:02d}:{minutes:02d}:{secs:02d}.{int(td.microseconds/1000):03d}"


def transcribe_file(model, audio_path: Path, output_dir: Path, verbose: bool = False):
    """Transcribe a single audio file and save results."""
    print(f"\nTranscribing: {audio_path.name}")

    segments, info = model.transcribe(
        str(audio_path),
        beam_size=5,
        vad_filter=True,  # Filter out non-speech sections
        vad_parameters=dict(
            min_silence_duration_ms=500,
            speech_pad_ms=200,
        ),
    )

    print(f"  Detected language: {info.language} (probability: {info.language_probability:.2f})")
    print(f"  Duration: {info.duration:.1f}s")

    # Collect all segments
    results = []
    full_text = []

    for segment in segments:
        results.append({
            "start": segment.start,
            "end": segment.end,
            "text": segment.text.strip(),
        })
        full_text.append(segment.text.strip())

        if verbose:
            print(f"  [{format_timestamp(segment.start)} -> {format_timestamp(segment.end)}] {segment.text.strip()}")

    # Save outputs
    stem = audio_path.stem

    # Save JSON with timestamps
    json_path = output_dir / f"{stem}.json"
    with open(json_path, "w") as f:
        json.dump({
            "source_file": str(audio_path),
            "language": info.language,
            "language_probability": info.language_probability,
            "duration": info.duration,
            "segments": results,
        }, f, indent=2)

    # Save plain text
    txt_path = output_dir / f"{stem}.txt"
    with open(txt_path, "w") as f:
        f.write("\n".join(full_text))

    print(f"  Saved: {json_path.name}, {txt_path.name}")
    print(f"  Found {len(results)} speech segments")

    return results


def main():
    parser = argparse.ArgumentParser(description="Transcribe audio files using Whisper")
    parser.add_argument(
        "input_dir",
        type=Path,
        nargs="?",
        default=None,
        help="Directory containing audio files",
    )
    parser.add_argument(
        "-o", "--output-dir",
        type=Path,
        default=None,
        help="Output directory for transcriptions (default: input_dir/transcriptions)",
    )
    parser.add_argument(
        "-m", "--model",
        default=os.getenv("WHISPER_MODEL", "large-v3"),
        choices=["tiny", "base", "small", "medium", "large-v3"],
        help="Whisper model size (default: large-v3, or WHISPER_MODEL env var)",
    )
    parser.add_argument(
        "--device",
        default="auto",
        choices=["auto", "cpu", "cuda"],
        help="Device to use (auto detects GPU)",
    )
    parser.add_argument(
        "-v", "--verbose",
        action="store_true",
        help="Print each segment as it's transcribed",
    )
    parser.add_argument(
        "--file",
        type=Path,
        default=None,
        help="Transcribe a single file instead of directory",
    )

    args = parser.parse_args()

    # Import here to show helpful error if not installed
    try:
        from faster_whisper import WhisperModel
    except ImportError:
        print("Error: faster-whisper not installed. Install with:")
        print("  uv add faster-whisper")
        print("  # or: pip install faster-whisper")
        sys.exit(1)

    # Determine device
    if args.device == "auto":
        import torch
        device = "cuda" if torch.cuda.is_available() else "cpu"
        compute_type = "float16" if device == "cuda" else "int8"
    elif args.device == "cuda":
        device = "cuda"
        compute_type = "float16"
    else:
        device = "cpu"
        compute_type = "int8"

    print(f"Using device: {device}, compute type: {compute_type}")
    print(f"Loading model: {args.model}")

    model = WhisperModel(args.model, device=device, compute_type=compute_type)

    # Handle single file or directory
    if args.file:
        audio_files = [args.file]
        output_dir = args.output_dir or args.file.parent / "transcriptions"
    elif args.input_dir:
        input_dir = args.input_dir
        if not input_dir.exists():
            print(f"Error: Input directory does not exist: {input_dir}")
            sys.exit(1)

        audio_files = sorted(
            p for p in input_dir.iterdir()
            if p.suffix.lower() in {".ogg", ".mp3", ".wav", ".m4a", ".flac", ".webm"}
        )
        output_dir = args.output_dir or input_dir / "transcriptions"
    else:
        print("Error: Must specify either input_dir or --file")
        sys.exit(1)

    if not audio_files:
        print("No audio files found")
        sys.exit(1)

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

    print(f"Found {len(audio_files)} audio file(s)")
    print(f"Output directory: {output_dir}")

    for audio_path in audio_files:
        try:
            transcribe_file(model, audio_path, output_dir, verbose=args.verbose)
        except Exception as e:
            print(f"  Error transcribing {audio_path.name}: {e}")

    print("\nDone!")


if __name__ == "__main__":
    main()