#!/usr/bin/env python3
"""Wandelt WebVTT-Untertitel (z. B. YouTube-Auto-Untertitel) in eine Rohabschrift um,
die anschließend mit abschrift_korrigieren.py bearbeitet werden kann."""
from pathlib import Path
import argparse
import glob
import html
import os
import re
import sys

TAG = re.compile(r"<[^>]*>")
GLOB_CHARS = re.compile(r"[*?\[]")
CUE_TIME = re.compile(r"(?:(\d+):)??(?:(\d+):)?(\d+(?:[.,]\d+)?)")
CUE_TIMES = re.compile(r"^\s*([\d:.,]+)\s*-->\s*([\d:.,]+)")
USER_TIME = re.compile(r"(\d+)(?::(\d+)(?::(\d+(?:[.,]\d+)?))?)?")
HEADER = re.compile(r"\[(\d{2,}:\d{2}:\d{2}) - (\d{2,}:\d{2}:\d{2})\]")
SENTENCE_END = re.compile(r"[.!?…][\"'»«“”)\]]*(?=\s|$)")
LINE_ENDS_SENTENCE = re.compile(r"[.!?…][\"'»«“”)\]]*$")

MAX_SNAP_SECONDS = 30  # so weit darf ein Zeitraum bis zur nächsten Satzgrenze wachsen
TIME_EXAMPLES = "z. B. 00:05:10, 00:20 (Minute 20) oder 25 (Minute 25)"


def parse_cue_time(text):
    """Seconds from a VTT time: "00:03:17.500" or, without hours, "03:17.500"."""
    match = CUE_TIME.fullmatch(text.strip())
    if not match:
        raise ValueError(f"Ungültige Zeit: {text}")
    hours, minutes, seconds = match.groups()
    return int(hours or 0) * 3600 + int(minutes or 0) * 60 + float(seconds.replace(",", "."))


def parse_user_time(text):
    """Seconds from "00:05:10" (h:m:s), "00:20" (h:m, minute 20) or "25" (minutes)."""
    match = USER_TIME.fullmatch(text.strip())
    if not match:
        raise ValueError(f"Ungültige Zeit: {text}")
    first, minutes, seconds = match.groups()
    if minutes is None:
        return int(first) * 60
    seconds = float(seconds.replace(",", ".")) if seconds is not None else 0
    if int(minutes) >= 60 or seconds >= 60:
        raise ValueError(f"Ungültige Zeit: {text}")
    return int(first) * 3600 + int(minutes) * 60 + seconds


def parse_ranges(text):
    """Sorted (start, end) seconds from "00:05:00-00:06:00 00:10-00:12 00:25".

    "A" or "A-" runs to the end, "0-B" or "-B" from the beginning; None marks an
    open side.
    """
    ranges = []
    for item in text.split():
        start, _, end = item.partition("-")
        if item == "-":
            raise ValueError(f"ungültiger Bereich „{item}“")
        try:
            span = (
                parse_user_time(start) or None if start else None,
                parse_user_time(end) if end else None,
            )
        except ValueError:
            raise ValueError(f"ungültiger Bereich „{item}“") from None
        if span == (None, None):
            raise ValueError(f"Bereich „{item}“ ist das ganze Video: -z weglassen")
        if span[1] is not None and (span[0] or 0) >= span[1]:
            raise ValueError(f"Bereich „{item}“: Start muss vor Ende liegen")
        ranges.append(span)
    if not ranges:
        raise ValueError("keine Bereiche angegeben")

    # Open start first, then by start time
    ranges.sort(key=lambda span: (span[0] is not None, span[0] or 0))
    for (_, end), (start, _) in zip(ranges, ranges[1:]):
        if end is None or start is None or start < end:
            raise ValueError("Bereiche überschneiden sich")
    return ranges


def format_time(seconds, separator=":"):
    seconds = int(seconds)
    return separator.join(
        f"{value:02d}" for value in (seconds // 3600, seconds // 60 % 60, seconds % 60)
    )


def format_header(start, end):
    """Header line in front of a block of text: "[00:04:57 - 00:06:04]"."""
    return f"[{format_time(start)} - {format_time(end)}]"


def range_name(start, end):
    if end is None:
        return f"ab_{format_time(start, '-')}"
    if start is None:
        return f"bis_{format_time(end, '-')}"
    return f"{format_time(start, '-')}_bis_{format_time(end, '-')}"


def time_suffix(ranges):
    """Part of a default file name: "_00-03-17_bis_00-05-25", several joined with "+"."""
    if not ranges:
        return ""
    return "_" + "+".join(range_name(start, end) for start, end in ranges)


RANGE_FORMS = (
    "A-B (von A bis B), A (ab A bis zum Ende) oder 0-B bzw. -B (vom Anfang bis B); "
    f"Zeiten {TIME_EXAMPLES}"
)


def add_time_options(parser):
    parser.add_argument(
        "-z", "--zeitraeume", nargs="+", metavar="BEREICH",
        help=f"nur Untertitel in diesen Bereichen: {RANGE_FORMS}; mehrere durch "
        "Leerzeichen getrennt, z. B. -z 00:05-00:06 00:10-00:12 00:25",
    )
    parser.add_argument(
        "--exakt", action="store_true",
        help="genau an den Zeiten schneiden, statt angeschnittene Sätze vollständig "
        f"aufzunehmen (bis {MAX_SNAP_SECONDS} s über den Bereich hinaus)",
    )


def range_argv(argv):
    """Copy of `argv` with "-00:20" after -z written as "0-00:20".

    argparse takes a word starting with "-" for an option and would reject it;
    no option starts with "-" and a digit, so such a word can only be a range.
    """
    result = []
    after_ranges = False
    for word in argv:
        if re.match(r"-\d", word):
            if after_ranges:
                word = f"0{word}"
        elif word.startswith("-"):
            after_ranges = word in ("-z", "--zeitraeume")
        result.append(word)
    return result


def resolve_time_ranges(parser, args):
    """Check -z and set args.bereiche: sorted list of (start, end), maybe empty."""
    if args.zeitraeume is None:
        args.bereiche = []
        return
    # Parts with and without quotes alike: "-z '0:05-0:06 0:25'" or "-z 0:05-0:06 0:25"
    try:
        args.bereiche = parse_ranges(" ".join(args.zeitraeume))
    except ValueError as error:
        parser.error(f"-z/--zeitraeume: {error} ({RANGE_FORMS})")


def read_cues(vtt):
    """Yield (start, end in seconds, non-empty tag-free text lines) of each cue."""
    cue = None
    start = end = None

    for line in vtt.splitlines():
        if "-->" in line:
            if cue is not None:
                yield start, end, cue
            match = CUE_TIMES.match(line)
            try:
                start, end = map(parse_cue_time, match.groups()) if match else (None, None)
            except ValueError:
                start = end = None
            cue = []
            continue
        # Only a truly empty line ends a cue; YouTube puts lines with a single
        # space inside cues, followed by more text
        if not line:
            if cue is not None:
                yield start, end, cue
            cue = None
            continue
        # Header, NOTE/STYLE blocks and cue identifiers
        if cue is None:
            continue

        text = " ".join(html.unescape(TAG.sub("", line)).split())
        if text:
            cue.append(text)

    if cue is not None:
        yield start, end, cue


def overlap(previous, cue):
    """Length of the longest end of `previous` that starts `cue`."""
    for size in range(min(len(previous), len(cue)), 0, -1):
        if previous[-size:] == cue[:size]:
            return size
    return 0


def in_range(time, start, end):
    if time is None:
        return False
    return (start is None or time >= start) and (end is None or time < end)


def timed_lines(vtt):
    """Return (cue start, cue end, line) for the new caption lines of a VTT document.

    Auto captions show a rolling window: each cue repeats the end of the previous
    cue before adding new text. Only the lines after that overlap are new, so a
    line the speaker really says twice in a row is kept.
    """
    lines = []
    previous = []

    for start, end, cue in read_cues(vtt):
        # Empty cues (YouTube transitions) say nothing about the window
        if not cue:
            continue
        lines.extend((start, end, line) for line in cue[overlap(previous, cue):])
        previous = cue

    return lines


# Positions in the caption text are (line index, character offset); an end
# position is exclusive.

def sentence_start(timed, first, start):
    """Position where the sentence cut by `start` begins, line `first` if none is cut."""
    if start is None or first == 0 or LINE_ENDS_SENTENCE.search(timed[first - 1][2]):
        return first, 0
    for index in range(first - 1, -1, -1):
        time, _, line = timed[index]
        if time is None or time < start - MAX_SNAP_SECONDS:
            return first, 0
        ends = list(SENTENCE_END.finditer(line))
        if ends:
            rest = line[ends[-1].end():]
            if not rest.strip():
                return index + 1, 0
            return index, len(line) - len(rest.lstrip())
    # Reached the first caption: the speech starts there
    return 0, 0


def sentence_end(timed, last, end):
    """Position after the sentence cut by `end`, after line `last` if none is cut."""
    if end is None or LINE_ENDS_SENTENCE.search(timed[last][2]):
        return last + 1, 0
    for index in range(last + 1, len(timed)):
        time, _, line = timed[index]
        if time is None or time >= end + MAX_SNAP_SECONDS:
            return last + 1, 0
        match = SENTENCE_END.search(line)
        if match:
            if match.end() == len(line):
                return index + 1, 0
            return index, match.end()
    # Reached the last caption: the speech ends there
    return len(timed), 0


def excerpt_bounds(timed, start, end, exact=False):
    """(begin, end) positions of a time range, None without captions in it.

    Unless `exact`, the range grows to whole sentences; without a sentence end
    within MAX_SNAP_SECONDS that side is cut exactly.
    """
    inside = [i for i, (time, _, _) in enumerate(timed) if in_range(time, start, end)]
    if not inside:
        return None
    first, last = inside[0], inside[-1]
    if exact:
        return (first, 0), (last + 1, 0)
    return sentence_start(timed, first, start), sentence_end(timed, last, end)


def lines_between(timed, begin, end):
    (first, begin_offset), (last, end_offset) = begin, end
    lines = []
    for index in range(first, min(last + 1, len(timed))):
        line = timed[index][2]
        part = line[begin_offset if index == first else 0:end_offset if index == last else None]
        if part.strip():
            lines.append(part.strip())
    return lines


def range_label(start, end):
    left = "Anfang" if start is None else format_time(start)
    right = "Ende" if end is None else format_time(end)
    return f"{left} – {right}"


def excerpts(timed, ranges, exact=False):
    """Return (start, end in seconds, lines) of the text block of each time range.

    Blocks that overlap after growing to whole sentences are merged; ranges
    without captions are skipped with a note.
    """
    bounds = []
    for start, end in ranges:
        found = excerpt_bounds(timed, start, end, exact)
        if found is None:
            print(f"Hinweis: keine Untertitel im Bereich {range_label(start, end)}")
            continue
        if bounds and found[0] < bounds[-1][1]:
            bounds[-1] = (bounds[-1][0], max(bounds[-1][1], found[1]))
        else:
            bounds.append(found)

    blocks = []
    for begin, end in bounds:
        last = end[0] if end[1] else end[0] - 1
        blocks.append((timed[begin[0]][0], timed[last][1], lines_between(timed, begin, end)))
    return blocks


def join_segments(segments, header_gap="\n"):
    """Text of (header or None, body) blocks, separated by blank lines."""
    return "\n\n".join(
        body if header is None else f"{header}{header_gap}{body}" for header, body in segments
    )


def split_segments(text):
    """Split text at header lines into [(header or None, body)]; see join_segments."""
    segments = []
    header, body = None, []

    def close():
        if header is not None or "".join(body).strip():
            segments.append((header, "\n".join(body).strip()))

    for line in text.splitlines():
        if HEADER.fullmatch(line.strip()):
            close()
            header, body = line.strip(), []
        else:
            body.append(line)
    close()
    return segments


def vtt_to_lines(vtt):
    """Return the caption text lines of a VTT document without rolling duplicates."""
    return [line for _, _, line in timed_lines(vtt)]


def vtt_to_text(vtt, ranges=(), exact=False):
    """Caption text of a VTT document; empty if there is none.

    With time ranges (seconds) one block per range under a header with its
    actual times, completed to whole sentences unless `exact`. The rolling
    window is followed through all cues, so a line repeated from before a
    range does not slip in.
    """
    timed = timed_lines(vtt)
    if not ranges:
        return "\n".join(line for _, _, line in timed)
    return join_segments(
        (format_header(start, end), "\n".join(lines))
        for start, end, lines in excerpts(timed, ranges, exact)
    )


def convert(source, target, ranges=(), exact=False):
    """Convert one file; return an error message or None on success."""
    try:
        vtt = source.read_text(encoding="utf-8-sig")
    except UnicodeDecodeError:
        return f"Keine UTF-8-Datei: {source}"

    text = vtt_to_text(vtt, ranges, exact)
    if not text:
        where = " im gewählten Zeitraum" if ranges else ""
        return f"Keine Untertitel{where} gefunden: {source}"

    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(text + "\n", encoding="utf-8")
    print(f"Erstellt: {target} ({len(text.splitlines())} Zeilen)")
    return None


def expand_inputs(patterns):
    """Resolve glob patterns (for shells that do not, e.g. on Windows).

    Returns the files in order without duplicates, or the first pattern
    without match.
    """
    sources = []
    for pattern in patterns:
        if GLOB_CHARS.search(pattern):
            matches = sorted(glob.glob(pattern))
            if not matches:
                return None, pattern
        else:
            matches = [pattern]
        for match in matches:
            path = Path(match)
            if path not in sources:
                sources.append(path)
    return sources, None


def parse_args():
    parser = argparse.ArgumentParser(
        description="Wandelt VTT-Untertitel in eine Textdatei für abschrift_korrigieren.py um."
    )
    parser.add_argument(
        "-i", "--input", required=True, nargs="+",
        help="VTT-Datei(en) oder Muster, z. B. \"ramstein.de*.vtt\"",
    )
    parser.add_argument(
        "-o", "--output",
        help="Ausgabedatei (nur bei einer Eingabedatei) oder Verzeichnis "
        "(vorhanden oder mit / am Ende); Standard: <Eingabe ohne Endung>.txt, "
        "mit -z zusätzlich die Bereiche im Namen",
    )
    add_time_options(parser)
    args = parser.parse_args(range_argv(sys.argv[1:]))
    resolve_time_ranges(parser, args)
    suffix = time_suffix(args.bereiche)

    sources, unmatched = expand_inputs(args.input)
    if unmatched is not None:
        parser.error(f"Keine Datei passt zu: {unmatched}")
    for source in sources:
        if not source.is_file():
            parser.error(f"Eingabedatei nicht gefunden: {source}")
        if source.suffix.lower() != ".vtt":
            parser.error(f"Keine VTT-Datei: {source}")

    output = args.output
    to_dir = output is not None and (
        output.endswith(("/", os.sep)) or Path(output).is_dir()
    )
    if output is not None and not to_dir and len(sources) > 1:
        parser.error("Bei mehreren Eingabedateien muss -o/--output ein Verzeichnis sein.")

    args.jobs = []
    for source in sources:
        name = f"{source.stem}{suffix}.txt"
        if output is None:
            target = source.with_name(name)
        elif to_dir:
            target = Path(output) / name
        else:
            target = Path(output)
        if target.resolve() == source.resolve():
            parser.error("Ausgabedatei darf nicht die Eingabedatei sein.")
        if any(target.resolve() == other.resolve() for _, other in args.jobs):
            parser.error(f"Mehrere Eingabedateien ergeben dieselbe Ausgabedatei: {target}")
        args.jobs.append((source, target))
    return args


def main():
    args = parse_args()
    failed = 0
    # Keep going on errors so one bad file does not stop a batch
    for source, target in args.jobs:
        error = convert(source, target, args.bereiche, args.exakt)
        if error:
            print(f"FEHLER: {error}", file=sys.stderr)
            failed += 1
    if failed:
        sys.exit(1)


if __name__ == "__main__":
    main()