import numpy as np
import os
import json
from suno_utils.audio import Audio
from util import get_json_file, get_genius_file, SAMPLE_RATE
from shazam import GPUShazam
from tqdm import tqdm 
import pickle
import argparse
import boto3

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Build a Shazam‐style index over metas and/or add/query songs."
    )
    parser.add_argument(
        "--batch_idx", type=int, default=0,
        help="start index of the metas"
    )
    parser.add_argument(
        "--size", type=int, default=5000,
        help="size of the dataset to process"
    )

    parser.add_argument(
        "--data", type=str, default="genius",
        help="Data bundle to process, 'genius' or 'youtube'"
    )
    
    args = parser.parse_args()
    # Process data type 
    if args.data not in ["genius", "youtube"]:
        raise ValueError(f"Invalid data type: {args.data}")
    get_data_metas_handler, local_data_name, index_file_prefix, audio_field_name = None, "", "", "" 
    if args.data == "genius":
        get_data_metas_handler = get_genius_file
        local_data_name = "genius_metas.jsonl"
        index_prefix = "genius_shazam_index_"
        audio_field_name = "audio_filepath"
    elif args.data == "youtube":
        get_data_metas_handler = get_json_file
        local_data_name = "youtube_metas.jsonl"
        index_prefix = "youtube_shazam_index_"
        audio_field_name = "audio_filepath"
    

    start_idx = args.batch_idx * args.size
    print(f"Processing batch {args.batch_idx} with size {args.size}")

    # If there is no genius_metas.jsonl in the current directory, download it from s3
    if not os.path.exists(local_data_name):
        print(f"Download {local_data_name} from s3")
        data_json = get_data_metas_handler()
        with open(local_data_name, 'w') as f:
            for item in data_json:
                f.write(json.dumps(item) + '\n')
    else:
        print(f"Load {local_data_name} from local file")
        with open(local_data_name, 'r') as f:
            data_json = [json.loads(line) for line in f]

    # Build the shazam index with your CLI parameters
    genius_shazam2 = GPUShazam(
        device='cpu',
        fan_value=5,
        kernal_size=5
    )

    print("Processing...")
    metas = data_json[start_idx:start_idx + args.size]
    print(f"Length of metas: {len(metas)}")

    for idx, item in tqdm(enumerate(metas), total=len(metas), desc=f"Processing batch {args.batch_idx}"):
        if not item.get(audio_field_name):
            print(f"No audio filepath given for idx {idx}, skipping")
            continue
        if not item.get("id"):
            print(f"No  id given for idx {idx}, skipping")
            continue
        audio_s3_path = item.get(audio_field_name)
        song_id = item.get("id")
        cur_audio = Audio.from_s3(audio_s3_path,sample_rate=SAMPLE_RATE)
        if cur_audio.duration_s > 6 * 60:
            continue
        try:
            genius_shazam2.add_song(song_id, cur_audio.array_float)
        except Exception as e:
            print(f"Fail to add song {song_id}, exception: {e}")

    tmp_dict_file = f"./data/{index_file_prefix}{args.batch_idx}.pkl"
    with open(tmp_dict_file, "wb") as f:
        pickle.dump(genius_shazam2.index, f)
    tmp_set_file = f"./data/{index_file_prefix}{args.batch_idx}_set.pkl"
    with open(tmp_set_file, "wb") as f:
        pickle.dump(genius_shazam2._set, f)

    # 2) Upload to S3
    s3 = boto3.client("s3")
    bucket = "suno-data"
    key = f"ashe/{args.data}_fingerprints/{index_file_prefix}{args.batch_idx}_dict.pkl"

    s3.upload_file(tmp_dict_file, bucket, key)
    print(f"Uploaded to s3://{bucket}/{key}")
    key = f"ashe/{args.data}_fingerprints/{index_file_prefix}{args.batch_idx}_set.pkl"
    s3.upload_file(tmp_set_file, bucket, key)
    print(f"Uploaded to s3://{bucket}/{key}")