import os
import boto3
import torch
import numpy as np
import torchaudio
import tempfile

from tqdm import tqdm
from suno_utils.utils.text import read_jsonl
from suno_utils.utils.s3 import read_from_s3
from concurrent.futures import ThreadPoolExecutor, as_completed


def get_s3_files(bucket_name, prefix, max_keys: int = 100000):
    all_files = []
    continuation_token = None

    while True:
        # Prepare the arguments for the request
        list_kwargs = {
            "Bucket": bucket_name,
            "Prefix": prefix,  # List objects under this prefix, or leave blank for all objects
        }

        if continuation_token:
            list_kwargs["ContinuationToken"] = continuation_token

        # Make the request to list objects
        response = s3.list_objects_v2(**list_kwargs)

        # Collect the file keys
        all_files += [obj["Key"] for obj in response.get("Contents", [])]

        # Check if more results are available
        if response.get("IsTruncated"):  # True if there are more results to fetch
            continuation_token = response["NextContinuationToken"]
        else:
            break  # No more results to fetch

    return all_files


def process_meta_id(meta_id, s3_paths, bucket_name):

    s3_path = list(s3_paths)[0]
    print(s3_path)
    full_path = f"s3://{bucket_name}/{s3_path}.mp3"
    audio_data, sr = read_from_s3(full_path, read_f=torchaudio.load)
    # apply highpass filter with torchaudio
    highpass_cutoff = np.random.uniform(500, 4000)
    audio_data_out = torchaudio.functional.highpass_biquad(
        audio_data, sr, highpass_cutoff
    )
    # save to tmp dir
    with tempfile.TemporaryDirectory() as tmp_dir:
        tmp_path = os.path.join(tmp_dir, f"{meta_id}.mp3")
        torchaudio.save(tmp_path, audio_data_out, sr)
        # upload to s3
        output_path = f"{s3_path}-highpass.mp3"
        print(output_path)

        s3.upload_file(tmp_path, bucket_name, output_path)


if __name__ == "__main__":

    bucket_name = "suno-data"
    base_dir = "christian/data/upsample_100z_v1"
    output_name = "v2"
    #
    # base_dir = "christian/data/upsample_100hz_v4_t_5_20241018"
    # output_name = "v2"

    base_metas_path = os.path.join("s3://", bucket_name, base_dir, "metas.jsonl")
    base_metas = read_from_s3(base_metas_path, read_f=read_jsonl)
    print(len(base_metas))

    # s3 client
    s3 = boto3.client("s3")

    # find all files on s3 with the pattern
    filepaths = get_s3_files(bucket_name, f"{base_dir}/{output_name}")
    print("total files on s3: ", len(filepaths))

    id_to_s3_paths = {}
    # create a dict with the id as key and the s3 paths a list of values
    for filepath in tqdm(filepaths):
        if "-text_cfg_" in filepath:
            meta_id = filepath.split("-text_cfg")[0].split("/")[-1]
        else:
            meta_id = filepath.split("/")[-1].split(".")[0].split("-")[:-1]
            meta_id = "-".join(meta_id)

        if meta_id not in id_to_s3_paths:
            id_to_s3_paths[meta_id] = set()

        s3_filepath_basename = filepath.split(".")[0]
        if (
            "highpass" not in s3_filepath_basename
            and "quality_scores" not in s3_filepath_basename
        ):
            id_to_s3_paths[meta_id].add(s3_filepath_basename)

    print("total ids: ", len(id_to_s3_paths))

    with ThreadPoolExecutor(max_workers=32) as executor:
        futures = [
            executor.submit(process_meta_id, meta_id, s3_paths, bucket_name)
            for meta_id, s3_paths in id_to_s3_paths.items()
        ]

        for future in as_completed(futures):
            future.result()  # This will raise an exception if the callable raised