import os
import re
import json
import torch
import boto3

from tqdm import tqdm

DEFAULT_N_CORES = 10
DEFAULT_CHUNKSIZE = 1_000
S3_BUCKET_PATH_RE = r"s3\:\/\/(.+?)\/"


def get_filename(filepath, keep_ext=True):
    if "http" in filepath:
        clean_filepath = filepath.split("?")[0]
    else:
        clean_filepath = filepath
    filename = clean_filepath.split("/")[-1]
    if "." not in filename:
        raise ValueError("filename does not seem to contain a period.")
    m = re.search(r"(.+)\.([^\.]+)$", filename)
    if not m:
        raise ValueError(f"filename could not be parsed for `{filepath}`")
    filename = m.group(1)
    file_ext = m.group(2).lower()
    if len(file_ext) > 10:
        raise ValueError(f"file extension suspiciously long for `{filepath}`")
    if keep_ext:
        filename = filename + "." + file_ext
    return filename


def get_file_ext(filepath):
    filename = get_filename(filepath, keep_ext=True)
    file_ext = filename.split(".")[-1]
    return file_ext


def _parse_s3_filepath(s3_filepath: str):
    bucket_name = re.search(S3_BUCKET_PATH_RE, s3_filepath).group(1)
    rel_s3_filepath = re.sub(S3_BUCKET_PATH_RE, "", s3_filepath)
    return bucket_name, rel_s3_filepath


def _get_client(client_config=None):
    if client_config is not None:
        client = boto3.client(
            "s3",
            endpoint_url=client_config["endpoint_url"],
            aws_access_key_id=client_config["aws_access_key_id"],
            aws_secret_access_key=client_config["aws_secret_access_key"],
            region_name=client_config["region_name"],
        )
    else:
        client = boto3.client("s3")
    return client


def _verify_s3_filepath(filepath):
    if re.search(S3_BUCKET_PATH_RE, filepath) is None:
        raise ValueError("not a valid s3 filepath")
    if get_file_ext(filepath) is None or len(get_file_ext(filepath)) == 0:
        raise ValueError("not a valid file extension")


def download_from_s3(s3_filepath: str, local_filepath: str):
    bucket_name, from_rel_s3_filepath = _parse_s3_filepath(s3_filepath)
    client = _get_client()
    client.download_file(bucket_name, from_rel_s3_filepath, local_filepath)


from suno_amp.utils import measure_energy_in_octave_band

# the goal of the script is to iterate over existing datasets with metdata
# and load the associated tags and audio.
# then we will store the tags and the octave band spectrum and compression factor
# which will be used for training the v1 model.

# we can connect tags to an audio file by iterating over the examples in the metas.jsonl file
# using the s3_filepath we can get the YouTube ID and then connect this to the webm file and construct a path to read audio

if __name__ == "__main__":

    meta_jsonl_filepath = "/app/suno/christian/data/metas.jsonl"
    audio_dir = "/app/suno/christian/data/ytm_audio"
    os.makedirs(audio_dir, exist_ok=True)

    count = 0
    with open(meta_jsonl_filepath, "r") as fp:
        while True:
            json_line = json.loads(fp.readline())
            for key, val in json_line.items():
                print(key, val)
                s3_filepath = json_line["s3_filepath"]
                tags = json_line["tags"]
            if not json_line:
                break

            # download from s3
            local_filepath = os.path.join(audio_dir, os.path.basename(s3_filepath))

            if os.path.isfile(local_filepath):
                print(f"{local_filepath} already downloaded...")
            else:
                print(f"Downloading {local_filepath}...")
                download_from_s3(s3_filepath, local_filepath)