#!/usr/bin/env python3
"""
Find audio files that need encoding by comparing audio list with existing NPZ files.

This script:
1. Reads the complete audio list JSONL
2. Checks which files are missing corresponding .npz files in the output directory
3. Creates a new JSONL with only the missing files for re-encoding

Usage:
    python3 find_missing_encodings.py \
        --audio-list /home/tony/Data/Preference/crow_t1/crow_t1_audio_list_complete.jsonl \
        --encoded-dir /app2/suno/data/semantic_code/crow \
        --output /home/tony/Data/Preference/crow_t1/crow_t1_missing_encodings.jsonl
"""

import argparse
import json
import os
from typing import List, Tuple

try:
    from tqdm import tqdm
except ImportError:

    def tqdm(iterable, **kwargs):
        return iterable


def check_encoding_exists(clip_id: str, encoded_dir: str) -> bool:
    """
    Check if NPZ encoding exists for given clip ID.

    Args:
        clip_id: The clip ID to check
        encoded_dir: Directory containing encoded NPZ files

    Returns:
        True if NPZ file exists, False otherwise
    """
    npz_path = os.path.join(encoded_dir, f"{clip_id}.npz")
    return os.path.exists(npz_path)


def find_missing_encodings(
    audio_list_jsonl: str, encoded_dir: str, output_jsonl: str, verbose: bool = False
) -> Tuple[int, int, int]:
    """
    Find audio files missing NPZ encodings.

    Args:
        audio_list_jsonl: Path to complete audio list JSONL
        encoded_dir: Directory containing encoded NPZ files
        output_jsonl: Path to output JSONL with missing files
        verbose: Whether to print detailed progress

    Returns:
        Tuple of (total_files, encoded_files, missing_files)
    """
    print(f"Loading audio list from: {audio_list_jsonl}")

    # Read audio list
    audio_entries = []
    with open(audio_list_jsonl, "r") as f:
        for line in f:
            audio_entries.append(json.loads(line.strip()))

    print(f"Total audio files in list: {len(audio_entries):,}")

    # Check which files are missing encodings
    missing_entries = []
    encoded_count = 0

    print(f"\nChecking for existing encodings in: {encoded_dir}")
    for entry in tqdm(audio_entries, desc="Checking encodings"):
        clip_id = entry["id"]

        if check_encoding_exists(clip_id, encoded_dir):
            encoded_count += 1
        else:
            missing_entries.append(entry)
            if verbose:
                print(f"Missing encoding for: {clip_id}")

    # Write missing entries to output
    if missing_entries:
        print(f"\nWriting missing files list to: {output_jsonl}")
        os.makedirs(os.path.dirname(output_jsonl), exist_ok=True)

        with open(output_jsonl, "w") as f:
            for entry in missing_entries:
                f.write(json.dumps(entry) + "\n")
    else:
        print("\nAll files are already encoded! No missing files.")

    # Report statistics
    total_files = len(audio_entries)
    missing_files = len(missing_entries)

    print(f"\n{'='*60}")
    print("Encoding Status Summary:")
    print(f"{'='*60}")
    print(f"Total audio files:        {total_files:>10,}")
    print(
        f"Already encoded:          {encoded_count:>10,} ({encoded_count/total_files*100:.2f}%)"
    )
    print(
        f"Missing encodings:        {missing_files:>10,} ({missing_files/total_files*100:.2f}%)"
    )
    print(f"{'='*60}")

    if missing_files > 0:
        print(f"Missing files list: {output_jsonl}")
        print(f"\nTo encode missing files, run:")
        print(f"  sbatch submit_encode_crow_t1_missing.sh")

    return total_files, encoded_count, missing_files


def main():
    parser = argparse.ArgumentParser(
        description="Find audio files missing NPZ encodings"
    )
    parser.add_argument(
        "--audio-list",
        type=str,
        default="/home/tony/Data/Preference/crow_t1/crow_t1_audio_list_complete.jsonl",
        help="Path to complete audio list JSONL",
    )
    parser.add_argument(
        "--encoded-dir",
        type=str,
        default="/app2/suno/data/semantic_code/crow",
        help="Directory containing encoded NPZ files",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="/home/tony/Data/Preference/crow_t1/crow_t1_missing_encodings.jsonl",
        help="Path to output JSONL with missing files",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print detailed progress including each missing file",
    )

    args = parser.parse_args()

    # Validate inputs
    if not os.path.exists(args.audio_list):
        print(f"Error: Audio list not found: {args.audio_list}")
        return 1

    if not os.path.exists(args.encoded_dir):
        print(f"Error: Encoded directory not found: {args.encoded_dir}")
        return 1

    # Find missing encodings
    total, encoded, missing = find_missing_encodings(
        args.audio_list, args.encoded_dir, args.output, args.verbose
    )

    return 0


if __name__ == "__main__":
    exit(main())