#!/usr/bin/python3

import argparse
import json
import subprocess
import sys


__author__ = 'Marius Gedminas <marius@gedmin.as>'
__version__ = '0.2'


def mkv_info(filename):
    result = subprocess.run(
        ["mkvmerge", "-J", filename], check=True, capture_output=True)
    return json.loads(result.stdout)


def extract_track(filename, track_id, output_filename):
    subprocess.run(
        ["mkvextract", filename, 'tracks', f'{track_id}:{output_filename}'],
        check=True)


def find_subtitle_tracks(filename):
    info = mkv_info(filename)
    for track in info['tracks']:
        if track['type'] == 'subtitles':
            yield track


def list_tracks(tracks, file=sys.stdout):
    for n, track in enumerate(tracks, 1):
        props = track['properties']
        title = props['language']
        if props.get('track_name'):
            title += f" - {props['track_name']}"
        if props.get('default_track'):
            title += " (default)"
        if props.get('forced_track'):
            title += " (forced)"
        print(f"{n}: {title}", file=file)


def choose_track(tracks):
    print("Multiple subtitle tracks found:", file=sys.stderr)
    list_tracks(tracks, file=sys.stderr)
    try:
        # XXX: how do I force input() to print its prompt to sys.stderr?
        which = int(input("Which track do you want to extract? "))
    except (ValueError, EOFError):
        return None
    if 1 <= which <= len(tracks):
        return tracks[which - 1]
    return None


def main():
    parser = argparse.ArgumentParser(
        description=(
            "Extract subtitles from .mkv files.  Convenience wrapper"
            " around mkvmerge and mkvextract."),
    )
    parser.add_argument("filename", metavar='filename.mkv')
    parser.add_argument("-l", "--list-tracks", action='store_true')
    parser.add_argument("-t", "--track", metavar="N", type=int)
    parser.add_argument("-o", "--output-file", metavar="filename.ass",
                        default="/dev/stdout")
    args = parser.parse_args()

    tracks = list(find_subtitle_tracks(args.filename))
    if not tracks:
        sys.exit("No subtitle tracks found")
    if args.list_tracks:
        list_tracks(tracks)
        sys.exit(0)
    if args.track and 1 <= args.track <= len(tracks):
        track = tracks[args.track - 1]
    elif args.track:
        list_tracks(tracks)
        sys.exit(f"There is no track {args.track}")
    elif len(tracks) > 1:
        track = choose_track(tracks)
        if not track:
            sys.exit(1)
    else:
        track = tracks[0]

    extract_track(args.filename, track['id'], args.output_file)


if __name__ == "__main__":
    main()
