import os
import pickle
import gc
from tqdm import tqdm

SHARD_SIZE = 20
RESHARD_SIZE = 100
TOTAL_SHARDS = 160
TEMP_FILE_NAME = "./index/temp_batch_{}_shard_{}.pkl"
INDEX_FILE_NAME = "./index/genius_shazam_index_dict_shard_{}.pkl"

def batch_merge_shards():
    """Merge temporary files into main shards periodically"""
    for shard_idx in range(SHARD_SIZE):
        main_dict = {}
        main_file = INDEX_FILE_NAME.format(shard_idx)
        
        # Merge all temp files for this shard
        for batch_idx in tqdm(range(TOTAL_SHARDS), desc=f"Merging shard {shard_idx}"):
            temp_file = TEMP_FILE_NAME.format(batch_idx, shard_idx)
            if os.path.exists(temp_file):
                with open(temp_file, "rb") as f:
                    temp_dict = pickle.load(f)
                
                for key, value in temp_dict.items():
                    if key in main_dict:
                        main_dict[key].extend(value)
                    else:
                        main_dict[key] = value
                
                # Remove temp file
                del temp_dict
        
        # reindex into smaller files and save merged result
        # For shard_idx = 0, we will reindex into 5 files, 0, 20, 40, 60, 80
        for i in tqdm(range(shard_idx, RESHARD_SIZE, SHARD_SIZE), desc=f"Dump dict to shard {shard_idx}"):
            main_file = INDEX_FILE_NAME.format(i)
            tmp_dict = {k: v for k, v in main_dict.items() if k % RESHARD_SIZE == i}
            with open(main_file, "wb") as f:
                pickle.dump(tmp_dict, f)
        
        del main_dict
        gc.collect()

if __name__ == "__main__":
    batch_merge_shards()