#!/usr/bin/env python3
"""
Helper script to create audio list files from various sources.

Examples:
    # From successful downloads JSON (like from download_audio_from_s3.py)
    python create_audio_list_from_json.py \
        --input successful_downloads.json \
        --audio-dir /app2/suno/data/dpo/audios/crow \
        --output audio_list.txt
    
    # From a directory (scan for audio files)
    python create_audio_list_from_json.py \
        --scan-dir /app2/suno/data/dpo/audios/crow \
        --output audio_list.txt \
        --extensions .opus .m4a .mp3
"""

import argparse
import json
import os
from pathlib import Path
from typing import List, Set


def load_clip_ids_from_json(json_path: str) -> List[str]:
    """Load clip IDs from JSON file (list of strings)."""
    with open(json_path, "r") as f:
        clip_ids = json.load(f)
    return clip_ids


def find_audio_file(clip_id: str, audio_dir: str, extensions: List[str]) -> str | None:
    """Find audio file for clip_id in audio_dir with any of the extensions."""
    for ext in extensions:
        filepath = os.path.join(audio_dir, f"{clip_id}{ext}")
        if os.path.exists(filepath):
            return filepath
    return None


def scan_directory(directory: str, extensions: List[str]) -> List[str]:
    """Scan directory for audio files with given extensions."""
    audio_files = []
    for ext in extensions:
        audio_files.extend(Path(directory).rglob(f"*{ext}"))
    return [str(f) for f in sorted(audio_files)]


def main():
    parser = argparse.ArgumentParser(
        description="Create audio list file from various sources"
    )
    parser.add_argument(
        "--input",
        type=str,
        help="Input JSON file with clip IDs (list of strings)",
    )
    parser.add_argument(
        "--audio-dir",
        type=str,
        help="Directory containing audio files (used with --input)",
    )
    parser.add_argument(
        "--scan-dir",
        type=str,
        help="Directory to scan for audio files (alternative to --input)",
    )
    parser.add_argument(
        "--output",
        type=str,
        required=True,
        help="Output file path (text or JSON)",
    )
    parser.add_argument(
        "--extensions",
        type=str,
        nargs="+",
        default=[".opus", ".m4a", ".mp3", ".wav"],
        help="Audio file extensions to look for (default: .opus .m4a .mp3 .wav)",
    )
    parser.add_argument(
        "--format",
        type=str,
        choices=["txt", "json"],
        default="txt",
        help="Output format: txt (one path per line) or json (JSON list)",
    )
    
    args = parser.parse_args()
    
    # Validate arguments
    if args.input and args.scan_dir:
        print("Error: Cannot specify both --input and --scan-dir")
        return
    
    if not args.input and not args.scan_dir:
        print("Error: Must specify either --input or --scan-dir")
        return
    
    if args.input and not args.audio_dir:
        print("Error: --audio-dir is required when using --input")
        return
    
    # Get audio paths
    audio_paths = []
    
    if args.input:
        # Load clip IDs from JSON
        print(f"Loading clip IDs from: {args.input}")
        clip_ids = load_clip_ids_from_json(args.input)
        print(f"Loaded {len(clip_ids):,} clip IDs")
        
        # Find corresponding audio files
        print(f"Searching for audio files in: {args.audio_dir}")
        not_found = []
        for clip_id in clip_ids:
            filepath = find_audio_file(clip_id, args.audio_dir, args.extensions)
            if filepath:
                audio_paths.append(filepath)
            else:
                not_found.append(clip_id)
        
        print(f"Found {len(audio_paths):,} audio files")
        if not_found:
            print(f"Warning: {len(not_found):,} clip IDs not found")
            if len(not_found) <= 10:
                print(f"  Missing: {not_found}")
    
    else:
        # Scan directory
        print(f"Scanning directory: {args.scan_dir}")
        print(f"Extensions: {args.extensions}")
        audio_paths = scan_directory(args.scan_dir, args.extensions)
        print(f"Found {len(audio_paths):,} audio files")
    
    # Write output
    print(f"Writing to: {args.output}")
    
    if args.format == "json":
        with open(args.output, "w") as f:
            json.dump(audio_paths, f, indent=2)
    else:
        with open(args.output, "w") as f:
            for path in audio_paths:
                f.write(f"{path}\n")
    
    print(f"✓ Wrote {len(audio_paths):,} audio paths to {args.output}")


if __name__ == "__main__":
    main()