import os
import boto3
import pandas as pd
from tqdm import tqdm
from mutagen.mp3 import MP3
from urllib.parse import urlparse
from multiprocessing import Pool, cpu_count

def parse_s3_path(s3_path):
    parsed_url = urlparse(s3_path)
    
    if parsed_url.scheme == 's3':
        bucket_name = parsed_url.netloc
        key = parsed_url.path.lstrip('/')
    else:
        raise ValueError("Invalid S3 path format. Must start with 's3://'")

    return bucket_name, key

def download_s3_file(s3_filepath: str, example_id: str, local_dir: str, duration_threshold: int):
    s3 = boto3.client('s3')
    out_filepath = os.path.join(local_dir, f"{example_id}.mp3")
    bucket_name, s3_key = parse_s3_path(s3_filepath)
    try:
        if not os.path.isfile(out_filepath):
            s3.download_file(bucket_name, s3_key, out_filepath)
            print(f"Downloaded {example_id} to {out_filepath}")
            audio = MP3(out_filepath)
            duration_secs = audio.info.length
            
            if duration_secs < duration_threshold:
                print(f"File {example_id} skipped due to short duration ({duration_secs} secs)")
                os.remove(out_filepath) 

    except FileNotFoundError:
        print(f"File {example_id} not found or directory {local_dir} does not exist.")
    except Exception as e:
        os.remove(out_filepath) 
        print(f"Error downloading {example_id}: {str(e)}")

def process_files(files_to_download, local_dir, duration_threshold):
    num_files = len(files_to_download)
    
    with tqdm(total=num_files, desc="Downloading files") as pbar:
        with Pool(processes=min(cpu_count(), 4)) as pool:
            results = []
            for example_id in files_to_download:
                s3_filepath = f's3://suno-data/datasets/harvest/tency_complete/audio/{example_id}/Cover Version.mp3'
                results.append(pool.apply_async(download_s3_file, args=(s3_filepath, example_id, local_dir, duration_threshold)))
            
            downloaded_files = []
            for result in tqdm(results, desc="Processing results", total=num_files, leave=False):
                file_key = result.get()
                if file_key:
                    downloaded_files.append(file_key)
                pbar.update(1)
    
    print("All files downloaded and processed.")

if __name__ == '__main__':
    bucket_name = 'suno-data'
    df_path = '/home/christian_c/christian_c/tency/tency_data.csv'
    local_dir = '/app/suno/christian_c/datasets/tency'
    duration_threshold = 8

    df = pd.read_csv(df_path, index_col='ID')

    files_to_download = list(df.index)

    process_files(files_to_download, local_dir, duration_threshold)