import re
from typing import Optional

TIMESTAMP = r"(?:\d+:)?\d{2}:\d{2}\.\d{3}"
# Timing line: optional hours, verbatim settings tail.
TIMING = re.compile(rf"^({TIMESTAMP})[ \t]+-->[ \t]+({TIMESTAMP})(.*)$")
CUE_TIMES = re.compile(rf"(?m)^([ \t]*)({TIMESTAMP})([ \t]+-->[ \t]+)({TIMESTAMP})")


def timestamp_ms(ts: str) -> int:
    hms, ms = ts.rsplit(".", 1)
    parts = [int(p) for p in hms.split(":")]
    while len(parts) < 3:
        parts.insert(0, 0)
    hours, minutes, seconds = parts
    return ((hours * 60 + minutes) * 60 + seconds) * 1000 + int(ms)


def merge_segmented_webvtt(vtt_raw: str, segment_durations: Optional[list[int]] = None, timescale: int = 1) -> str:
    """
    Merge concatenated Segmented WebVTT documents into a single document.

    Strips the per-segment WEBVTT and X-TIMESTAMP-MAP headers and splices cues
    repeated across a segment boundary (identical payload, <=1ms gap or overlap).
    Timing lines, cue settings and payload are kept verbatim, so inline formatting
    (span tags, entities, ASS-style overrides) survives untouched.

    This function applies no timestamp offset: the segments must already carry absolute
    cue times. This is the same guard as shaka-player and N_m3u8DL-RE, because
    re-offsetting double-shifts such documents. This function accepts
    segment_durations/timescale for call-site compatibility. They only matter if
    someone reintroduces offsetting.

    This function keeps unique STYLE/REGION blocks (players that cannot read them ignore
    them). It drops cue identifiers and NOTE blocks. It skips a cue with a malformed
    timing line (downstream sanitizers catch worse).
    """
    # each kept cue: [start, end, settings, payload, normalized payload for dedupe]
    cues: list[list[str]] = []
    headers: list[str] = []  # unique STYLE/REGION blocks, first-seen order
    prev: Optional[list[str]] = None
    prev_end_ms = 0

    for block in re.split(r"\n[ \t]*\n", vtt_raw.replace("\r\n", "\n").replace("\r", "\n")):
        lines = [
            line
            for line in block.split("\n")
            if not (line.startswith("WEBVTT") or line.lstrip().startswith("X-TIMESTAMP-MAP"))
        ]
        timing_i = next((i for i, line in enumerate(lines) if "-->" in line), None)
        if timing_i is None:  # NOTE or other metadata block
            first = next((line.strip() for line in lines if line.strip()), "")
            if first.startswith(("STYLE", "REGION")):
                header = "\n".join(lines).strip("\n")
                if header not in headers:  # segments repeat the same block
                    headers.append(header)
            continue
        timing = TIMING.match(lines[timing_i].strip())
        if not timing:
            continue
        start, end, settings = timing.groups()
        payload = [line for line in lines[timing_i + 1 :] if line.strip()]
        if not payload:
            continue

        normalized = "\n".join(line.strip() for line in payload)
        if prev is not None and timestamp_ms(start) - prev_end_ms <= 1 and normalized == prev[4]:
            end_ms = timestamp_ms(end)
            if end_ms > prev_end_ms:  # splice: extend only, never shorten the kept cue
                prev[1] = end
                prev_end_ms = end_ms
            continue

        prev = [start, end, settings, "\n".join(payload), normalized]
        prev_end_ms = timestamp_ms(end)
        cues.append(prev)

    blocks = headers + [f"{c[0]} --> {c[1]}{c[2]}\n{c[3]}" for c in cues]
    return "WEBVTT\n\n" + "\n\n".join(blocks) + "\n"


def rebase_webvtt(text: str, offset_ms: int) -> str:
    """
    Subtract offset_ms from every cue time when the cues sit on the media timeline of a DASH
    period whose presentationTimeOffset is offset_ms, as a live replay's do. Cues that start
    below the offset are already presentation times, so the text stays as it is.
    """
    first = CUE_TIMES.search(text)
    if offset_ms <= 0 or not first or timestamp_ms(first[2]) < offset_ms:
        return text

    def shift(ts: str) -> str:
        ms = max(timestamp_ms(ts) - offset_ms, 0)
        return f"{ms // 3_600_000:02}:{ms // 60_000 % 60:02}:{ms // 1000 % 60:02}.{ms % 1000:03}"

    return CUE_TIMES.sub(lambda m: f"{m[1]}{shift(m[2])}{m[3]}{shift(m[4])}", text)