#!/usr/bin/env python3
"""
Parallel script to update sfx_metas with opus file paths.
Checks for files in opus_audio_truncated (preferred) or opus_audio directories.
"""

import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm
from suno_utils.utils.text import read_jsonl, write_jsonl


def check_and_update_meta(meta, opus_path, opus_truncated_path):
    """
    Check if opus file exists and update meta with the filepath.
    
    Args:
        meta: Dictionary containing metadata including 'id'
        opus_path: Path to opus_audio directory
        opus_truncated_path: Path to opus_audio_truncated directory
    
    Returns:
        tuple: (updated_meta or None, song_id if missing else None)
    """
    song_id = meta["id"]
    opus_filename = song_id + ".opus"
    
    opus_cand = os.path.join(opus_path, opus_filename)
    trunc_cand = os.path.join(opus_truncated_path, opus_filename)
    
    # Prefer truncated version first
    if os.path.exists(trunc_cand):
        meta["s3_filepath"] = trunc_cand
        return meta, None
    elif os.path.exists(opus_cand):
        meta["s3_filepath"] = opus_cand
        return meta, None
    else:
        return None, song_id  # Return None for missing files


def main():
    # Configuration
    input_jsonl = "/app2/suno/data/sara/sfx_get_beats/combined_v3_w_extreme_metas_v0.jsonl"
    output_jsonl = "/app2/suno/data/sara/sfx_get_beats/combined_v3_w_extreme_metas_v0_updated.jsonl"
    opus_path = "/app2/suno/data/sara/sfx_get_beats/opus_audio"
    opus_truncated_path = "/app2/suno/data/sara/sfx_get_beats/opus_audio_truncated"
    missing_output = "/app2/suno/data/sara/sfx_get_beats/missing_files.txt"
    
    # Number of worker threads
    max_workers = 32  # Adjust based on your system
    
    print(f"Reading metadata from {input_jsonl}...")
    sfx_metas = read_jsonl(input_jsonl)
    print(f"Loaded {len(sfx_metas)} metadata entries")
    
    # Process in parallel
    updated_metas = []
    missing = []
    
    print(f"Processing with {max_workers} workers...")
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # Submit all tasks
        future_to_meta = {
            executor.submit(check_and_update_meta, meta, opus_path, opus_truncated_path): meta
            for meta in sfx_metas
        }
        
        # Collect results with progress bar
        for future in tqdm(as_completed(future_to_meta), total=len(sfx_metas)):
            try:
                updated_meta, missing_id = future.result()
                if updated_meta is not None:
                    updated_metas.append(updated_meta)
                if missing_id is not None:
                    missing.append(missing_id)
            except Exception as e:
                original_meta = future_to_meta[future]
                print(f"Error processing {original_meta.get('id', 'unknown')}: {e}")
    
    print(f"\nProcessing complete!")
    print(f"Total entries processed: {len(sfx_metas)}")
    print(f"Missing files: {len(missing)}")
    print(f"Found files (written to output): {len(updated_metas)}")
    
    # Save updated metadata
    print(f"\nSaving updated metadata to {output_jsonl}...")
    write_jsonl(updated_metas, output_jsonl)
    
    # Save missing file list
    if missing:
        print(f"Saving missing file list to {missing_output}...")
        with open(missing_output, 'w') as f:
            for song_id in missing:
                f.write(f"{song_id}\n")
    
    print("Done!")


if __name__ == "__main__":
    main()