import os
import re
import boto3
import argparse
import sqlalchemy
import pandas as pd

from tqdm import tqdm
from sqlalchemy import create_engine

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)


def grab(output_dir: str, num: int = 10):
    query = f"select * from bots_generatedclip limit {num}"
    df = pd.read_sql(query, engine)
    print(df)

    # construct s3 filepath
    s3_basepath = "s3://suno-data-uploads/studio/uploads/"

    for row_idx, row in tqdm(df.iterrows()):
        s3_filepath = os.path.join(s3_basepath, f"""{(row["id"])}.mp3""")
        print(s3_filepath)
        output_filepath = os.path.join(output_dir, f"""{row["id"]}.mp3""")
        try:
            _verify_s3_filepath(s3_filepath)
            download_from_s3(s3_filepath, output_filepath)
        except Exception as e:
            print(e)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--output_dir", help="Path to directory to store downloaded audio files."
    )
    parser.add_argument(
        "--num", help="Number of audio files to pull from database.", default=10000
    )
    args = parser.parse_args()

    db_password = os.getenv("PGPASSWORD")

    # Ensure the password is actually retrieved, otherwise raise an error
    if db_password is None:
        raise ValueError("PGPASSWORD environment variable not set")

    # Build the connection string using the password from the environment variable
    database_url = f"postgresql+psycopg2://studio_hga1_user:{db_password}@dpg-cgfrde82qv28tc0tavcg-d.replica-cyan.ohio-postgres.render.com/studio_hga1"
    engine = create_engine(database_url)

    grab(args.output_dir, args.num)