#!/usr/bin/env python
# coding=utf-8
"""
Inspect WebDataset tar file contents to understand the data format.
"""
import json
import tarfile
from PIL import Image
import io

wdsmeta_path = '/lustre/fsw/portfolios/nvr/users/ligengz/datasets/vila-webds-meta/--home--ligengz--nvr_elm_llm--dataset--coyo-25m-vila.wdsmeta.json'

print("Loading metadata...")
with open(wdsmeta_path) as f:
    meta = json.load(f)

first_shard = meta['shardlist'][0]['url']
print(f'First shard: {first_shard}')
print(f'Total shards: {len(meta["shardlist"])}')

# Open the tar file and inspect a few samples
with tarfile.open(first_shard, 'r') as tar:
    members = tar.getmembers()
    print(f'\nTotal files in first shard: {len(members)}')
    print(f'Expected samples: {len(members) // 3} (each sample has 3 files: .jpg, .json, .txt)')
    
    # Group files by sample ID
    samples = {}
    for m in members:
        # Extract sample ID (e.g., "2001454978075" from "2001454978075.jpg")
        sample_id = m.name.rsplit('.', 1)[0]
        ext = m.name.rsplit('.', 1)[1]
        
        if sample_id not in samples:
            samples[sample_id] = {}
        samples[sample_id][ext] = m
    
    print(f'\nFound {len(samples)} unique samples')
    
    # Inspect first 3 samples
    print('\n' + '='*80)
    print('Inspecting first 3 samples:')
    print('='*80)
    
    for i, (sample_id, files) in enumerate(list(samples.items())[:3]):
        print(f'\n--- Sample {i+1}: {sample_id} ---')
        
        # Read .jpg file
        if 'jpg' in files:
            jpg_member = files['jpg']
            jpg_file = tar.extractfile(jpg_member)
            if jpg_file:
                img_data = jpg_file.read()
                img = Image.open(io.BytesIO(img_data))
                print(f'Image: {jpg_member.name} ({jpg_member.size} bytes)')
                print(f'  Size: {img.size}, Mode: {img.mode}')
        
        # Read .json file
        if 'json' in files:
            json_member = files['json']
            json_file = tar.extractfile(json_member)
            if json_file:
                json_data = json.loads(json_file.read().decode('utf-8'))
                print(f'JSON: {json_member.name} ({json_member.size} bytes)')
                print(f'  Content: {json.dumps(json_data, indent=2, ensure_ascii=False)}')
        
        # Read .txt file
        if 'txt' in files:
            txt_member = files['txt']
            txt_file = tar.extractfile(txt_member)
            if txt_file:
                txt_data = txt_file.read().decode('utf-8')
                print(f'TXT: {txt_member.name} ({txt_member.size} bytes)')
                print(f'  Content: {txt_data[:200]}...' if len(txt_data) > 200 else f'  Content: {txt_data}')
