import pickle
import gc
import os
import time
from tqdm import tqdm

SHARD_SIZE = 20
BATCH_SIZE = 10
INDEX_FILE_NAME = "./index/genius_shazam_index_dict_shard_{}.pkl"
TEMP_FILE_NAME = "./index/temp_batch_{}_shard_{}.pkl"

def batch_merge_shards(batch_size=10):
    """Merge temporary files into main shards periodically"""
    for shard_idx in range(SHARD_SIZE):
        # Load main shard
        main_file = INDEX_FILE_NAME.format(shard_idx)
        if os.path.exists(main_file):
            with open(main_file, "rb") as f:
                main_dict = pickle.load(f)
        else:
            main_dict = {}
        
        # Merge all temp files for this shard
        for batch_idx in range(batch_size):
            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
                os.remove(temp_file)
                del temp_dict
        
        # Save merged result
        with open(main_file, "wb") as f:
            pickle.dump(main_dict, f)
        
        del main_dict
        gc.collect()

def dump_dict_to_temp_files(file_name, batch_index):
    """Save to temporary files instead of merging immediately"""
    index_dicts = [{} for _ in range(SHARD_SIZE)]
    
    with open(file_name, "rb") as f:
        index_dict = pickle.load(f)
    
    for key, value in tqdm(index_dict.items(), f"Sharding file {file_name}..."):
        shard_index = hash(key) % SHARD_SIZE
        index_dicts[shard_index][key] = value
    
    del index_dict
    gc.collect()
    
    # Save to temporary files
    for i in range(SHARD_SIZE):
        if index_dicts[i]:
            temp_file = TEMP_FILE_NAME.format(batch_index, i)
            with open(temp_file, "wb") as f:
                pickle.dump(index_dicts[i], f)
        index_dicts[i] = None
    
    gc.collect()

if __name__ == "__main__":
    # Process files in batches
    from multiprocessing import Pool

    def process_file(i):
        print(f"\nProcessing file {i}/160")
        file = f"./data/genius_shazam_{i}.pkl"
        
        start_time = time.time()
        dump_dict_to_temp_files(file, i)
        
        # Merge every BATCH_SIZE files
        # if (i + 1) % BATCH_SIZE == 0:
        #     print(f"Merging batch {i//BATCH_SIZE}")
        #     batch_merge_shards(BATCH_SIZE)
        
        end_time = time.time()
        print(f"Time taken: {end_time - start_time:.2f} seconds")

    with Pool(10) as pool:
        pool.map(process_file, range(160))