#!/usr/bin/env python3
"""
Debug script to check which shards are empty in WebDataset.
This script simulates the exact shard splitting logic used in training
and tests each shard to see if it yields any samples.

Usage:
    python3 scripts/debug_webdataset_shards.py \
        --wdsmeta-path /path/to/metadata.wdsmeta.json \
        --world-size 8 \
        --check-all-ranks
"""

import argparse
import json
import os
import sys
from pathlib import Path

# Add LMFlow to path
sys.path.insert(0, str(Path(__file__).parent.parent))

try:
    import webdataset as wds
except ImportError:
    print("ERROR: webdataset not installed. Run: pip install webdataset>=0.2.48")
    sys.exit(1)


def resolve_shard_path(shard_url, mirror_roots=None):
    """Simple path resolution (replicates logic from CustomWebDatasetMultiModalDataset)"""
    if not isinstance(shard_url, str):
        return shard_url
    
    shard_path = Path(shard_url)
    if shard_path.exists():
        return shard_path.as_posix()
    
    # Try common remappings
    replacements = [
        ("/home/ligengz/nvr_elm_llm/dataset", "/lustre/fsw/portfolios/nvr/users/ligengz/datasets"),
        ("/home/ligengz/datasets", "/lustre/fsw/portfolios/nvr/users/ligengz/datasets"),
    ]
    
    if mirror_roots:
        for old, new in mirror_roots:
            replacements.append((old, new))
    
    for old_prefix, new_prefix in replacements:
        if shard_url.startswith(old_prefix):
            candidate = Path(shard_url.replace(old_prefix, new_prefix, 1))
            if candidate.exists():
                return candidate.as_posix()
    
    return shard_url


def check_shard(shard_url, shard_idx, rank=None):
    """Check if a shard yields any samples"""
    if not Path(shard_url).exists():
        return {
            'shard_idx': shard_idx,
            'shard_url': shard_url,
            'exists': False,
            'error': 'File not found',
            'iterated_samples': 0,
            'valid_samples': 0,
        }
    
    try:
        dataset = wds.WebDataset([shard_url])
        iterated = 0
        valid = 0
        
        for sample in dataset:
            iterated += 1
            sample_id = sample.get('__key__')
            if sample_id is None:
                continue
            if ('jpg' in sample or 'png' in sample) and 'txt' in sample:
                valid += 1
            
            # Only check first 100 samples for speed
            if iterated >= 100:
                break
        
        return {
            'shard_idx': shard_idx,
            'shard_url': shard_url,
            'exists': True,
            'error': None,
            'iterated_samples': iterated,
            'valid_samples': valid,
        }
    
    except Exception as e:
        return {
            'shard_idx': shard_idx,
            'shard_url': shard_url,
            'exists': True,
            'error': str(e),
            'iterated_samples': 0,
            'valid_samples': 0,
        }


def split_shards_by_rank(all_shards, rank, world_size):
    """Replicate the shard splitting logic from CustomWebDatasetMultiModalDataset"""
    if world_size <= 1:
        return all_shards
    
    shards_per_rank = len(all_shards) // world_size
    rank_start = rank * shards_per_rank
    rank_end = (rank + 1) * shards_per_rank if rank < world_size - 1 else len(all_shards)
    return all_shards[rank_start:rank_end]


def main():
    parser = argparse.ArgumentParser(description='Debug WebDataset shards')
    parser.add_argument('--wdsmeta-path', required=True, help='Path to .wdsmeta.json file')
    parser.add_argument('--world-size', type=int, default=1, help='Number of ranks (GPUs)')
    parser.add_argument('--check-all-ranks', action='store_true', help='Check all ranks (otherwise only rank 0)')
    parser.add_argument('--max-shards-per-rank', type=int, default=5, help='Max shards to check per rank')
    parser.add_argument('--path-map', type=str, help='Path mappings in format "old=new,old2=new2"')
    
    args = parser.parse_args()
    
    # Parse path mappings
    mirror_roots = []
    if args.path_map:
        for pair in args.path_map.split(','):
            if '=' in pair:
                old, new = pair.split('=', 1)
                mirror_roots.append((old.strip(), new.strip()))
    
    # Load metadata
    meta_path = Path(args.wdsmeta_path)
    print(f"Loading metadata from: {meta_path}")
    if not meta_path.exists():
        print(f"ERROR: Metadata file not found: {meta_path}")
        sys.exit(1)
    
    with open(meta_path, 'r') as f:
        meta = json.load(f)
    
    if 'shardlist' not in meta:
        print(f"ERROR: No 'shardlist' in metadata. Keys: {list(meta.keys())}")
        sys.exit(1)
    
    # Get all shard URLs
    all_shard_urls = [shard['url'] for shard in meta['shardlist']]
    print(f"Total shards in metadata: {len(all_shard_urls)}")
    
    # Resolve paths
    all_shard_urls = [resolve_shard_path(url, mirror_roots) for url in all_shard_urls]
    print(f"First shard URL: {all_shard_urls[0]}")
    print(f"First shard exists: {Path(all_shard_urls[0]).exists()}")
    print()
    
    # Check shards for each rank
    ranks_to_check = range(args.world_size) if args.check_all_ranks else [0]
    
    for rank in ranks_to_check:
        print(f"{'='*80}")
        print(f"RANK {rank}/{args.world_size}")
        print(f"{'='*80}")
        
        # Split shards for this rank
        rank_shards = split_shards_by_rank(all_shard_urls, rank, args.world_size)
        print(f"Assigned shards: {len(rank_shards)}")
        
        if len(rank_shards) == 0:
            print(f"ERROR: Rank {rank} has NO SHARDS assigned!")
            print()
            continue
        
        # Check first N shards
        shards_to_check = rank_shards[:args.max_shards_per_rank]
        print(f"Checking first {len(shards_to_check)} shards...\n")
        
        total_iterated = 0
        total_valid = 0
        errors = []
        
        for i, shard_url in enumerate(shards_to_check):
            global_shard_idx = rank * (len(all_shard_urls) // args.world_size) + i
            print(f"  [{i}/{len(shards_to_check)}] Shard {global_shard_idx}: {Path(shard_url).name}")
            
            result = check_shard(shard_url, global_shard_idx, rank)
            
            if not result['exists']:
                print(f"    ❌ FILE NOT FOUND: {shard_url}")
                errors.append(result)
            elif result['error']:
                print(f"    ❌ ERROR: {result['error']}")
                errors.append(result)
            elif result['iterated_samples'] == 0:
                print(f"    ❌ EMPTY: No samples yielded by WebDataset")
                errors.append(result)
            elif result['valid_samples'] == 0:
                print(f"    ⚠️  WARNING: Iterated {result['iterated_samples']} samples but 0 valid (missing jpg/png/txt)")
                errors.append(result)
            else:
                print(f"    ✅ OK: {result['iterated_samples']} iterated, {result['valid_samples']} valid")
                total_iterated += result['iterated_samples']
                total_valid += result['valid_samples']
        
        print()
        print(f"Summary for Rank {rank}:")
        print(f"  Total iterated samples: {total_iterated}")
        print(f"  Total valid samples: {total_valid}")
        print(f"  Errors/warnings: {len(errors)}")
        
        if errors:
            print(f"\n  Problem shards:")
            for err in errors:
                print(f"    - Shard {err['shard_idx']}: {err.get('error', 'Empty or invalid')}")
        
        print()


if __name__ == '__main__':
    main()

