# Code Contributor - Ankit Shah - ankit.tronix@gmail.com - Baole Fang - baole.fang@gmail.com
import time
import datetime
import itertools
import os
import sys
import multiprocessing
from multiprocessing import Pool

import yt_dlp
from tqdm import tqdm


# Format audio - 16 bit Signed PCM audio sampled at 44.1kHz
def format_audio(input_audio_file, output_audio_file):
    temp_audio_file = output_audio_file.split('.wav')[0] + '_temp.wav'
    cmdstring = "ffmpeg -loglevel panic -i %s -ac 1 -ar 44100 %s" % (input_audio_file, temp_audio_file)
    os.system(cmdstring)
    cmdstring1 = "sox %s -G -b 16 -r 44100 %s" % (temp_audio_file, output_audio_file)
    os.system(cmdstring1)
    cmdstring2 = "rm -rf %s" % (temp_audio_file)
    os.system(cmdstring2)


# Trim audio based on start time and duration of audio.
def trim_audio(input_audio_file, output_audio_file, start_time, duration):
    # print input_audio_file
    # print output_audio_file
    # print input_audio_file
    # print output_audio_file
    # print start_time
    # print duration
    # ffmpeg -i Y--aaILOrkII.m4a -ss 10 -t 10 -c copy -y Y--aaILOrkII_trimmed.m4a
    # cmdstring = "sox %s %s trim %s %s" %(input_audio_file,output_audio_file,start_time,duration)
    cmdstring = "ffmpeg -loglevel panic -i %s -ss %s -t %s -c copy -y %s" % (
        input_audio_file, start_time, duration, output_audio_file)
    os.system(cmdstring)


def multi_run_wrapper(args):
    return download_audio_method(*args)


# Method to download audio - Downloads the best audio available for audio id, calls the formatting audio function and then segments the audio formatted based on start and end time.
def download_audio_method(line, csv_file):
    query_id = line.split(",")[0];
    # Hacck to get the start time and end time
    start_seconds = line.split(",")[1].split(" ")[1];
    end_seconds = line.split(",")[2].split(" ")[1];
    audio_duration = float(end_seconds) - float(start_seconds)
    # positive_labels = ','.join(line.split(",")[3:]);
    print("Query -> " + query_id)
    # print "start_time -> " + start_seconds
    # print "end_time -> " + end_seconds
    # print "positive_labels -> " + positive_labels
    url = "https://www.youtube.com/watch?v=" + query_id
    # url = "https://youtu.be/" + query_id
    try:
        # video = pafy.new(url)
        # audiostreams = video.audiostreams
        # bestaudio = video.getbestaudio(preftype="m4a")
        # .csv - split - to get the folder information. As path is also passed for the audio - creating the directory from the path where this audio script is present. THus using second split to get the folder name where output files shall be downloaded
        output_folder = sys.argv[1].split('.csv')[0].split("/")[-1] + "_" + csv_file.split('.csv')[
            0] + "_" + "audio_downloaded"
        # print output_folder
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)
        path_to_download = output_folder + "/Y" + query_id + ".m4a"

        URLS = [url]
        ydl_opts = {
            'format': 'm4a/bestaudio/best',
            'outtmpl': path_to_download,
            # ℹ️ See help(yt_dlp.postprocessor) for a list of available Postprocessors and their arguments
            'postprocessors': [{  # Extract audio using ffmpeg
                'key': 'FFmpegExtractAudio',
                'preferredcodec': 'm4a',
            }]
        }
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            error_code = ydl.download(URLS)
        # print path_to_download

        formatted_folder = sys.argv[1].split('.csv')[0].split("/")[-1] + "_" + csv_file.split('.csv')[
            0] + "_" + "audio_formatted_downloaded"
        if not os.path.exists(formatted_folder):
            os.makedirs(formatted_folder)
        path_to_formatted_audio = formatted_folder + "/Y" + query_id + ".wav"
        format_audio(path_to_download, path_to_formatted_audio)
        # Trimming code
        segmented_folder = sys.argv[1].split('.csv')[0].split("/")[-1] + "_" + csv_file.split('.csv')[
            0] + "_" + "audio_formatted_and_segmented_downloads"
        if not os.path.exists(segmented_folder):
            os.makedirs(segmented_folder)
        path_to_segmented_audio = segmented_folder + "/Y" + query_id + '_' + start_seconds + '_' + end_seconds + ".m4a"

        # print path_to_segmented_audio
        trim_audio(path_to_download, path_to_segmented_audio, start_seconds, audio_duration)

        # Remove the original audio and the formatted audio. Comment line to keep both. Delete "output_folder" or "formatted_folder" to keep one.
        # cmdstring2="rm -rf %s %s" %(output_folder,formatted_folder)
        # os.system(cmdstring2)
        # Remove formatted audio. Comment the line to keep the formatted files as well. Deleting as we have original - thus formatted_files could be generated easily
        # cmdstring3="rm -rf %s" %(formatted_folder)
        # os.system(cmdstring3)

        ex1 = ""
    except Exception as ex:
        ex1 = str(ex) + ',' + str(query_id)
        print("Error is ---> " + str(ex))
    return ex1


# Download audio - Reads 3 lines of input csv file at a time and passes them to multi_run wrapper which calls download_audio_method to download the file based on id.
# Multiprocessing module spawns 3 process in parallel which runs download_audio_method. Multiprocessing, thus allows downloading process to happen in 40 percent of the time approximately to downloading sequentially - processing line by line of input csv file.
def download_audio(csv_file, timestamp):
    error_log = 'error' + timestamp + '.log'
    with open(csv_file, "r") as segments_info_file:
        with open(error_log, "a") as fo:
            for line in tqdm(segments_info_file):
                line = (line, csv_file)
                lines_list = []
                lines_list.append(line)
                try:
                    next_line = next(segments_info_file)
                    next_line = (next_line, csv_file)
                    lines_list.append(next_line)
                except:
                    print("end of file")
                try:
                    next_line = next(segments_info_file)
                    next_line = (next_line, csv_file)
                    lines_list.append(next_line)
                except:
                    print("end of file")
                # print lines_list
                P = multiprocessing.Pool(3)

                exception = P.map(multi_run_wrapper, lines_list)
                for item in exception:
                    if item:
                        line = fo.writelines(str(item) + '\n')
                P.close()
                P.join()
        fo.close()


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print('takes arg1 as csv file to downloaded')
    else:

        ts = time.time()
        timestamp = datetime.datetime.fromtimestamp(ts).strftime('%Y_%m_%d_%H_%M_%S')
        download_audio(sys.argv[1], timestamp)
