import os
import re
import subprocess
import sys
import json


def get_duration(file):
    """Get duration in seconds using ffprobe"""
    result = subprocess.run(
        [
            "ffprobe",
            "-v", "error",
            "-select_streams", "v:0",
            "-show_entries", "format=duration",
            "-of", "json",
            file
        ],
        capture_output=True,
        text=True,
        check=True
    )
    data = json.loads(result.stdout)
    return float(data["format"]["duration"])


def seconds_to_timestamp(seconds):
    """Convert seconds to MKV chapter timestamp"""
    h = int(seconds // 3600)
    m = int((seconds % 3600) // 60)
    s = seconds % 60
    return f"{h:02}:{m:02}:{s:06.3f}"


def create_chapter_file(files, chapter_file):
    current = 0.0

    with open(chapter_file, "w", encoding="utf-8") as f:
        f.write("<?xml version=\"1.0\"?>\n")
        f.write("<!DOCTYPE Chapters SYSTEM \"matroskachapters.dtd\">\n")
        f.write("<Chapters>\n  <EditionEntry>\n")

        for i, file in enumerate(files):
            start = seconds_to_timestamp(current)
            title = f"Part {i+1}"

            f.write("    <ChapterAtom>\n")
            f.write(f"      <ChapterTimeStart>{start}</ChapterTimeStart>\n")
            f.write("      <ChapterDisplay>\n")
            f.write(f"        <ChapterString>{title}</ChapterString>\n")
            f.write("        <ChapterLanguage>eng</ChapterLanguage>\n")
            f.write("      </ChapterDisplay>\n")
            f.write("    </ChapterAtom>\n")

            current += get_duration(file)

        f.write("  </EditionEntry>\n</Chapters>\n")


def append_videos_in_folder(folder_path):
    episode_pattern = re.compile(r"(S\d+E\d+)")
    video_groups = {}

    for root, dirs, files in os.walk(folder_path):
        for file in files:
            match = episode_pattern.search(file)
            if match:
                episode = match.group(1)
                full_path = os.path.join(root, file)
                video_groups.setdefault(episode, []).append(full_path)

    for episode, files in video_groups.items():
        if len(files) > 1:
            files.sort()
            output_file = os.path.join(folder_path, f"{episode}.mkv")
            chapter_file = os.path.join(folder_path, f"{episode}_chapters.xml")

            print(f"Processing {episode}...")

            create_chapter_file(files, chapter_file)

            # mkvmerge append syntax
            cmd = ["mkvmerge", "-o", output_file]

            for i, f in enumerate(files):
                if i == 0:
                    cmd.append(f)
                else:
                    cmd.extend(["+", f])

            cmd.extend(["--chapters", chapter_file])

            try:
                subprocess.run(cmd, check=True)
                print(f"Created {output_file}")
            except subprocess.CalledProcessError as e:
                print(f"Error processing {episode}: {e}")


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} <folder_path>")
        sys.exit(1)

    append_videos_in_folder(sys.argv[1])