#!/bin/python3

# external modules

from argparse import ArgumentParser
from mutagen.flac import FLAC
from os import cpu_count, path, walk
from pathlib import Path
from subprocess import check_output, Popen
from sys import exit as s_exit
from time import sleep


# utility functions

def scan_source():
    _artist_dirs, _album_dirs, _full_paths = [], [], []
    for _root, _directories, _files in walk(path.expanduser(source)):
        for _dir in sorted(_directories):
            if not _root.endswith(path.expanduser(source)):
                _artist_dirs.append(_root)
    _artist_dirs = list(set(_artist_dirs))
    for _artist in sorted(_artist_dirs):
        for _root, _directories, _files in walk(path.expanduser(_artist)):
            if not _root.endswith(path.expanduser(_artist)):
                _album_dirs.append(_root)
    _album_dirs = list(set(_album_dirs))
    for _album in sorted(_album_dirs):
        for _root, _directories, _files in walk(path.expanduser(_album)):
            for _filename in _files:
                if _filename.endswith('.flac'):
                    _full_paths.append(f"{_album}/{_filename[:-5]}")
    return _album_dirs, _full_paths


def create_destination():
    for _dir in source_directories[0]:
        Path(path.expanduser(_dir.replace(path.expanduser(source), path.expanduser(destination))))\
            .mkdir(0o700, parents=True, exist_ok=True)


def scan_destination_convert():
    _i, _i_total, _total_files = 0, 0, len(source_directories[1])
    for _file in source_directories[1]:
        _i_total += 1
        _progress = round((_i_total / _total_files) * 100, 2)
        if not Path(f"{_file.replace(path.expanduser(source), path.expanduser(destination))}.mp3").is_file() and not \
                Path(f"{_file.replace(path.expanduser(source), path.expanduser(destination))}.m4a").is_file():
            _file_s = _file.replace(' ', '\\ ').replace("'", "\\'").replace(')', '\\)').replace('(', '\\(')\
                          .replace('&', '\\&').replace('`', '\\`').replace('$', '\\$') + '.flac'
            if args.mp3:
                _file_d = _file.replace(path.expanduser(source), path.expanduser(destination)).replace(' ', '\\ ') \
                              .replace("'", "\\'").replace(')', '\\)').replace('(', '\\(').replace('&', '\\&') \
                              .replace('`', '\\`').replace('$', '\\$') + '.mp3'
            else:
                _file_d = _file.replace(path.expanduser(source), path.expanduser(destination)).replace(' ', '\\ ') \
                              .replace("'", "\\'").replace(')', '\\)').replace('(', '\\(').replace('&', '\\&') \
                              .replace('`', '\\`').replace('$', '\\$') + '.m4a'
            _active_processes = int(check_output('ps -C ffmpeg | wc -l', shell=True)) - 1
            if _active_processes >= cpu_count():
                while _active_processes >= cpu_count():
                    print(f"\nThere are already {cpu_count()} processes running...\n")
                    sleep(1)
                    _active_processes = int(check_output('ps -C ffmpeg | wc -l', shell=True)) - 1
            _i += 1
            print(f"\n{_progress}% | \u001b[38;5;0;48;5;15mConverting {_file_s} to {_file_d}...\u001b[0m\n")
            if args.mp3:
                if args.strip:
                    _cmd = f"screen -DmS flac2pod{_i} ffmpeg -i {_file_s} -map 0:a -ab 320k -map_metadata 0" \
                           f" -id3v2_version 3 {_file_d} </dev/null; rg2sc -f -s --mp3 {_file_d}"
                else:
                    _cmd = f"screen -DmS flac2pod{_i} ffmpeg -i {_file_s} -c:v copy -ab 320k -map_metadata 0" \
                           f" -id3v2_version 3 {_file_d} </dev/null; rg2sc -f --mp3 {_file_d}"
            else:
                _audio = FLAC(f"{_file}.flac")
                if _audio.get('replaygain_track_gain'):
                    _rggain = float(_audio.get('replaygain_track_gain')[0][:-3])
                elif _audio.get('REPLAYGAIN_TRACK_GAIN'):
                    _rggain = float(_audio.get('REPLAYGAIN_TRACK_GAIN')[0][:-3])
                else:
                    _rggain = None
                if _audio.get('replaygain_track_peak'):
                    _rgpeak = float(_audio.get('replaygain_track_peak')[0])
                elif _audio.get('REPLAYGAIN_TRACK_PEAK'):
                    _rgpeak = float(_audio.get('REPLAYGAIN_TRACK_PEAK')[0])
                else:
                    _rgpeak = None
                if args.strip:
                    _cmd = f"screen -DmS flac2pod{_i} ffmpeg -i {_file_s} -c:a aac -ab 256k -map_metadata 0 -metadata" \
                           f" comment='FLAC2PODRG#{_rggain}#{_rgpeak}#' -aac_pns 0 -movflags +faststart -vn" \
                           f" {_file_d} </dev/null; rg2sc -f -s {_file_d}"
                else:
                    _cmd = f"screen -DmS flac2pod{_i} ffmpeg -i {_file_s} -c:a aac -ab 256k -c:v copy -map_metadata" \
                           f" 0 -metadata comment='FLAC2PODRG#{_rggain}#{_rgpeak}#' -aac_pns 0 -movflags +faststart" \
                           f" {_file_d} </dev/null; rg2sc -f {_file_d}"
            Popen(_cmd, shell=True)
    _active_processes = int(check_output('ps -C ffmpeg | wc -l', shell=True)) - 1
    while _active_processes != 0:
        print('\nPlease wait for the remaining conversions to complete...\n')
        sleep(1)
        _active_processes = int(check_output('ps -C ffmpeg | wc -l', shell=True)) - 1
    sleep(1)
    print('\niPod conversion complete!\n')
    s_exit()


# argument parsing
if __name__ == "__main__":
    parser = ArgumentParser(description='Convert your existing FLAC library to be played on an iPod Classic.')
    parser.add_argument('source_dir', nargs='+',
                        help='directory tree containing music files - folder must contain library sorted into artists '
                             'and albums')
    parser.add_argument('destination_dir', nargs='+',
                        help='destination parent directory for output - will be created if it does not already exist')
    parser.add_argument('-s', '--strip', action='store_true',
                        help='strip all TXXX, APIC, and covr tags from the output files')
    parser.add_argument('--mp3', action='store_true',
                        help='convert to MP3, as opposed to the default of AAC .M4A')

    args = parser.parse_args()

    if args.source_dir[0].endswith('/'):
        source = args.source_dir[0][:-1]
    else:
        source = args.source_dir[0]
    if args.destination_dir[0].endswith('/'):
        destination = args.destination_dir[0][:-1]
    else:
        destination = args.destination_dir[0]

    source_directories = scan_source()
    create_destination()
    scan_destination_convert()
