#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Script to maintain a copy of a music collection, in which FLAC files are recoded to OGG Vorbis for
space reasons. If possible, all other files are hard linked, so that no additional space is wasted
for them.
"""

import os
import subprocess
import sys
import time
import shutil
from multiprocessing import Process, Value, Event, Queue
from os.path import relpath, join, exists, getmtime, getsize, basename
import queue
import argparse


DEFAULT_OGG_OPTIONS = ['-Q', '-q3'] # quiet mode, normal quality


def runRsync(options):
    print('running rsync to copy non-flac files...')
    subprocess.call(['rsync',
                     '--exclude=*.flac', '--exclude=unsortiert & neu/*',
                     options.source + '/',
                     options.destination,
                     '-va'] + ([] if options.copy else ['-H']))
    
def startTranssync(options):
    encodingJobs = Queue()
    startTime = time.time()

    toEncode = 0
    flacSize = Value('d', 0)
    oggSize = Value('d', 0)
    encoded = Value('i', 0)
    terminate = Event()
    errorFiles = []
    if not options.no_rsync:
        runRsync(options)
    bla = 0
    
    def inSource(path):
        return join(options.source, relpath(path, options.destination))
    
    def inDestination(path):
        return join(options.destination, relpath(path, options.source))

    # delete files without corresponding sources and empty folders
    for dirpath, _, files in os.walk(options.destination, topdown=False):
        dirInSource = inSource(dirpath)
        if not exists(dirInSource):
            shutil.rmtree(dirpath)
            continue
        for file in files:
            fileInSource = join(dirInSource, file)
            fileInDest = join(dirpath, file)
            if exists(fileInSource):
                continue
            base, extension = file.rsplit('.', 1)
            if extension == 'ogg' and exists(join(dirInSource, base + '.flac')):
                continue
            print('deleting {}'.format(fileInDest))
            os.remove(fileInDest)

    # add encoding and copying jobs
    for dirpath, dirnames, files in os.walk(options.source):
        for file in files:
            sourcefile = join(dirpath, file)
            try:
                base, extension = file.rsplit('.', 1)
            except IndexError:
                continue  # no audio file
            extension = extension.lower()
            if extension == 'flac':
                targetfile = join(inDestination(dirpath), base) + '.ogg'
                if not exists(targetfile) or getmtime(sourcefile) > getmtime(targetfile):
                    encodingJobs.put((sourcefile, targetfile))
                    toEncode += 1

    print(f'starting {options.threads} encoding threads for {toEncode} FLAC files')
    # create and start threads
    threads = []
    for i in range(options.threads):
        t = Process(target=encoderWorker,
                    kwargs=dict(
                        name=i,
                        jobs=encodingJobs,
                        options=options,
                        flacSize=flacSize,
                        oggSize=oggSize,
                        encoded=encoded,
                        terminate=terminate))
        t.start()
        threads.append(t)

    # main loop: fancy output
    try:
        ratio = 0
        while not encodingJobs.empty():
            sys.stdout.write('\r' + ' '*bla)
            if flacSize.value > 0:
                ratio = oggSize.value / flacSize.value
            output = ('Encoded: {}/{} files, shrunk {:.1f}MB FLAC -> '
                      '{:.1f}MB OGG (ratio: {:.3f}). {:.2f}s elapsed.'
                      .format(encoded.value, toEncode,
                              flacSize.value/(1024*1024), oggSize.value/(1024*1024),
                              ratio, time.time() - startTime))
            bla = len(output)
            print('\r' + output, end='')
            sys.stdout.flush()
            time.sleep(.5)
    except KeyboardInterrupt:
        terminate.set()
    print('Waiting for threads to finish...')
    for t in threads:
        t.join()


def encoderWorker(name, jobs, options, flacSize, oggSize, encoded, terminate):
    """Worker function for an encoder thread."""
    while not terminate.is_set():
        try:
            source, target = jobs.get(timeout=0.1)
        except queue.Empty:
            return
        try:
            retcode = subprocess.call(['oggenc'] + DEFAULT_OGG_OPTIONS + ['-o' + target + '.transtmp', source])
        except KeyboardInterrupt:
            terminate.set()
            retcode = 99
        if retcode != 0:
            if terminate.is_set():
                try:
                    os.remove(target + '.transtmp')
                    print('Cleaning unfinished file {}'.format(target + '.transtmp'))
                finally:
                    print('\nThread {} returned'.format(name))
                    return
            else:
                print('ERROR encoding {}'.format(source))
        else:
            print('\nthread {} has encoded {}'.format(name, basename(target)))
            os.rename(target + '.transtmp', target)
            oggSize.value += getsize(target)
            flacSize.value += getsize(source)
            encoded.value += 1


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Maintain a copy of a music collection in which '
        'FLAC files are recoded to ogg vorbis.')
    parser.add_argument('source', help='source directory')
    parser.add_argument('destination', help='target directory')
    parser.add_argument('-t', '--threads', type=int, default=os.cpu_count() - 1)
    parser.add_argument('-c', '--copy', action='store_true', help='always copy non-FLAC files '
        'instead of using hard links')
    parser.add_argument('-n', '--no-rsync', action='store_true')

    options = parser.parse_args()
    startTranssync(options)
