#!/usr/bin/env cached-nix-shell
#!nix-shell -i python3 -p "python3.withPackages(ps: [ ps.tqdm ])" mkvtoolnix-cli
import glob
import multiprocessing
import os.path
from subprocess import run

from tqdm import tqdm


class SubtitleTrack:
    def __init__(self, filename):
        self.filename = filename
        self.language = self._parse_filename()

    def _parse_filename(self):
        basename = self.filename.split(".")[0]
        try:
            language = basename.split("_")[1]
        except IndexError:
            # older format
            language = "eng"

        return language


class Attachment:
    def __init__(self, filename):
        self.filename = filename


class AnimeToshoFile:
    def __init__(self, base):
        self.base = base
        self.outfile = os.path.splitext(self.base)[0] + ".mks"
        self.subtitle_tracks = list(map(SubtitleTrack, self._relative_glob("*.ass")))
        self.attachments = list(map(Attachment, self._relative_glob("attachments/*")))
        self.chapters = "chapters.xml"
        if not os.path.isfile(os.path.join(self.base, self.chapters)):
            self.chapters = None

    def _relative_glob(self, pattern):
        for file in glob.glob(os.path.join(glob.escape(self.base), pattern)):
            yield os.path.basename(file)

    def _full_path(self, filename):
        return os.path.join(self.base, filename)

    def merge(self):
        command = ["mkvmerge", "-q", "-o", self.outfile]
        if self.chapters:
            command.extend(["--chapters", os.path.join(self.base, "chapters.xml")])
        for subtitle_track in self.subtitle_tracks:
            command.extend(
                [
                    "--language",
                    f"0:{subtitle_track.language}",
                    os.path.join(self.base, subtitle_track.filename),
                ]
            )
        for attachment in self.attachments:
            command.extend(
                [
                    "--attachment-mime-type",
                    "font/sfnt",
                    "--attach-file",
                    os.path.join(self.base, "attachments", attachment.filename),
                ]
            )

        run(command, check=True)


files = list(map(AnimeToshoFile, glob.glob("*.mkv")))


with multiprocessing.Pool() as pool:
    list(tqdm(pool.imap_unordered(AnimeToshoFile.merge, files), total=len(files)))