from json import load
from os import listdir, makedirs
from os.path import isdir, join, basename, splitext, dirname, exists


def merge_m3t(source_json: str, trans_json: str, output_m3t: str):
    with open(source_json, 'r', encoding='UTF-8') as f:
        sources = load(f)
    with open(trans_json, 'r', encoding='UTF-8') as f:
        trans = load(f)
    if len(sources) != len(trans):
        print(f"message not matched. source {source_json} length {len(sources)}, trans {trans_json} length {len(trans)}")
    with open(output_m3t, 'w', encoding='UTF-8') as f:
        for src, tran in zip(sources, trans):
            if 'name' in src:
                f.write('○ NAME: ' + src['name'] + '\n\n')
            f.write('○' + src['message'].replace('\n', '\\n') + '\n')
            f.write('●' + tran['message'].replace('\n', '\\n') + '\n\n')


def merge_m3t_auto(source_json: str, trans_json: str, output_m3t: str):
    if isdir(source_json):
        makedirs(output_m3t, exist_ok=True)
        for file in listdir(source_json):
            if not file.lower().endswith(".json"):
                continue
            source_full = join(source_json, file)
            trans_full = join(trans_json, file)
            if not exists(trans_full):
                continue
            output_full = join(output_m3t, splitext(basename(file))[0] + ".m3t")
            merge_m3t(source_full, trans_full, output_full)
    else:
        pdir = dirname(output_m3t)
        if pdir and not isdir(pdir):
            makedirs(pdir, exist_ok=True)
        merge_m3t(source_json, trans_json, output_m3t)


if __name__ == "__main__":
    from argparse import ArgumentParser
    parser = ArgumentParser(description="Merge source JSON and translation JSON into m3t format")
    parser.add_argument("source_json", help="Path to source JSON file or directory")
    parser.add_argument("trans_json", help="Path to translation JSON file or directory")
    parser.add_argument("output_m3t", help="Path to output m3t file or directory")
    args = parser.parse_args()
    merge_m3t_auto(args.source_json, args.trans_json, args.output_m3t)