import json
import random
import argparse
from pathlib import Path

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True, help="Input clustered manifest")
    parser.add_argument("--train-out", type=Path, required=True, help="Output Train manifest")
    parser.add_argument("--val-out", type=Path, required=True, help="Output Val manifest")
    parser.add_argument("--val-count", type=int, default=1000, help="Number of files to hold out for validation")
    args = parser.parse_args()

    print(f"Reading {args.input}...")
    with args.input.open("r", encoding="utf-8") as f:
        lines = [line.strip() for line in f if line.strip()]

    print(f"Total entries: {len(lines)}")
    
    # Shuffle nicely
    random.seed(42)
    random.shuffle(lines)

    # Split
    val_lines = lines[:args.val_count]
    train_lines = lines[args.val_count:]

    print(f"Writing {len(train_lines)} to {args.train_out}...")
    with args.train_out.open("w", encoding="utf-8") as f:
        f.write("\n".join(train_lines))
        
    print(f"Writing {len(val_lines)} to {args.val_out}...")
    with args.val_out.open("w", encoding="utf-8") as f:
        f.write("\n".join(val_lines))
    
    print("Done! Speaker IDs are now consistent across both files.")

if __name__ == "__main__":
    main()