#!/usr/bin/env python3
"""Check, flip, merge, and style two-language SRT files. Python standard library only.

check  SOURCE DUAL   compare a bilingual SRT with the one-language SRT it came from
flip   SOURCE DUAL   write the same file with the translation on top
merge  A B           pair two separately timed SRTs by time overlap (--index shows
                     what pairing by cue number would have done)
to-vtt SOURCE DUAL   write WebVTT with a separate CSS class for each language
                     (removes other tags such as <i> and <v>)

check assumes the original lines come first, so run it before flip. A row marked
"original lines not on top" means its per-language numbers are not reliable.
"""
import argparse, html, re, sys

TIME = re.compile(r"(?:(\d+):)?(\d\d):(\d\d)[,.](\d{3})\s*-->\s*(?:(\d+):)?(\d\d):(\d\d)[,.](\d{3})")
TAG = re.compile(r"<[^>]+>|\{\\[^}]*\}")
ENTITY = re.compile(r"&(?:[a-zA-Z]+|#\d+|#x[0-9a-fA-F]+);")
BIDI = re.compile("[‎‏‪-‮⁦-⁩]")


def ms(h, m, s, f):
    return int(h or 0) * 3600000 + int(m) * 60000 + int(s) * 1000 + int(f)


def stamp(t, sep=","):
    return f"{t // 3600000:02}:{t // 60000 % 60:02}:{t // 1000 % 60:02}{sep}{t % 1000:03}"


def read(path):
    try:
        text = open(path, encoding="utf-8-sig").read().replace("\r\n", "\n")
    except UnicodeDecodeError:
        sys.exit(f"{path} is not UTF-8. Convert it to UTF-8 first.")
    cues = []
    for block in re.split(r"\n\s*\n", text.strip()):
        lines = block.split("\n")
        for i, line in enumerate(lines):
            m = TIME.search(line)
            if m:
                g = m.groups()
                cues.append({"start": ms(*g[:4]), "end": ms(*g[4:]), "lines": lines[i + 1:]})
                break
    return cues


def plain(line):
    return BIDI.sub("", html.unescape(TAG.sub("", line))).strip()


def chars(lines):
    return sum(len(plain(l)) for l in lines)


def split_dual(source, dual):
    """Return (original lines, translated lines), assuming the original comes first."""
    n = len(source["lines"])
    return dual["lines"][:n], dual["lines"][n:]


def write(path, cues, sep=",", header=""):
    with open(path, "w", encoding="utf-8") as f:
        f.write(header)
        for i, c in enumerate(cues, 1):
            if not header:
                f.write(f"{i}\n")
            f.write(f"{stamp(c['start'], sep)} --> {stamp(c['end'], sep)}\n" + "\n".join(c["lines"]) + "\n\n")


def check(a):
    source, dual = read(a.source), read(a.dual)
    if not source or not dual:
        sys.exit(f"no cues found: source {len(source)}, bilingual {len(dual)}")
    problems = {}
    def add(kind, n):
        problems.setdefault(kind, []).append(n)
    if len(source) != len(dual):
        print(f"cue count differs: source {len(source)}, bilingual {len(dual)}")
    print("cue  sec  lines  orig_cps  tr_cps  both_cps  notes")
    for n, (s, d) in enumerate(zip(source, dual), 1):
        notes = []
        if (s["start"], s["end"]) != (d["start"], d["end"]):
            notes.append("timing changed"); add("timing changed", n)
        orig, tr = split_dual(s, d)
        if [plain(l) for l in orig] != [plain(l) for l in s["lines"]]:
            notes.append("original lines not on top"); add("original lines do not match the source", n)
        secs = max(d["end"] - d["start"], 1) / 1000
        ocps, tcps = chars(orig) / secs, chars(tr) / secs
        if [plain(l) for l in orig] == [plain(l) for l in tr]:
            notes.append("same text twice"); add("same text twice", n)
        if len(d["lines"]) > a.max_lines:
            notes.append(f"{len(d['lines'])} lines"); add(f"more than {a.max_lines} lines on screen", n)
        if ocps > a.max_cps:
            notes.append("original fast"); add(f"original over {a.max_cps} cps", n)
        if tcps > a.max_target_cps:
            notes.append("translation fast"); add(f"translation over {a.max_target_cps} cps", n)
        if any(len(plain(l)) > a.max_line for l in orig):
            notes.append("long original line"); add(f"original line over {a.max_line} characters", n)
        if any(len(plain(l)) > a.max_target_line for l in tr):
            notes.append("long translated line"); add(f"translated line over {a.max_target_line} characters", n)
        if any(ENTITY.search(l) for l in d["lines"]):
            notes.append("HTML entity"); add("HTML entity in the text", n)
        if n > 1 and d["start"] < dual[n - 2]["end"]:
            notes.append("overlaps previous"); add("overlaps previous cue", n)
        print(f"{n:>3} {secs:>4.1f} {len(d['lines']):>6} {ocps:>9.1f} {tcps:>7.1f} {ocps + tcps:>9.1f}  {', '.join(notes)}")
    print()
    for kind, cues in sorted(problems.items()):
        print(f"{kind}: {len(cues)} ({', '.join(map(str, cues))})")
    if not problems:
        print("no problems found")


def flip(a):
    out = []
    for s, d in zip(read(a.source), read(a.dual)):
        orig, tr = split_dual(s, d)
        out.append({**d, "lines": tr + orig})
    write(a.output, out)
    print(f"wrote {a.output}: {len(out)} cues, translation on top")


def overlap(x, y):
    return max(0, min(x["end"], y["end"]) - max(x["start"], y["start"]))


def merge(a):
    first, second = read(a.first), read(a.second)
    if a.index:
        out = [{**x, "lines": x["lines"] + (second[i]["lines"] if i < len(second) else [])} for i, x in enumerate(first)]
        unpaired = max(0, len(second) - len(first))
    else:
        out = [{**x, "lines": list(x["lines"])} for x in first]
        unpaired = 0
        for y in second:
            best = max(out, key=lambda x: overlap(x, y))
            if overlap(best, y) / max(1, y["end"] - y["start"]) >= a.min_overlap:
                best["lines"] += y["lines"]
            else:
                unpaired += 1
                print(f"not paired: {stamp(y['start'])} {' / '.join(map(plain, y['lines']))}")
    write(a.output, out)
    print(f"wrote {a.output}: {len(out)} cues, {unpaired} second-track cues not paired")


def to_vtt(a):
    style = (
        "WEBVTT\n\nSTYLE\n::cue(.orig) { color: #d8d8d8; font-size: 80%; }\n\n"
        "STYLE\n::cue(.tr) { color: #ffffff; }\n\n"
    )
    out = []
    for s, d in zip(read(a.source), read(a.dual)):
        orig, tr = split_dual(s, d)
        wrap = lambda cls, lines: [f"<c.{cls}>{html.escape(html.unescape(TAG.sub('', l)), quote=False)}</c>" for l in lines]
        out.append({**d, "lines": wrap("orig", orig) + wrap("tr", tr)})
    write(a.output, out, sep=".", header=style)
    print(f"wrote {a.output}: {len(out)} cues")


def main():
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = p.add_subparsers(dest="cmd", required=True)
    c = sub.add_parser("check")
    c.add_argument("source"); c.add_argument("dual")
    c.add_argument("--max-cps", type=float, default=20)
    c.add_argument("--max-target-cps", type=float, default=20)
    c.add_argument("--max-line", type=int, default=42)
    c.add_argument("--max-target-line", type=int, default=42)
    c.add_argument("--max-lines", type=int, default=3)
    c.set_defaults(fn=check)
    for name, fn in (("flip", flip), ("to-vtt", to_vtt)):
        f = sub.add_parser(name)
        f.add_argument("source"); f.add_argument("dual"); f.add_argument("-o", "--output", required=True)
        f.set_defaults(fn=fn)
    m = sub.add_parser("merge")
    m.add_argument("first"); m.add_argument("second"); m.add_argument("-o", "--output", required=True)
    m.add_argument("--min-overlap", type=float, default=0.5)
    m.add_argument("--index", action="store_true", help="pair by cue number instead of time")
    m.set_defaults(fn=merge)
    a = p.parse_args()
    a.fn(a)


if __name__ == "__main__":
    sys.exit(main())