import argparse
import subprocess
import shutil
from subprocess import CalledProcessError
from pathlib import Path
from multiprocessing import Pool


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('indir', type=Path, help='Directory containing artist directories')
    parser.add_argument('outdir', type=Path, help='Empty directory where transcodes will be placed')
    parser.add_argument('encoder', type=Path, help='Location of encoder')
    return parser.parse_args()


def skip_all_albums(path: Path):
    for directory in path.iterdir():
        if directory.is_dir():
            if not skip_album(directory):
                return False
        else:
            print(f"Warning, skipping non-dir '{directory}' found in '{path}'")
    return True


def skip_album(path: Path):
    for file in path.iterdir():
        if file.name == "DONE":
            return True
    return False


def transcode(transcode_list: list, encoder: Path, workers=16):
    worker_args = [(row[0], row[1], encoder) for row in transcode_list]
    with Pool(workers) as p:
        results = p.starmap_async(transcode_worker, worker_args)
        p.close()
        p.join()
        for result in results.get():
            print(result)


def transcode_worker(in_track, out_track, encoder):
    enc_filename = encoder.parts[-1]
    if enc_filename == "opusenc.exe":
        additional_args = ('--bitrate', '128', '--music')
        subprocess_args = (str(encoder),) + additional_args + (in_track, out_track)
    elif enc_filename == "qaac64.exe":
        additional_args = (in_track, '-o', out_track, '--threading')
        subprocess_args = (str(encoder),) + additional_args
    try:
        subprocess.run(subprocess_args, capture_output=True, text=True, check=True)
        return f"Transcoded '{in_track}' successfully."
    except Exception:
        return f"ERROR: Transcoding of '{in_track}' failed."


def main(input_dir: Path, output_dir: Path, encoder: Path, out_extension: str = 'opus'):
    enc_filename = encoder.parts[-1]
    if enc_filename == "qaac64.exe":
        out_extension = "m4a"
    transcode_list = []
    file_whitelist = ['.jpg', '.jpeg', '.png']
    for artist in input_dir.iterdir():
        if artist.is_dir():
            artist_out = Path(output_dir) / artist.name
            if False: #skip_all_albums(artist):
                print(f"Skipping '{artist.parts[-1]}'")
                continue
            else:
                artist_out.mkdir()
                for album in artist.iterdir():
                    if album.is_dir():
                        album_out = artist_out / album.parts[-1]
                        if album_out.exists():
                            print(f"Skipping '{artist.parts[-1]} / {album.parts[-1]}'")
                        else:
                            album_out.mkdir()
                            done_file = album / 'DONE'
                            open(done_file, 'a').close()
                            for file in album.iterdir():
                                if file.is_file() and file.suffix.lower() == '.flac':
                                    transcode_list.append((str(file), str(album_out / f"{file.stem}.{out_extension}")))
                                elif file.is_file() and file.suffix.lower() in file_whitelist:
                                    shutil.copy(file, album_out / file.name)
                    else:
                        print(f"Warning, skipping non-dir '{album}' found in artist '{artist.parts[-1]}'")
                        continue
        else:
            print(f"Warning, skipping non-dir '{artist}' found in root")
            continue
    transcode(transcode_list, encoder)


if __name__ == '__main__':
    args = get_args()
    main(args.indir, args.outdir, args.encoder)