from __future__ import annotations

import re
from dataclasses import dataclass

from caption_profile import (
    CaptionProfile,
    characters_per_second,
    display_unit_count,
    restore_cue_wide_markup,
    split_cue_wide_markup,
    wrap_text,
)
from timed_text import (
    Cue,
    TimedTextDocument,
    TimedWord,
    TimingFidelity,
    WordTimingSource,
)
from timebase import RationalTimebase


SENTENCE_END_RE = re.compile(r"""[.!?]["')\]}]*$""")
CJK_RE = re.compile(r"[\u3040-\u30ff\u3400-\u9fff\uf900-\ufaff]")
WEAK_CAPTION_END_WORDS = {
    # Articles, conjunctions, and short prepositions are normally understood with
    # the words that follow. Keeping this compact multilingual set here prevents
    # obviously awkward caption endings without pretending to parse grammar.
    "a", "an", "the", "and", "or", "but", "for", "nor", "so", "yet",
    "of", "to", "in", "on", "at", "by", "from", "with", "without", "into",
    "onto", "as",
    "der", "die", "das", "den", "dem", "des", "ein", "eine", "einer",
    "eines", "einem", "einen", "und", "oder", "aber", "für", "von", "zu",
    "im", "am", "an", "auf", "aus", "mit", "ohne", "als",
    "le", "la", "les", "un", "une", "des", "de", "du", "et", "ou", "mais",
    "à", "au", "aux", "en", "pour", "par", "avec", "sans",
    "el", "los", "las", "unos", "unas", "del", "y", "o", "pero", "por",
    "para", "con", "sin",
    "il", "lo", "i", "gli", "uno", "una", "di", "della", "e", "ma", "per",
    "senza",
}
WEAK_CAPTION_START_WORDS = {
    # A caption that begins with one of these words usually continues a phrase
    # from the preceding caption. Articles are intentionally excluded because
    # they read naturally at the start together with their noun.
    "of", "to", "in", "on", "at", "by", "from", "with", "without", "into",
    "onto", "von", "zu", "im", "am", "an", "auf", "aus", "mit", "ohne",
    "de", "du", "à", "au", "aux", "en", "pour", "par", "avec", "sans",
    "del", "por", "para", "con", "sin", "di", "della", "per", "senza",
}
LATIN_RE = re.compile(r"[A-Za-z]")


@dataclass(frozen=True)
class Finding:
    severity: str
    code: str
    message: str
    cue_ids: tuple[str, ...] = ()


@dataclass(frozen=True)
class SentenceGroup:
    id: str
    cue_ids: tuple[str, ...]
    start_ms: int
    end_ms: int
    plain_text: str
    wrappers: tuple[str, ...]
    warnings: tuple[str, ...]


@dataclass(frozen=True)
class ReflowToken:
    text: str
    start_ms: int
    end_ms: int
    source_cue_id: str
    timed_word: TimedWord | None = None


def find_overlaps(document: TimedTextDocument) -> list[tuple[str, str, int, int]]:
    overlaps: list[tuple[str, str, int, int]] = []
    active: list[Cue] = []
    for cue in sorted(document.cues, key=lambda item: (item.start_ms, item.end_ms, item.srt_index)):
        active = [candidate for candidate in active if candidate.end_ms > cue.start_ms]
        for candidate in active:
            start = max(candidate.start_ms, cue.start_ms)
            end = min(candidate.end_ms, cue.end_ms)
            if end > start:
                overlaps.append((candidate.id, cue.id, start, end))
        active.append(cue)
    return overlaps


def form_sentence_groups(
    document: TimedTextDocument,
    profile: CaptionProfile,
    *,
    max_gap_ms: int = 500,
) -> list[SentenceGroup]:
    groups: list[list[Cue]] = []
    current: list[Cue] = []
    for cue in document.cues:
        if not current:
            current = [cue]
            continue
        previous = current[-1]
        current_plain = " ".join(split_cue_wide_markup(item.text)[0].replace("\n", " ") for item in current)
        cue_plain = split_cue_wide_markup(cue.text)[0].replace("\n", " ")
        boundary = (
            cue.start_ms < previous.end_ms
            or cue.start_ms - previous.end_ms > max_gap_ms
            or (previous.speaker and cue.speaker and previous.speaker != cue.speaker)
            or (previous.source_track and cue.source_track and previous.source_track != cue.source_track)
            or SENTENCE_END_RE.search(current_plain.strip()) is not None
            or cue.end_ms - current[0].start_ms > profile.group_max_duration_ms
            or display_unit_count(current_plain + " " + cue_plain) > profile.group_max_units
        )
        if boundary:
            groups.append(current)
            current = [cue]
        else:
            current.append(cue)
    if current:
        groups.append(current)

    result: list[SentenceGroup] = []
    for index, cues in enumerate(groups, start=1):
        parts: list[str] = []
        wrapper_sets: list[tuple[str, ...]] = []
        warnings: list[str] = []
        for cue in cues:
            plain, wrappers, markup_warnings = split_cue_wide_markup(cue.text)
            parts.append(plain.replace("\n", " ").strip())
            wrapper_sets.append(wrappers)
            warnings.extend(markup_warnings)
        uniform = bool(wrapper_sets) and all(value == wrapper_sets[0] for value in wrapper_sets)
        wrappers = wrapper_sets[0] if uniform else ()
        if not uniform and any(wrapper_sets):
            warnings.append("Mixed cue-wide markup was stripped for this sentence group.")
        result.append(
            SentenceGroup(
                id=f"group_{index:06d}",
                cue_ids=tuple(cue.id for cue in cues),
                start_ms=cues[0].start_ms,
                end_ms=max(cue.end_ms for cue in cues),
                plain_text=" ".join(part for part in parts if part),
                wrappers=wrappers,
                warnings=tuple(dict.fromkeys(warnings)),
            )
        )
    return result


def validate_document(
    document: TimedTextDocument,
    profile: CaptionProfile,
    *,
    explicit_gap_limit: bool = False,
    gap_timebase: RationalTimebase | None = None,
    required_gap_frames: int | None = None,
) -> list[Finding]:
    findings: list[Finding] = []
    try:
        document.validate()
    except ValueError as exc:
        findings.append(Finding("ERROR", "INVALID_DOCUMENT", str(exc)))
        return findings
    for cue in document.cues:
        lines = cue.text.splitlines() or [cue.text]
        _, _, markup_warnings = split_cue_wide_markup(cue.text)
        findings.extend(
            Finding("WARNING", "MARKUP_WARNING", warning, (cue.id,))
            for warning in markup_warnings
        )
        duration = cue.end_ms - cue.start_ms
        if len(lines) > profile.max_lines:
            findings.append(
                Finding("WARNING", "TOO_MANY_LINES", f"Cue has {len(lines)} lines; limit is {profile.max_lines}.", (cue.id,))
            )
        for line_number, line in enumerate(lines, start=1):
            units = display_unit_count(line)
            if units > profile.max_units_per_line:
                findings.append(
                    Finding(
                        "WARNING",
                        "LINE_TOO_LONG",
                        f"Cue line {line_number} has {units} display units; limit is {profile.max_units_per_line}.",
                        (cue.id,),
                    )
                )
        cps = characters_per_second(cue.text, duration)
        if cps > profile.max_cps:
            findings.append(
                Finding("WARNING", "CPS_TOO_HIGH", f"Cue CPS is {cps:.2f}; limit is {profile.max_cps:g}.", (cue.id,))
            )
        if duration < profile.min_duration_ms:
            findings.append(Finding("WARNING", "CUE_TOO_SHORT", "Cue duration is below the configured minimum.", (cue.id,)))
        if duration > profile.max_duration_ms:
            findings.append(Finding("WARNING", "CUE_TOO_LONG", "Cue duration exceeds the configured maximum.", (cue.id,)))
    for left_id, right_id, start_ms, end_ms in find_overlaps(document):
        findings.append(
            Finding(
                "ERROR",
                "CUE_OVERLAP",
                f"Cues overlap from {start_ms} ms to {end_ms} ms.",
                (left_id, right_id),
            )
        )
    for previous, current in zip(document.cues, document.cues[1:]):
        gap = current.start_ms - previous.end_ms
        if 0 <= gap < profile.preferred_gap_ms:
            gap_code = (
                "GAP_BELOW_PROFILE_LIMIT"
                if explicit_gap_limit
                else "GAP_BELOW_PREFERENCE"
            )
            gap_severity = "ERROR" if explicit_gap_limit else "WARNING"
            threshold_label = (
                "required minimum" if explicit_gap_limit else "preferred minimum"
            )
            if gap_timebase is not None:
                measured_frames = float(
                    gap_timebase.milliseconds_to_frames(gap)
                )
                if required_gap_frames is not None:
                    threshold = (
                        f"{required_gap_frames} frame(s) (at least "
                        f"{profile.preferred_gap_ms} ms)"
                    )
                else:
                    threshold_frames = float(
                        gap_timebase.milliseconds_to_frames(
                            profile.preferred_gap_ms
                        )
                    )
                    threshold = (
                        f"{profile.preferred_gap_ms} ms "
                        f"({threshold_frames:.2f} frames)"
                    )
                gap_message = (
                    f"Adjacent cue gap is {gap} ms "
                    f"({measured_frames:.2f} frames); {threshold_label} is "
                    f"{threshold}."
                )
            else:
                gap_message = (
                    f"Adjacent cue gap is {gap} ms; {threshold_label} is "
                    f"{profile.preferred_gap_ms} ms."
                )
            findings.append(
                Finding(
                    gap_severity,
                    gap_code,
                    gap_message,
                    (previous.id, current.id),
                )
            )
        if (
            previous.text.replace("\n", " ").strip().casefold()
            == current.text.replace("\n", " ").strip().casefold()
        ):
            findings.append(
                Finding(
                    "WARNING",
                    "SUSPICIOUS_DUPLICATE_TEXT",
                    "Adjacent cues contain identical normalized text.",
                    (previous.id, current.id),
                )
            )
    return findings


def line_only_reflow(
    document: TimedTextDocument, profile: CaptionProfile
) -> tuple[TimedTextDocument, list[Finding]]:
    cues: list[Cue] = []
    findings: list[Finding] = []
    for cue in document.cues:
        wrapped, warnings = wrap_text(cue.text, profile.max_units_per_line, profile.max_lines)
        cues.append(
            Cue(
                id=cue.id,
                srt_index=cue.srt_index,
                start_ms=cue.start_ms,
                end_ms=cue.end_ms,
                text=wrapped,
                speaker=cue.speaker,
                source_track=cue.source_track,
                source_item=cue.source_item,
                words=list(cue.words),
                metadata=dict(cue.metadata),
            )
        )
        findings.extend(Finding("WARNING", "REFLOW_WARNING", warning, (cue.id,)) for warning in warnings)
    result = TimedTextDocument(
        version=document.version,
        language=document.language,
        source_kind=document.source_kind,
        source_name=document.source_name,
        timing_fidelity=document.timing_fidelity,
        cues=cues,
        metadata=dict(document.metadata),
    )
    return result, findings


def allocate_estimated_boundaries(
    texts: list[str], start_ms: int, end_ms: int
) -> list[tuple[int, int]]:
    if not texts or any(not text.strip() for text in texts):
        raise ValueError("Estimated timing requires non-empty text segments.")
    if end_ms - start_ms < len(texts):
        raise ValueError("Group window is too short for positive cue durations.")
    weights = [max(1, display_unit_count(text, exclude_whitespace=True)) for text in texts]
    total_weight = sum(weights)
    duration = end_ms - start_ms
    boundaries = [start_ms]
    cumulative = 0
    for weight in weights[:-1]:
        cumulative += weight
        boundary = start_ms + int(round(duration * cumulative / total_weight))
        boundary = max(boundaries[-1] + 1, min(boundary, end_ms - (len(texts) - len(boundaries))))
        boundaries.append(boundary)
    boundaries.append(end_ms)
    return list(zip(boundaries, boundaries[1:]))


def mixed_script_interpolation_warning(text: str) -> str | None:
    if CJK_RE.search(text) and LATIN_RE.search(text):
        return (
            "Estimated timing uses display-unit text weights; mixed CJK and Latin text can differ "
            "substantially from speaking duration."
        )
    return None


def _join_tokens(tokens: list[str]) -> str:
    value = ""
    for token in tokens:
        if (
            not value
            or re.fullmatch(r"[,.;:!?%)\]}]+", token)
            or (CJK_RE.search(value[-1:]) and CJK_RE.search(token[:1]))
        ):
            value += token
        else:
            value += " " + token
    return value.strip()


def _estimated_tokens(cue: Cue, plain: str) -> list[ReflowToken]:
    normalized = plain.replace("\n", " ").strip()
    values = normalized.split()
    if len(values) == 1 and CJK_RE.search(normalized):
        values = [character for character in normalized if not character.isspace()]
    if not values:
        return []
    positions = allocate_estimated_boundaries(values, cue.start_ms, cue.end_ms)
    return [
        ReflowToken(text, start_ms, end_ms, cue.id)
        for text, (start_ms, end_ms) in zip(values, positions)
    ]


def _reflow_regions(
    document: TimedTextDocument,
    *,
    max_join_gap_ms: int = 1000,
) -> tuple[list[tuple[list[ReflowToken], tuple[str, ...]]], list[Finding], bool]:
    regions: list[tuple[list[ReflowToken], tuple[str, ...]]] = []
    findings: list[Finding] = []
    current: list[ReflowToken] = []
    current_wrappers: tuple[str, ...] = ()
    previous_cue: Cue | None = None
    used_estimation = False
    for cue in document.cues:
        plain, wrappers, markup_warnings = split_cue_wide_markup(cue.text)
        findings.extend(
            Finding("WARNING", "MARKUP_WARNING", warning, (cue.id,))
            for warning in markup_warnings
        )
        stored_word_timing = (
            document.timing_fidelity
            in {
                TimingFidelity.WORD_ALIGNED_EXACT,
                TimingFidelity.WORD_ALIGNED_ESTIMATED,
            }
            and bool(cue.words)
        )
        if stored_word_timing:
            tokens = [
                ReflowToken(word.text, word.start_ms, word.end_ms, cue.id, word)
                for word in cue.words
            ]
        else:
            tokens = _estimated_tokens(cue, plain)
            used_estimation = True
            warning = mixed_script_interpolation_warning(plain)
            if warning:
                findings.append(
                    Finding("WARNING", "MIXED_SCRIPT_ESTIMATION", warning, (cue.id,))
                )
        if not tokens:
            continue
        if previous_cue is not None and cue.start_ms < previous_cue.end_ms:
            raise ValueError(
                "Full reflow cannot safely resolve overlapping source captions. "
                "Review the overlap findings before reflowing."
            )
        starts_new_region = (
            previous_cue is not None
            and (
                cue.start_ms - previous_cue.end_ms > max_join_gap_ms
                or wrappers != current_wrappers
                or (
                    previous_cue.speaker
                    and cue.speaker
                    and previous_cue.speaker != cue.speaker
                )
                or (
                    previous_cue.source_track
                    and cue.source_track
                    and previous_cue.source_track != cue.source_track
                )
            )
        )
        if starts_new_region and current:
            regions.append((current, current_wrappers))
            current = []
        if not current:
            current_wrappers = wrappers
        current.extend(tokens)
        previous_cue = cue
    if current:
        regions.append((current, current_wrappers))
    return regions, findings, used_estimation


def _segment_cost(
    tokens: list[ReflowToken],
    start: int,
    end: int,
    profile: CaptionProfile,
    *,
    sentence_end_min_fill: float,
    comma_min_fill: float,
) -> tuple[float, str, list[str]]:
    values = [token.text for token in tokens[start:end]]
    text = _join_tokens(values)
    capacity = profile.max_units_per_line * profile.max_lines
    units = display_unit_count(text)
    duration = tokens[end - 1].end_ms - tokens[start].start_ms
    wrapped, warnings = wrap_text(
        text, profile.max_units_per_line, profile.max_lines
    )
    line_units = [display_unit_count(line) for line in wrapped.splitlines()]
    overflow = sum(max(0, value - profile.max_units_per_line) ** 2 for value in line_units)
    fill = units / max(capacity, 1)
    score = overflow * 100_000
    score += max(0.0, 1.0 - fill) ** 2 * 45
    if duration > profile.max_duration_ms:
        score += (duration - profile.max_duration_ms) ** 2
    elif duration < profile.min_duration_ms:
        score += ((profile.min_duration_ms - duration) / 100.0) ** 2
    if len(line_units) > 1:
        average = sum(line_units) / len(line_units)
        score += sum((value - average) ** 2 for value in line_units) * 0.2

    final = tokens[end - 1].text
    sentence_end = SENTENCE_END_RE.search(final) is not None
    comma_end = re.search(r"[,;:]$", final) is not None
    if sentence_end and fill >= sentence_end_min_fill:
        score -= 240
    elif comma_end and fill >= comma_min_fill:
        score -= 80
    else:
        score += 90
    if end < len(tokens):
        pause = tokens[end].start_ms - tokens[end - 1].end_ms
        if pause >= 250:
            score -= min(80, pause / 10)

    sentence_positions = [
        index for index in range(start, end)
        if SENTENCE_END_RE.search(tokens[index].text)
    ]
    if sentence_positions:
        words_after_last_sentence = end - sentence_positions[-1] - 1
        if words_after_last_sentence == 1 and not sentence_end:
            score += 1500
    normalized_final = re.sub(r"^[^\wÀ-ÿ]+|[^\wÀ-ÿ]+$", "", final, flags=re.UNICODE).casefold()
    if end < len(tokens) and normalized_final in WEAK_CAPTION_END_WORDS:
        score += 2000
    if end < len(tokens):
        normalized_next = re.sub(
            r"^[^\wÀ-ÿ]+|[^\wÀ-ÿ]+$",
            "",
            tokens[end].text,
            flags=re.UNICODE,
        ).casefold()
        if normalized_next in WEAK_CAPTION_START_WORDS:
            score += 1200
    if end - start == 1:
        genuine_one_word_utterance = sentence_end and (
            end == len(tokens)
            or tokens[end].start_ms - tokens[end - 1].end_ms >= 500
        )
        if not genuine_one_word_utterance:
            score += 1200
    return score, wrapped, warnings


def _is_shared_window_boundary(
    tokens: list[ReflowToken], boundary: int
) -> bool:
    previous_word = tokens[boundary - 1].timed_word
    next_word = tokens[boundary].timed_word
    return bool(
        previous_word is not None
        and next_word is not None
        and previous_word.timing_source
        == WordTimingSource.SHARED_ANCHOR_WINDOW
        and next_word.timing_source
        == WordTimingSource.SHARED_ANCHOR_WINDOW
        and next_word.start_ms < previous_word.end_ms
    )


def _optimize_region(
    tokens: list[ReflowToken],
    profile: CaptionProfile,
    *,
    sentence_end_min_fill: float,
    comma_min_fill: float,
) -> list[tuple[int, int, str, list[str]]] | None:
    capacity = profile.max_units_per_line * profile.max_lines
    count = len(tokens)
    best_cost = [float("inf")] * (count + 1)
    next_end = [-1] * (count + 1)
    segment_text: list[str] = [""] * (count + 1)
    segment_warnings: list[list[str]] = [[] for _ in range(count + 1)]
    best_cost[count] = 0.0
    for start in range(count - 1, -1, -1):
        for end in range(start + 1, count + 1):
            text = _join_tokens([token.text for token in tokens[start:end]])
            units = display_unit_count(text)
            duration = tokens[end - 1].end_ms - tokens[start].start_ms
            if end > start + 1 and (
                units > capacity or duration > profile.max_duration_ms
            ):
                break
            if end < count and _is_shared_window_boundary(tokens, end):
                # Overlapping shared-window words are one indivisible timing
                # run. A caption boundary inside it would produce overlapping
                # captions and invalidate the v2 timing contract.
                continue
            local_cost, wrapped, warnings = _segment_cost(
                tokens,
                start,
                end,
                profile,
                sentence_end_min_fill=sentence_end_min_fill,
                comma_min_fill=comma_min_fill,
            )
            total = local_cost + best_cost[end]
            if (
                total < best_cost[start]
                or (
                    abs(total - best_cost[start]) < 0.0001
                    and end > next_end[start]
                )
            ):
                best_cost[start] = total
                next_end[start] = end
                segment_text[start] = wrapped
                segment_warnings[start] = warnings
        if next_end[start] < 0:
            return None

    result: list[tuple[int, int, str, list[str]]] = []
    start = 0
    while start < count:
        end = next_end[start]
        result.append((start, end, segment_text[start], segment_warnings[start]))
        start = end
    return result


def _estimated_region_tokens(
    document: TimedTextDocument,
    tokens: list[ReflowToken],
) -> list[ReflowToken]:
    # Reflow regions are complete runs of source cues. Reconstructing fallback
    # timing cue by cue therefore preserves each source window and the region's
    # outer bounds. Revisit this if region formation ever stops following cue
    # boundaries.
    cue_ids = tuple(dict.fromkeys(token.source_cue_id for token in tokens))
    cues_by_id = {cue.id: cue for cue in document.cues}
    result: list[ReflowToken] = []
    for cue_id in cue_ids:
        cue = cues_by_id[cue_id]
        plain, _, _ = split_cue_wide_markup(cue.text)
        result.extend(_estimated_tokens(cue, plain))
    return result


def global_reflow(
    document: TimedTextDocument,
    profile: CaptionProfile,
    *,
    sentence_end_min_fill: float = 0.60,
    comma_min_fill: float = 0.80,
) -> tuple[TimedTextDocument, list[Finding]]:
    """Globally optimize caption boundaries inside continuous timed-text regions."""
    profile.validate()
    if not 0 <= sentence_end_min_fill <= 1 or not 0 <= comma_min_fill <= 1:
        raise ValueError("Punctuation fill thresholds must be between 0 and 1.")
    regions, findings, used_estimation = _reflow_regions(document)
    output: list[Cue] = []
    for tokens, wrappers in regions:
        segments = _optimize_region(
            tokens,
            profile,
            sentence_end_min_fill=sentence_end_min_fill,
            comma_min_fill=comma_min_fill,
        )
        if segments is None:
            has_shared_run = any(
                _is_shared_window_boundary(tokens, boundary)
                for boundary in range(1, len(tokens))
            )
            if not has_shared_run:
                raise ValueError("Could not find a complete caption reflow path.")
            source_ids = tuple(
                dict.fromkeys(token.source_cue_id for token in tokens)
            )
            tokens = _estimated_region_tokens(document, tokens)
            segments = _optimize_region(
                tokens,
                profile,
                sentence_end_min_fill=sentence_end_min_fill,
                comma_min_fill=comma_min_fill,
            )
            if segments is None:
                raise ValueError("Could not find a complete caption reflow path.")
            used_estimation = True
            findings.append(
                Finding(
                    "WARNING",
                    "SHARED_WORD_TIMING_FALLBACK",
                    "Stored estimated word timing could not satisfy the selected layout while keeping a shared timing window together. Reflow used character-weighted timing for this region instead, so no output word-timing metadata is created.",
                    source_ids,
                )
            )
        for start, end, wrapped, wrap_warnings in segments:
            selected = tokens[start:end]
            source_ids = tuple(dict.fromkeys(token.source_cue_id for token in selected))
            timed_words = [
                token.timed_word for token in selected if token.timed_word is not None
            ]
            output.append(
                Cue(
                    id=f"cue_{len(output) + 1:06d}",
                    srt_index=len(output) + 1,
                    start_ms=selected[0].start_ms,
                    end_ms=selected[-1].end_ms,
                    text=restore_cue_wide_markup(wrapped, wrappers),
                    words=timed_words,
                    metadata={"sourceCueIds": list(source_ids)},
                )
            )
            findings.extend(
                Finding("WARNING", "REFLOW_WARNING", warning, source_ids)
                for warning in wrap_warnings
            )
    result = TimedTextDocument(
        version=document.version,
        language=document.language,
        source_kind=document.source_kind,
        source_name=document.source_name,
        timing_fidelity=(
            TimingFidelity.ESTIMATED_WITHIN_GROUP
            if used_estimation
            else document.timing_fidelity
        ),
        cues=output,
        metadata=dict(document.metadata),
    )
    result.validate(allow_overlaps=False)
    findings.append(
        Finding(
            "INFO",
            "GLOBAL_REFLOW",
            (
                "Caption boundaries were globally optimized at verified word times."
                if not used_estimation
                else "Caption boundaries were globally optimized using deterministic estimated word times."
            ),
        )
    )
    return result, findings


def boundary_reflow(
    document: TimedTextDocument, profile: CaptionProfile
) -> tuple[TimedTextDocument, list[Finding]]:
    """Backward-compatible alias for the full global reflow operation."""
    return global_reflow(document, profile)
