#!/usr/bin/env python3
from datetime import datetime
import itertools
import math
import re
import os
import subprocess
import tempfile

import click


FILE_TYPE = click.Path(file_okay=True, dir_okay=False, exists=True)
DIRECTORY_TYPE = click.Path(file_okay=False, dir_okay=True, exists=True)
DEVICE_OR_DIRECTORY_TYPE = click.Path(file_okay=True, dir_okay=True, exists=True)


@click.group()
def cli():
    pass


def with_fifo(func):
    def f(*args, **kwargs):
        with tempfile.TemporaryDirectory() as tmpdir:
            fifo_path = os.path.join(tmpdir, 'rip.fifo')
            try:
                os.mkfifo(fifo_path)
            except FileExistsError:
                pass
            kwargs['fifo_path'] = fifo_path
            return func(*args, **kwargs)
    return f


def _get_audio_map(title=None, mpv_output=None, device=None):
    if (None, None) == (title, mpv_output):
        raise RuntimeError('Either `title` or `mpv_output` is required.')

    if mpv_output is None:
        mpv_cmd = ['mpv', '--quiet', '--frames=0',
                   'dvdread://{}'.format(title)]
        if device:
            mpv_cmd.append('--dvd-device={}'.format(device))
        mpv_output = subprocess.check_output(mpv_cmd,
                                             stderr=subprocess.DEVNULL)
        mpv_output = mpv_output.decode('utf-8')

    audio_map = []
    for line in mpv_output.splitlines():
        # [dvd] audio stream: 0 format: mpeg1 (stereo) language: de aid: 0.
        m = re.match(r'\[dvd\] audio stream: (\d+) .* language: (\S+) .*',
                     line)
        if m:
            audio_map.append((int(m.group(1)), m.group(2)))

    if not audio_map:
        print('#' * 15)
        print('WARNING: Could not find any audio streams.')
        print('#' * 15)

    return audio_map


@with_fifo
def _rip_title(title, output_dir, fifo_path, device=None, media_type=None,
               keep_audio=False):
    out_file = os.path.join(output_dir, '{}.mkv'.format(title))

    mpv_cmd = ['mpv', '--quiet', '--stream-dump={}'.format(fifo_path),
               'dvdread://{}'.format(title)]
    if device:
        mpv_cmd.append('--dvd-device={}'.format(device))
    mpv_proc = subprocess.Popen(mpv_cmd, stdout=subprocess.PIPE,
                                stderr=subprocess.STDOUT)

    ffmpeg_cmd = ['ffmpeg',
                  '-probesize', '101M',
                  '-analyzeduration', '150M',
                  '-i', fifo_path,
                  '-map', '0:v:0', 
                  #'-c:a', 'flac',
                  '-c:a', 'copy',
                  '-c:v', 'h264',
                  '-crf', '18',
                  '-pix_fmt', 'yuv420p',
                  #'-x264-params', 'keyint=240:min-keyint=20',
                  '-preset:v', 'slower',
                  '-movflags', '+faststart',
                  #'-profile:v', 'baseline',
                  #'-level', '3.0',
                  ]

    if keep_audio:
        audio_map = _get_audio_map(title, device=device)
        for channel, lang in audio_map:
            ffmpeg_cmd.extend([
                '-map', '0:a:{}'.format(channel),
                '-metadata:s:a:{}'.format(channel), 'language={}'.format(lang)
                ])
    else:
        ffmpeg_cmd.extend(['-map', '0:a:0', '-metadata:s:a:0', 'language=de'])

    if media_type:
        ffmpeg_cmd.extend(['-tune:v', media_type])

    ffmpeg_cmd.append(out_file)
    subprocess.check_call(ffmpeg_cmd)

    mpv_output = mpv_proc.stdout.read().decode('utf-8')
    chapters = _get_title_chapters(mpv_output=mpv_output)
    if chapters:
        _set_chapters(out_file, chapters)


def _get_title_chapters(title=None, mpv_output=None, device=None):
    if (None, None) == (title, mpv_output):
        raise RuntimeError('Either `title` or `mpv_output` is required.')

    if mpv_output is None:
        mpv_cmd = ['mpv', '--quiet', '--frames=0',
                   'dvdread://{}'.format(title)]
        if device:
            mpv_cmd.append('--dvd-device={}'.format(device))
        mpv_output = subprocess.check_output(mpv_cmd,
                                             stderr=subprocess.DEVNULL)
        mpv_output = mpv_output.decode('utf-8')

    for line in mpv_output.splitlines():
        m = re.match(r'\[dvd\] CHAPTERS: (\S*)', line)
        if m:
            chapters = m.group(1)
            break
    else:
        raise RuntimeError('Could not find CHAPTERS in mpv output.')

    chapters = chapters.split(',')
    chapters = [c for c in chapters if c]   # filter empty/last

    return chapters


def _convert_chapters_to_mkvmerge_format(chapters):
    lines =  []
    number_of_zeros = 0
    if len(chapters) > 1:
        number_of_zeros = int(math.log10(len(chapters) - 1))
    chapter_template = 'CHAPTER{{:0{}d}}'.format(number_of_zeros)
    for i, timestamp in enumerate(chapters):
        chapter_str = chapter_template.format(i)
        lines.append('{}={}'.format(chapter_str, timestamp))
        lines.append('{}NAME='.format(chapter_str))
    return os.linesep.join(lines)


def _set_chapters(mkv_file, chapters):
    if not os.path.exists(mkv_file):
        raise RuntimeError('MKV file does not exist.')

    with tempfile.NamedTemporaryFile() as f:
        mkvmerge_chapters = _convert_chapters_to_mkvmerge_format(chapters)
        f.write(mkvmerge_chapters.encode('utf-8'))
        f.flush()

        mkvpropedit_cmd = ['mkvpropedit', mkv_file, '--chapters', f.name]
        subprocess.check_call(mkvpropedit_cmd)


def _get_dvd_titles(device=None):
    mpv_cmd = ['mpv', '--vo', 'null', '--quiet', '--frames=0', 'dvdread://']
    if device:
        mpv_cmd.append('--dvd-device={}'.format(device))
    out = subprocess.check_output(mpv_cmd, stderr=subprocess.STDOUT)
    out = out.decode('utf-8')

    for line in out.splitlines():
        m = re.match(r'^\[dvd\] There are (\d+) titles on this DVD.', line)
        # next line doesn't work on DVDs with one titleset having multiple
        # titles
        # m = re.match(r'^libdvdread: Found (\d+) VTS', line)
        if m:
            titles = int(m.group(1))
            break
    else:
        raise RuntimeError('Could not find the number of VTS in mpv output.')

    return titles


@cli.command()
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-d', '--device', type=DEVICE_OR_DIRECTORY_TYPE)
@click.option('--media-type', type=click.Choice(['film', 'animation']))
@click.option('--keep-audio', is_flag=True)
def all(output_dir, device=None, media_type=None, keep_audio=False):
    titles = _get_dvd_titles(device=device)
    for title in range(0, titles):
        _rip_title(title, output_dir, device=device, media_type=media_type,
                   keep_audio=keep_audio)


@cli.command()
@click.option('-t', '--title', type=int, required=True)
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
@click.option('-d', '--device', type=DEVICE_OR_DIRECTORY_TYPE)
@click.option('--media-type', type=click.Choice(['film', 'animation']))
@click.option('--keep-audio', is_flag=True)
def single(title, output_dir, device=None, media_type=None, keep_audio=False):
    _rip_title(title, output_dir, device=device, media_type=media_type,
               keep_audio=keep_audio)


@cli.command('get-chapters')
@click.option('-t', '--title', type=int, required=True)
@click.option('-d', '--device', type=DEVICE_OR_DIRECTORY_TYPE)
def get_chapters(title, device=None):
    print(','.join(_get_title_chapters(title=title, device=device)))


@cli.command('set-chapters')
@click.option('-o', '--output-file', type=FILE_TYPE, required=True)
@click.option('-t', '--title', type=int)
@click.option('-c', '--chapters')
@click.option('-d', '--device', type=DEVICE_OR_DIRECTORY_TYPE)
def set_chapters(output_file, title=None, chapters=None, device=None):
    if (None, None) == (title, chapters):
        raise RuntimeError('One of `title` and `chapters` is required.')

    if chapters is None:
        chapters = _get_title_chapters(title=title, device=device)
    else:
        chapters = chapters.split(',')

    _set_chapters(output_file, chapters)


@cli.command('find-right-title')
@click.option('--start-title', type=int, default=0)
@click.option('--end-title', type=int)
@click.option('-d', '--device', type=DEVICE_OR_DIRECTORY_TYPE)
def find_right_title(start_title, end_title, device):
    if end_title and start_title > end_title:
        raise click.ClickException('--start-title must not be bigger than --end-title')

    titles = _get_dvd_titles(device=device)
    if start_title > titles:
        raise RuntimeError('DVD has less titles than start title.')

    if end_title:
        end_title = min(titles - 1, end_title)
    else:
        end_title = titles - 1

    title_chapter_lengths = {}
    for title in range(start_title, end_title + 1):
        chapters = _get_title_chapters(title, device=device)

        timestamps = []
        for timestamp in chapters:
            timestamp, microseconds = timestamp.split('.')
            dt = datetime.strptime(timestamp, '%H:%M:%S')
            dt = dt.replace(microsecond=int(microseconds) * 1000)
            timestamps.append(dt)

        lengths = []
        if len(timestamps) > 1:
            for i, timestamp in enumerate(timestamps[1:]):
                lengths.append(timestamp - timestamps[i])

        if not lengths:
            continue

        title_chapter_lengths[title] = lengths

    # remove titles with duplicate chapters by looking at the length
    title_chapter_lengths = {title: lengths for title, lengths in title_chapter_lengths.items()
                             if len(lengths) == len(set(lengths))}

    if title_chapter_lengths:
        print('The titles {} contain no duplicate chapter lengths.'.format(
              ', '.join(str(i) for i in title_chapter_lengths.keys())))
    else:
        print('No title without duplicates found.')


@cli.command('split-file')
@click.option('-c', '--chapters-file', type=FILE_TYPE, required=True,
    help='File with timestamps separated by comma')
@click.option('-i', '--input-file', type=FILE_TYPE, required=True)
@click.option('-o', '--output-dir', type=DIRECTORY_TYPE, required=True)
def split_file(output_dir, input_file, chapters_file):
    extension = input_file.split('.')[-1]

    with open(chapters_file) as f:
        chapters = f.read().strip().split(',')
    
    for i, (start, end) in enumerate(itertools.zip_longest(chapters, chapters[1:])):
        ffmpeg_cmd = ['ffmpeg',
                      '-ss', start]
        if end is not None:
            ffmpeg_cmd += ['-to', end]

        ffmpeg_cmd += [
            '-i', input_file,
            '-c:a', 'flac',
            #'-c:v', 'copy',
            os.path.join(output_dir, '{}.{}'.format(i + 1, extension))
            ]
        #print(ffmpeg_cmd)
        subprocess.check_call(ffmpeg_cmd)


if __name__ == '__main__':
    cli()