#!/usr/bin/env python3
"""
Fetch Apollo audio recordings from Archive.org.

Downloads the Apollo 11 audio highlights compilation and extracts individual
clips using ffmpeg. Source material is from the Internet Archive's Apollo 11
collection.

Clips are saved as 48 kHz mono WAV files in examples/audio/ for use with
the gr-apollo signal processing demos.

Usage:
    uv run python examples/fetch_apollo_audio.py --list
    uv run python examples/fetch_apollo_audio.py --all
    uv run python examples/fetch_apollo_audio.py --clip eagle_has_landed
    uv run python examples/fetch_apollo_audio.py --clip liftoff --force
"""

import argparse
import os
import shutil
import subprocess
import sys
import urllib.request

# Output directory: examples/audio/ relative to this script
AUDIO_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "audio")

# Source FLAC from the Internet Archive
FLAC_URL = "https://archive.org/download/Apollo11AudioHighlights/Apollo11Highlights.flac"
FLAC_FILENAME = "Apollo11Highlights.flac"

# Clip definitions -- timestamps are approximate offsets into the highlights reel.
# The important thing is having a working extraction pipeline; timestamps can be
# refined once someone listens through the actual source file.
CLIPS = {
    "liftoff": {
        "start": "00:00:05",
        "duration": "00:00:30",
        "description": "Apollo 11 liftoff",
    },
    "eagle_has_landed": {
        "start": "00:06:45",
        "duration": "00:00:30",
        "description": "The Eagle has landed",
    },
    "one_small_step": {
        "start": "00:15:30",
        "duration": "00:00:25",
        "description": "One small step for man",
    },
    "houston_problem": {
        "start": "00:20:00",
        "duration": "00:00:15",
        "description": "Houston, we've had a problem",
    },
    "splashdown": {
        "start": "00:42:00",
        "duration": "00:00:20",
        "description": "Splashdown",
    },
}


def check_ffmpeg():
    """Verify ffmpeg is available on PATH."""
    if shutil.which("ffmpeg") is None:
        print("ERROR: ffmpeg not found on PATH.", file=sys.stderr)
        print("Install it with your package manager:", file=sys.stderr)
        print("  Arch:   pacman -S ffmpeg", file=sys.stderr)
        print("  Debian: apt install ffmpeg", file=sys.stderr)
        print("  macOS:  brew install ffmpeg", file=sys.stderr)
        sys.exit(1)


def _progress_hook(block_num, block_size, total_size):
    """Report download progress to stderr."""
    downloaded = block_num * block_size
    if total_size > 0:
        pct = min(100.0, downloaded * 100.0 / total_size)
        mb_down = downloaded / (1024 * 1024)
        mb_total = total_size / (1024 * 1024)
        bar_width = 40
        filled = int(bar_width * pct / 100.0)
        bar = "#" * filled + "-" * (bar_width - filled)
        sys.stderr.write(f"\r  [{bar}] {pct:5.1f}%  {mb_down:.1f}/{mb_total:.1f} MB")
        sys.stderr.flush()
    else:
        mb_down = downloaded / (1024 * 1024)
        sys.stderr.write(f"\r  Downloaded {mb_down:.1f} MB (unknown total)")
        sys.stderr.flush()


def download_flac(output_dir, force=False):
    """Download the FLAC source file with progress reporting.

    Returns the path to the downloaded file, or None on failure.
    """
    os.makedirs(output_dir, exist_ok=True)
    flac_path = os.path.join(output_dir, FLAC_FILENAME)

    if os.path.exists(flac_path) and not force:
        size_mb = os.path.getsize(flac_path) / (1024 * 1024)
        print(f"  FLAC already exists: {flac_path} ({size_mb:.1f} MB)")
        print("  Use --force to re-download.")
        return flac_path

    print(f"  Downloading: {FLAC_URL}")
    print(f"  Saving to:   {flac_path}")
    print()

    try:
        urllib.request.urlretrieve(FLAC_URL, flac_path, reporthook=_progress_hook)
        sys.stderr.write("\n")
        sys.stderr.flush()
    except (urllib.error.URLError, OSError) as exc:
        print(f"\n  Download failed: {exc}", file=sys.stderr)
        # Clean up partial file
        if os.path.exists(flac_path):
            os.remove(flac_path)
        return None

    size_mb = os.path.getsize(flac_path) / (1024 * 1024)
    print(f"  Downloaded {size_mb:.1f} MB")
    return flac_path


def extract_clip(flac_path, clip_name, clip_info, output_dir, force=False):
    """Extract a clip segment from the FLAC source using ffmpeg.

    Outputs a 48 kHz mono WAV file.
    Returns True on success, False on failure.
    """
    out_path = os.path.join(output_dir, f"apollo11_{clip_name}.wav")

    if os.path.exists(out_path) and not force:
        print(f"  [{clip_name}] Already exists: {out_path}")
        return True

    cmd = [
        "ffmpeg",
        "-y",                        # overwrite without asking
        "-ss", clip_info["start"],   # seek to start
        "-t", clip_info["duration"], # extract duration
        "-i", flac_path,             # input
        "-ac", "1",                  # mono
        "-ar", "48000",              # 48 kHz
        "-sample_fmt", "s16",        # 16-bit
        out_path,
    ]

    print(f"  [{clip_name}] Extracting: {clip_info['description']}")
    print(f"    start={clip_info['start']}  duration={clip_info['duration']}")

    result = subprocess.run(cmd, capture_output=True)

    if result.returncode != 0:
        print(f"  [{clip_name}] ffmpeg failed (exit {result.returncode}):", file=sys.stderr)
        stderr_text = result.stderr.decode("utf-8", errors="replace")
        # Print last few lines of ffmpeg output for diagnostics
        for line in stderr_text.strip().splitlines()[-5:]:
            print(f"    {line}", file=sys.stderr)
        return False

    size_kb = os.path.getsize(out_path) / 1024
    print(f"    -> {out_path} ({size_kb:.0f} KB)")
    return True


def list_clips():
    """Print available clip names and descriptions."""
    print("Available clips:")
    print()
    max_name = max(len(n) for n in CLIPS)
    for name, info in CLIPS.items():
        print(f"  {name:<{max_name}}  {info['start']} ({info['duration']})  {info['description']}")
    print()
    print(f"  {len(CLIPS)} clips defined.")
    print("  Extract with: --clip NAME  or  --all")


def main():
    parser = argparse.ArgumentParser(
        description="Fetch Apollo 11 audio from Archive.org and extract clips.",
        epilog="Clips are saved as 48 kHz mono WAV in examples/audio/.",
    )
    parser.add_argument(
        "--list",
        action="store_true",
        help="List available clip names and timestamps",
    )
    parser.add_argument(
        "--clip",
        metavar="NAME",
        help="Extract a specific clip by name",
    )
    parser.add_argument(
        "--all",
        action="store_true",
        help="Extract all defined clips",
    )
    parser.add_argument(
        "--keep-flac",
        action="store_true",
        help="Keep the downloaded FLAC file after extraction",
    )
    parser.add_argument(
        "--force",
        action="store_true",
        help="Re-download and re-extract even if files already exist",
    )
    parser.add_argument(
        "--output-dir",
        default=AUDIO_DIR,
        help=f"Output directory (default: {AUDIO_DIR})",
    )
    args = parser.parse_args()

    # --list doesn't need ffmpeg
    if args.list:
        list_clips()
        return

    # Everything else requires ffmpeg
    check_ffmpeg()

    # Validate arguments
    if not args.clip and not args.all:
        parser.print_help()
        print()
        print("Specify --clip NAME, --all, or --list.")
        sys.exit(1)

    if args.clip and args.clip not in CLIPS:
        print(f"Unknown clip: {args.clip}", file=sys.stderr)
        print(f"Available: {', '.join(CLIPS.keys())}", file=sys.stderr)
        sys.exit(1)

    # Determine which clips to extract
    clip_names = list(CLIPS.keys()) if args.all else [args.clip]

    print("=" * 60)
    print("Apollo 11 Audio Fetch")
    print("=" * 60)
    print()

    # Download the source FLAC
    print("Step 1: Download source FLAC")
    flac_path = download_flac(args.output_dir, force=args.force)
    if flac_path is None:
        sys.exit(1)
    print()

    # Extract clips
    print(f"Step 2: Extract {len(clip_names)} clip(s)")
    print()

    ok_count = 0
    fail_count = 0
    for name in clip_names:
        success = extract_clip(flac_path, name, CLIPS[name], args.output_dir, force=args.force)
        if success:
            ok_count += 1
        else:
            fail_count += 1
    print()

    # Clean up FLAC unless --keep-flac
    if not args.keep_flac and os.path.exists(flac_path):
        size_mb = os.path.getsize(flac_path) / (1024 * 1024)
        os.remove(flac_path)
        print(f"Removed source FLAC ({size_mb:.1f} MB). Use --keep-flac to retain.")
    elif args.keep_flac and os.path.exists(flac_path):
        print(f"Kept source FLAC: {flac_path}")
    print()

    # Summary
    print("=" * 60)
    print(f"  Extracted: {ok_count}  Failed: {fail_count}")
    if ok_count > 0:
        print(f"  Output:    {args.output_dir}/apollo11_*.wav")
    print("=" * 60)

    if fail_count > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()