import argparse
import numpy as np
import random
import os
import time
import math
import glob
import soundfile as sf
from scipy import signal
from scipy.io import wavfile
import librosa
import io

import torch
# import torchaudio
import torchaudio.transforms as T
import torch.nn.functional as F

from tqdm import tqdm
import multiprocessing as mp
from multiprocessing import Pool
from typing import Any, Coroutine, Iterable, List, Tuple


def loadWav_samplerate(filepath, sample_rate, transform=None):
    waveform, _sample_rate = sf.read(filepath)

    if _sample_rate != sample_rate:
        t = torch.from_numpy(waveform).type(torch.float32)
        if transform is None:
            transform = T.Resample(_sample_rate, sample_rate)
        waveform = transform(t).numpy()

    return waveform


def check_ogg(root_dir, wav_path, sample_rate = 16000):
    try:
        sub_path = wav_path.rsplit('.', 1)[0].split(root_dir)[-1] + '.ogg' ## wav_path = root_dir + sub_path
        if sub_path[0] == '/':
            sub_path = sub_path[1:]
        ogg_path = os.path.join(root_dir, 'ogg', sub_path)

        errors = []

        if not os.path.isfile(ogg_path):
            error = {'status' : "ogg file not exist", 'detail' : ogg_path}
            errors.append(error)

            return {'status' : 'fail', 'file' : wav_path, 'errors' : errors}
        
        wav_sound = loadWav_samplerate(wav_path, sample_rate)

        file_format = "OGG"
        memory_file = io.BytesIO( )
        memory_file.name = "test.ogg"
        sf.write( memory_file, wav_sound, sample_rate, format = file_format )

        memory_file.seek( 0 )
        temp_data, temp_sr = sf.read( memory_file )


        ogg_sound, sr = sf.read(ogg_path)

        if sr != sample_rate:
            error = {'status' : "sample rate not match", 'detail' : str(sr) + "\t" + str(sample_rate)}
            errors.append(error)
        
        if not np.array_equal(temp_data, ogg_sound):
            error = {'status' : "wav and ogg not equal"}
            errors.append(error)

            return {'status' : 'fail', 'file' : wav_path, 'errors' : errors}
        
    except Exception as e:
        return {'status' : 'exception', 'file' : wav_path, 'errors' : [str(e)]}

    return {'status' : 'success', 'file' : wav_path}
    


def str2bool(v):
    if isinstance(v, bool):
       return v
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')



if __name__ == '__main__':
    parser = argparse.ArgumentParser()
   
    parser.add_argument('--root-dir',   default='/root/nas/data/',   type=str, help='setting to the root directory')
    parser.add_argument('--sub-dir',   default='clovacall',   type=str, help='setting to the root directory')
    parser.add_argument('--dest-dir',       default='/root/preprocess/backup', type=str, help='The destination save path of ogg audio file')
    parser.add_argument('--use-list',   default=False,   type=str2bool, help='use a files.txt')
    
    args = parser.parse_args()

    print("task dir : ", os.path.join(args.root_dir, args.sub_dir))


    wav_files = []
    files_txt = args.sub_dir + "_wavs.txt"
    if args.use_list and os.path.exists(files_txt):
        with open(files_txt, 'r') as f:
            for line in f:
                wav_files.append(line.split('\n')[0])
    else:
        print("glob start !!")
        wav_files = glob.glob(os.path.join(args.root_dir, args.sub_dir, '**/*.wav'), recursive=True)
        print("glob end !!")

        if args.use_list and not os.path.isfile(files_txt):
            with open(files_txt, 'w+') as f:
                for wav in wav_files:
                    f.write(f'{wav}\n')

    # wav_files = wav_files[:100] ## for test

    # for wav_path in wav_files:
    #     r = os.path.join(args.root_dir, args.sub_dir)
    #     sub_path = wav_path.rsplit('.', 1)[0].split(r)[-1] + '.ogg' ## wav_path = root_dir + sub_path
    #     if sub_path[0] == '/':
    #         sub_path = sub_path[1:]

    #     ogg_path = os.path.join(r, 'ogg', sub_path)
    #     # ogg_path = f"{os.path.join(args.root_dir, args.sub_dir)}/ogg/{sub_path}"

    #     print(sub_path + "\t" + ogg_path + "\t" + r)

    print("wav len = ", len(wav_files))
    # exit()
    
    num_cores = mp.cpu_count()
    def work_func(path):
        return check_ogg(os.path.join(args.root_dir, args.sub_dir), path, 16000)

    fail_list = []
    try:
        pool = Pool((num_cores * 2) // 3)

        with tqdm(total=len(wav_files)) as pbar:
            for res in tqdm(pool.imap_unordered(work_func, wav_files)):
                # print(res)
                if res['status'] != 'success':
                    fail_list.append(res)
                    
                pbar.update()

        pool.close()
        pool.join()

        fail_txt = "fail_" + args.sub_dir + ".txt"
        if len(fail_list) >= 1:
            with open(fail_txt, 'w+') as f:
                for fail in fail_list:
                    content = res["file"] + "\t" + str(res['errors'])
                    f.write(content + "\n")

        print("Fail len ", len(fail_list))
        print("Script complete!") 
    except KeyboardInterrupt:
        pool.close()
        pool.join()