import json
import random
import os
import shutil
from pathlib import Path

# --- CONFIG ---
input_manifest = "thai_processed_data/manifest_clustered.jsonl"
output_dir = "check_clusters"  # Where we will copy files for you to listen
samples_to_check = 5           # How many groups to check
files_per_group = 3            # How many audio files per group to copy
# --------------

def main():
    print(f"Reading {input_manifest}...")
    
    # 1. Load and Group Data
    clusters = {}
    with open(input_manifest, 'r', encoding='utf-8') as f:
        for line in f:
            data = json.loads(line)
            speaker_id = data['speaker']
            audio_path = data['audio_path']
            
            if speaker_id not in clusters:
                clusters[speaker_id] = []
            clusters[speaker_id].append(audio_path)

    # Filter for groups that actually have pairs (size > 1)
    valid_clusters = [k for k, v in clusters.items() if len(v) >= 2]
    
    print(f"Found {len(clusters)} total groups.")
    print(f"Found {len(valid_clusters)} groups with at least 2 files (valid for pairing).")

    if not valid_clusters:
        print("Error: No valid clusters found! Did you run the clustering script?")
        return

    # 2. Pick Random Groups to Audit
    selected_speakers = random.sample(valid_clusters, min(samples_to_check, len(valid_clusters)))

    # 3. Copy Files for Manual Checking
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)

    print(f"\n--- GENERATING SAMPLES IN '{output_dir}' ---")
    
    for i, spk_id in enumerate(selected_speakers):
        # Create a subfolder for this speaker
        folder_name = f"Group_{i+1}_{spk_id}"
        folder_path = os.path.join(output_dir, folder_name)
        os.makedirs(folder_path)
        
        # Get random files from this speaker
        files = clusters[spk_id]
        selected_files = random.sample(files, min(files_per_group, len(files)))
        
        print(f"\nGroup {i+1} (ID: {spk_id}):")
        
        for src_path in selected_files:
            # Handle path logic (in case paths are relative)
            if not os.path.exists(src_path):
                # Try prepending 'thaitts_emilia_dataset' or common roots if not found
                # But usually manifest paths are relative to where you ran preprocess
                print(f"  [WARNING] File not found: {src_path}")
                continue
                
            file_name = os.path.basename(src_path)
            dst_path = os.path.join(folder_path, file_name)
            shutil.copy2(src_path, dst_path)
            print(f"  - Copied: {file_name}")

    print(f"\nDONE! Go open the folder '{output_dir}' and listen to the files.")
    print("If files in 'Group_1' sound like the same person, the clustering worked!")

if __name__ == "__main__":
    main()