"""
* Version: 0.1.0.0.0

* Purposes:
    * Process the SRT files.
    * Write each intermediate result to file for caching.
"""

import re
import pingcap.pe.flash.ai_editor as editor
import pingcap.pe.flash.util as util
import pingcap.pe.flash.util_en2ja as util_en2ja
import pingcap.pe.flash.util_en2zh as util_en2zh


def process_srt_to_paragraph(lines: str) -> dict:
    # The parsing, creates two types of keys: <sequence> and <sequence>-time
    srt_dict = {}
    segment_idx = None
    for line in lines.split("\n"):
        if re.match("^[0-9]*$", line):
            segment_idx = line
            srt_dict.update({segment_idx: ""})
        elif re.match("^[0-9][0-9]:[0-9][0-9]:[0-9][0-9],[0-9][0-9][0-9]", line):
            srt_dict.update({segment_idx + "-TIME": line})
        elif re.match("^$", line):
            None
        else:
            srt_dict[segment_idx] = (
                srt_dict[segment_idx] + " " + line.replace("\xa0", "")
            ).strip()

    # Clean-up
    del srt_dict[""]

    # The leading three, creates a new type of key: <sequence>-l3
    l3_dict = {}
    for k in srt_dict.keys():
        if not k.endswith("-TIME"):
            l3 = " ".join(
                [
                    w.strip(",").strip(".").strip(";").lower()
                    for w in srt_dict[k].split(" ")[0:3]
                ]
            )
            l3_dict.update({k + "-L3": l3})
    srt_dict.update(l3_dict)

    # The entire paragraph, create a new key: paragraph
    p = []
    for k in srt_dict.keys():
        if not k.endswith("-TIME") and not k.endswith("-L3"):
            p += [srt_dict[k]]
    paragraph = " ".join(p)

    # The entire time paragraph, create a new key: paragraph
    tp = []
    for k in srt_dict.keys():
        if not k.endswith("-L3"):
            tp += [srt_dict[k]]
    tparagraph = " ".join(tp)

    srt_dict.update({"time-paragraph": tparagraph})
    srt_dict.update({"paragraph": paragraph})
    return srt_dict


def split_en_paragraph_to_sentences(cache_dict: dict) -> dict:
    sp = [".", ";", "?"]

    sentences = []
    for w in cache_dict["paragraph"].split(" "):
        sentences += [w]
        if w[-1] in sp:
            sentences += ["\n"]
    cache_dict["sentences"] = " ".join(sentences).strip()
    sentences = cache_dict["sentences"]
    i = 1
    for s in sentences.split("\n"):
        if not s == "":
            cache_dict["S-" + str(i)] = s.strip()
            i += 1

    tsentences = []
    for w in cache_dict["time-paragraph"].split(" "):
        tsentences += [w]
        if w[-1] in sp:
            tsentences += ["\n"]
    cache_dict["time-sentences"] = " ".join(tsentences).strip()
    tsentences = cache_dict["time-sentences"]
    i = 1
    real_tsentences = []
    for s in tsentences.split("\n"):
        if not s == "":
            real_tsentences += [s]
    ls = len(real_tsentences)
    for s in real_tsentences:
        matches = re.findall(
            "[0-9][0-9]:[0-9][0-9]:[0-9][0-9],[0-9][0-9][0-9] --> [0-9][0-9]:[0-9][0-9]:[0-9][0-9],[0-9][0-9][0-9]",
            s,
        )
        ys = re.sub(
            "[0-9][0-9]:[0-9][0-9]:[0-9][0-9],[0-9][0-9][0-9] --> [0-9][0-9]:[0-9][0-9]:[0-9][0-9],[0-9][0-9][0-9]",
            "|+|",
            s,
        ).strip()
        if i == 1:
            cache_dict["TS-" + str(i)] = "EN: |+| " + ys
            cache_dict["TSS-" + str(i)] = len(matches) + 1
        elif i == ls:
            None
        else:
            cache_dict["TS-" + str(i)] = "EN: " + ys
            cache_dict["TSS-" + str(i)] = len(matches)
        i += 1

    return cache_dict


def _align_time_marker_count(source_line: str, target_line: str) -> str:
    source_mc = source_line.count("|+|")  # marker count in source line
    target_mc = target_line.count("|+|")  # marker count in target line
    result = target_line
    if source_mc > target_mc:
        if source_line.startswith("|+|") and not target_line.startswith("|+|"):
            result = "|+|" + target_line  # append a leading marker
        else:
            result = target_line + " |+|"  # append a trailing marker
    elif source_mc < target_mc:
        if target_line.endswith("|+|"):
            result = re.sub("\|\+\|$", "", target_line)  # remove the trailing marker
        else:
            result = target_line.replace("|+|", "", 1)  # remove the first marker
    else:
        return result
    return _align_time_marker_count(source_line, result)


def translate_sentence_en2ja(cache_dict: dict) -> dict:
    en_1gram_list = util.load_1gram_terms_en("pe")
    ja_sentences = []
    for k in cache_dict.keys():
        if k.startswith("TS-"):
            en_srt_line = cache_dict[k].replace("EN: ", "")
            ja_injected_srt_line = util_en2ja.translate_en2ja_md_line(
                en_srt_line,
                en_1gram_list,
                "standard",
                True,
                None,
                True,
            ).strip()
            ja_injected_srt_line = (
                ja_injected_srt_line.replace(
                    "`sign |+|` を見えないものとして扱いますが、翻訳結果には入れます。", ""
                )
                .replace("この文章を日本語に翻訳すると、次のようになります:", "")
                .replace("|+| もし元の英文に存在した場合は、翻訳結果に |+| を挿入してください。", "")
                .replace("が元の英語の文章に含まれている場合は、結果にもそれらを戻す必要があります。", "")
                .replace("この場合、英語の文でが存在するので、翻訳結果でもそれを入れる必要があります。", "")
                .replace("「|+|」は英語の文章にある場合、翻訳結果に再度挿入する必要があります。", "")
            )
            ja_injected_srt_line = re.sub("^私はあなたに", "", ja_injected_srt_line)
            result = _align_time_marker_count(en_srt_line, ja_injected_srt_line)
            ja_sentences += [result]
    i = 1
    for j in ja_sentences:
        if j != None:
            cache_dict["JS-" + str(i)] = j
            cache_dict["EJS-" + str(i)] = cache_dict["TS-" + str(i)] + "\nJA: " + j
            i += 1
    return cache_dict


def translate_sentence_en2zh(cache_dict: dict) -> dict:
    en_1gram_list = util.load_1gram_terms_en("pe")
    ja_sentences = []
    for k in cache_dict.keys():
        if k.startswith("S-"):
            ja_sentences += [
                util_en2zh.translate_en2zh_md_line(
                    cache_dict[k], en_1gram_list, "standard", True, None
                ).strip()
            ]
    i = 1
    for j in ja_sentences:
        cache_dict["ZS-" + str(i)] = j
        cache_dict["EZS-" + str(i)] = cache_dict["TS-" + str(i)] + "\nZH: " + j
        i += 1
    return cache_dict