import os
import struct
from tqdm import tqdm
from collections import Counter
import hashlib

def get_top_identities_binary(rec_path, idx_path, top_n=51):
    """
    Pass 1: Scans the actual BINARY HEADERS in the .rec file.
    This is the only way to be 100% sure which image belongs to whom.
    """
    identity_counts = Counter()
    
    with open(idx_path, 'r') as f:
        offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()]

    print("Pass 1: Scanning binary headers to count identities...")
    with open(rec_path, 'rb') as f:
        for offset in tqdm(offsets):
            f.seek(offset)
            header_bin = f.read(32) # Read enough for the header
            if len(header_bin) < 32: continue
            
            # MXNet Header format: [Flag, Label (float), ID, ID]
            # The label is at offset 12 (float32)
            label = int(struct.unpack('f', header_bin[12:16])[0])
            identity_counts[label] += 1

    top_stats = identity_counts.most_common(top_n)
    top_labels = {label for label, count in top_stats}

    print(f"\nTop {top_n} Identities by Binary Label:")
    for label, count in top_stats:
        print(f"ID: {label:<10} | Count: {count:<10}")
    
    return top_labels

def extract_selected_binary(rec_path, idx_path, output_dir, top_labels):
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    with open(idx_path, 'r') as f:
        offsets = [int(line.strip().split('\t')[1]) for line in f.readlines()]

    print(f"\nPass 2: Extracting verified images...")
    
    # NEW: Keep track of how many images we've saved for each ID
    # to avoid overwriting files.
    save_counters = {label: 0 for label in top_labels}
    total_extracted = 0
    
    with open(rec_path, 'rb') as f:
        for offset in tqdm(offsets):
            f.seek(offset)
            header_bin = f.read(32)
            if len(header_bin) < 32: continue
            
            label = int(struct.unpack('f', header_bin[12:16])[0])
            
            if label not in top_labels:
                continue
            
            # Read image content
            _, length_flag = struct.unpack('II', header_bin[:8])
            content_length = length_flag & ((1 << 31) - 1)
            content = f.read(content_length)
            
            img_start = content.find(b'\xff\xd8')
            if img_start == -1: continue

            target_folder = os.path.join(output_dir, str(label))
            os.makedirs(target_folder, exist_ok=True)
            
            # Use the counter for this specific label
            current_count = save_counters[label]
            img_filename = f"{current_count}.jpg"
            img_path = os.path.join(target_folder, img_filename)
            if(current_count > 405):
                continue
            
            with open(img_path, 'wb') as img_f:
                img_f.write(content[img_start:])
            
            save_counters[label] += 1
            total_extracted += 1

    print(f"\nDone! Extracted {total_extracted} total images.")



def remove_duplicates(root_dir):
    hashes = {} # hash -> first_filepath
    duplicates_removed = 0

    # Walk through every identity folder
    for subdir, dirs, files in os.walk(root_dir):
        for filename in tqdm(files, desc=f"Checking {os.path.basename(subdir)}"):
            filepath = os.path.join(subdir, filename)
            
            # Calculate MD5 hash of the file
            with open(filepath, 'rb') as f:
                file_hash = hashlib.md5(f.read()).hexdigest()
            
            if file_hash in hashes:
                # We've seen this image before!
                os.remove(filepath)
                duplicates_removed += 1
            else:
                hashes[file_hash] = filepath

    print(f"\nClean-up complete. Removed {duplicates_removed} duplicate images.")


'''
if __name__ == "__main__":
    # Point this to your unpacked Top 50 folder
    target_dir = "./datasets/casia-set"
    remove_duplicates(target_dir)
'''
if __name__ == "__main__":
    base_dir = os.path.dirname(os.path.abspath(__file__))
    REC = os.path.join(base_dir, '../data/casia-set', 'train.rec')
    IDX = os.path.join(base_dir, '../data/casia-set', 'train.idx')
    OUT = os.path.join(base_dir, '../data/casia-set')

    # Step 1: Trust the binary, not the text file
    top_verified_labels = get_top_identities_binary(REC, IDX, top_n=50)
    
    # Step 2: Extract
    extract_selected_binary(REC, IDX, OUT, top_verified_labels)