import argparse
from pathlib import Path
import subprocess
import json
from argparse import ArgumentParser
from datetime import datetime

def get_ext_by_codec(codec_name):
    if codec_name in ['SubStationAlpha', 'S_TEXT/ASS', 'S_TEXT/SSA']:
        ext = 'ass'
    elif codec_name in ['SubRip/SRT', 'S_TEXT/UTF8', 'S_ARIBSUB']:
        ext = 'srt'
    elif codec_name in ['S_TEXT/WEBVTT', 'D_WEBVTT/SUBTITLES']:
        ext = 'vtt'
    elif codec_name in ['S_HDMV/PGS', 'S_HDMV/TEXTST', 'S_DVBSUB']:
        ext = 'sup'
    else:
        ext = 'ass'
    return ext


def extract_all_subtitles(file: Path):
    print(f'file {file.name}')
    command = ["mkvmerge", "-J", file]
    result = subprocess.run(command, capture_output=True, text=True)

    # Check if the command was successful
    if result.returncode == 0:
        # Parse the JSON response
        mkv_info = json.loads(result.stdout)

        # Now 'mkv_info' is a Python dictionary containing information about the MKV file
        # You can access specific fields like mkv_info['tracks'], etc.
        for track in mkv_info['tracks']:
            track_codec = track['codec']
            track_id = track['id']
            track_type = track['type']
            track_language = track['properties']['language']
            if track_type == 'subtitles':
                extension = get_ext_by_codec(track_codec)
                subprocess.run(['mkvextract', 'tracks', file, f'{track_id}:subs/{file.name[:-4]}.{track_id}.{track_language}.{extension}'])
                print(f'Track {track_id} codec: {track_codec} language: {track_language} type: {track_type}')
    else:
        time_now = datetime.now().strftime("%Y-%m-%d %H:%M")
        # Print an error message if the command failed
        with open('error.txt', 'a') as f:
            f.write(result.stderr)
            f.write(f'[{time_now}] Error: cannot process file {file} {result.stderr}\n')
            f.close()

        print(f"[{time_now}] Error: cannot process file {file} {result.stderr}")


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('filename', metavar='filename', type=str)
    args = parser.parse_args()
    print(args.filename)
    file = Path(args.filename)
    extract_all_subtitles(file)