from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any

from caption_profile import CaptionProfile, characters_per_second, display_unit_count
from operations import Finding, validate_document
from reference_script_alignment import (
    MatchKind,
    ReferenceScriptAlignment,
    ReferenceTimingFidelity,
    normalize_reference_whitespace,
)
from timed_text import (
    Cue,
    TimedTextDocument,
    TimedWord,
    TimingFidelity,
    WordTimingSource,
)
from transcript import join_transcript_word_text


SENTENCE_END_RE = re.compile(r"[.!?][\"')\]}]*$")
COMMA_END_RE = re.compile(r"[,;:][\"')\]}]*$")
CJK_RE = re.compile(r"[\u3040-\u30ff\u3400-\u9fff\uf900-\ufaff]")
TAG_LIKE_RE = re.compile(r"<[^>\r\n]+>")
LEADING_PUNCTUATION = frozenset("([{\"'“‘$€£¥#@")


@dataclass(frozen=True)
class ReferenceCaptionBuild:
    status: str
    status_message: str
    document: TimedTextDocument | None
    findings: tuple[Finding, ...]
    text_invariant_verified: bool
    exact_timing_token_count: int
    estimated_timing_token_count: int


@dataclass(frozen=True)
class _CaptionToken:
    reference_index: int
    text: str
    start_character: int
    end_character: int
    start_ms: int
    end_ms: int
    timing_fidelity: ReferenceTimingFidelity
    timing_source: WordTimingSource
    match_kind: MatchKind | None


def _plain_text_units(value: str) -> int:
    return display_unit_count(value)


def _split_unspaced_text(value: str, capacity: int) -> list[str]:
    if capacity < 1:
        raise ValueError("Line capacity must be positive.")
    lines: list[str] = []
    current = ""
    for character in value:
        candidate = current + character
        if current and _plain_text_units(candidate) > capacity:
            lines.append(current)
            current = character
        else:
            current = candidate
    if current:
        lines.append(current)
    return lines


def _wrap_plain_reference_text(
    text: str, max_units_per_line: int, max_lines: int
) -> tuple[str, list[str]]:
    """Wrap plain authoritative text without interpreting markup-like content."""
    value = normalize_reference_whitespace(text)
    if not value:
        return "", []
    warnings: list[str] = []
    words = value.split(" ")
    if len(words) == 1:
        if _plain_text_units(value) <= max_units_per_line:
            return value, warnings
        if CJK_RE.search(value):
            lines = _split_unspaced_text(value, max_units_per_line)
            if len(lines) > max_lines:
                warnings.append(
                    "Reference text cannot satisfy the configured line count without changing caption boundaries."
                )
            return "\n".join(lines), warnings
        warnings.append("An indivisible reference token exceeds the configured line length.")
        return value, warnings

    if any(_plain_text_units(word) > max_units_per_line for word in words):
        warnings.append("An indivisible reference token exceeds the configured line length.")
    line_count = min(
        max_lines,
        len(words),
        max(1, (_plain_text_units(value) + max_units_per_line - 1) // max_units_per_line),
    )
    if line_count == 1:
        return value, warnings
    prefixes = [0]
    for word in words:
        prefixes.append(prefixes[-1] + _plain_text_units(word))

    def units(start: int, end: int) -> int:
        return prefixes[end] - prefixes[start] + max(0, end - start - 1)

    target = _plain_text_units(value) / line_count
    states: dict[tuple[int, int], tuple[float, list[int]]] = {(0, 0): (0.0, [])}
    for used_lines in range(line_count):
        for position in range(len(words) + 1):
            state = states.get((used_lines, position))
            if state is None:
                continue
            for end in range(position + 1, len(words) + 1):
                length = units(position, end)
                overflow = max(0, length - max_units_per_line)
                score = state[0] + overflow * overflow * 1000 + (length - target) ** 2
                key = (used_lines + 1, end)
                candidate = (score, state[1] + [end])
                if key not in states or candidate[0] < states[key][0]:
                    states[key] = candidate
    _, boundaries = states[(line_count, len(words))]
    lines: list[str] = []
    start = 0
    for end in boundaries:
        lines.append(" ".join(words[start:end]))
        start = end
    if any(_plain_text_units(line) > max_units_per_line for line in lines):
        warnings.append(
            "Reference text cannot satisfy the configured line length without changing an indivisible token."
        )
    return "\n".join(lines), warnings


def _caption_tokens(alignment: ReferenceScriptAlignment) -> list[_CaptionToken]:
    timing_by_reference = {
        timing.reference_index: timing for timing in alignment.timings
    }
    match_by_reference = {
        match.reference_index: match for match in alignment.matches
    }
    display_starts = [0]
    for previous, current in zip(
        alignment.reference_tokens, alignment.reference_tokens[1:]
    ):
        gap = alignment.reference_text[previous.end_character : current.start_character]
        whitespace_positions = [
            position for position, character in enumerate(gap) if character.isspace()
        ]
        display_starts.append(
            (
                previous.end_character + whitespace_positions[-1] + 1
                if whitespace_positions
                else previous.end_character
                if gap and all(character in LEADING_PUNCTUATION for character in gap)
                else current.start_character
            )
        )
    tokens: list[_CaptionToken] = []
    for index, reference in enumerate(alignment.reference_tokens):
        timing = timing_by_reference[index]
        if timing.start_ms is None or timing.end_ms is None:
            raise ValueError(
                "Reference caption construction requires positive timing for every token."
            )
        start_character = display_starts[index]
        end_character = (
            display_starts[index + 1]
            if index + 1 < len(alignment.reference_tokens)
            else len(alignment.reference_text)
        )
        text = alignment.reference_text[start_character:end_character].strip()
        if not text:
            raise ValueError("A reference token produced an empty authoritative text span.")
        match = match_by_reference.get(index)
        if timing.timing_source is None:
            raise ValueError(
                "Positive reference timing requires explicit word timing provenance."
            )
        tokens.append(
            _CaptionToken(
                reference_index=index,
                text=text,
                start_character=start_character,
                end_character=end_character,
                start_ms=timing.start_ms,
                end_ms=timing.end_ms,
                timing_fidelity=timing.fidelity,
                timing_source=timing.timing_source,
                match_kind=match.kind if match else None,
            )
        )
    return tokens


def _segment_text(
    alignment: ReferenceScriptAlignment,
    tokens: list[_CaptionToken],
    start: int,
    end: int,
) -> tuple[str, int, int]:
    start_character = tokens[start].start_character
    end_character = tokens[end - 1].end_character
    return (
        alignment.reference_text[start_character:end_character].strip(),
        start_character,
        end_character,
    )


def _segment_score(
    text: str,
    wrapped: str,
    duration_ms: int,
    profile: CaptionProfile,
    *,
    sentence_end_min_fill: float,
    comma_min_fill: float,
) -> float:
    line_units = [_plain_text_units(line) for line in wrapped.splitlines()]
    capacity = profile.max_units_per_line * profile.max_lines
    units = _plain_text_units(text)
    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_ms > profile.max_duration_ms:
        score += (duration_ms - profile.max_duration_ms) ** 2
    elif duration_ms < profile.min_duration_ms:
        score += ((profile.min_duration_ms - duration_ms) / 100.0) ** 2
    cps = characters_per_second(text, max(1, duration_ms))
    if cps > profile.max_cps:
        score += (cps - profile.max_cps) ** 2 * 20
    if len(line_units) > 1:
        average = sum(line_units) / len(line_units)
        score += sum((value - average) ** 2 for value in line_units) * 0.2
    if SENTENCE_END_RE.search(text) and fill >= sentence_end_min_fill:
        score -= 240
    elif COMMA_END_RE.search(text) and fill >= comma_min_fill:
        score -= 80
    else:
        score += 90
    return score


def _optimize_region(
    alignment: ReferenceScriptAlignment,
    tokens: list[_CaptionToken],
    profile: CaptionProfile,
    *,
    sentence_end_min_fill: float,
    comma_min_fill: float,
) -> list[tuple[int, int, str, int, int, list[str]]] | None:
    count = len(tokens)
    capacity = profile.max_units_per_line * profile.max_lines
    best_cost = [float("inf")] * (count + 1)
    choices: list[tuple[int, str, int, int, list[str]] | None] = [None] * (count + 1)
    best_cost[count] = 0.0
    for start in range(count - 1, -1, -1):
        for end in range(start + 1, count + 1):
            text, start_character, end_character = _segment_text(
                alignment, tokens, start, end
            )
            duration = tokens[end - 1].end_ms - tokens[start].start_ms
            units = _plain_text_units(text)
            if end > start + 1 and (
                units > capacity or duration > profile.max_duration_ms
            ):
                break
            wrapped, warnings = _wrap_plain_reference_text(
                text, profile.max_units_per_line, profile.max_lines
            )
            local = _segment_score(
                text,
                wrapped,
                duration,
                profile,
                sentence_end_min_fill=sentence_end_min_fill,
                comma_min_fill=comma_min_fill,
            )
            if end < count:
                pause = tokens[end].start_ms - tokens[end - 1].end_ms
                if pause >= 250:
                    local -= min(80, pause / 10)
                if tokens[end].start_ms < tokens[end - 1].end_ms:
                    # Tokens originating in one very short Premiere word can
                    # share a window. They must stay in one cue rather than
                    # creating overlapping SRT captions.
                    continue
            if end - start == 1 and not SENTENCE_END_RE.search(text):
                local += 1200
            total = local + best_cost[end]
            current = choices[start]
            if total < best_cost[start] or (
                abs(total - best_cost[start]) < 0.0001
                and (current is None or end > current[0])
            ):
                best_cost[start] = total
                choices[start] = (
                    end,
                    wrapped,
                    start_character,
                    end_character,
                    warnings,
                )
        if choices[start] is None:
            return None
    output: list[tuple[int, int, str, int, int, list[str]]] = []
    start = 0
    while start < count:
        choice = choices[start]
        assert choice is not None
        end, wrapped, start_character, end_character, warnings = choice
        output.append(
            (start, end, wrapped, start_character, end_character, warnings)
        )
        start = end
    return output


def _regions(tokens: list[_CaptionToken], *, max_join_gap_ms: int) -> list[tuple[int, int]]:
    if not tokens:
        return []
    regions: list[tuple[int, int]] = []
    start = 0
    for index in range(1, len(tokens)):
        if tokens[index].start_ms - tokens[index - 1].end_ms > max_join_gap_ms:
            regions.append((start, index))
            start = index
    regions.append((start, len(tokens)))
    return regions


def verify_reference_text_invariant(
    document: TimedTextDocument, reference_text: str
) -> tuple[bool, str]:
    normalized_reference = normalize_reference_whitespace(reference_text)
    cursor = 0
    for cue in document.cues:
        start = cue.metadata.get("referenceStartCharacter")
        end = cue.metadata.get("referenceEndCharacter")
        if not isinstance(start, int) or not isinstance(end, int):
            return False, f"{cue.id} has no authoritative reference character range."
        if start != cursor or end <= start or end > len(normalized_reference):
            return False, f"{cue.id} does not continue the authoritative reference range."
        expected = normalize_reference_whitespace(normalized_reference[start:end])
        actual_lines = [normalize_reference_whitespace(line) for line in cue.text.splitlines()]
        actual_with_breaks = "\n".join(actual_lines)
        expected_position = 0
        matches = True
        for character in actual_with_breaks:
            if character == "\n":
                if expected_position < len(expected) and expected[expected_position] == " ":
                    expected_position += 1
                continue
            if (
                expected_position >= len(expected)
                or expected[expected_position] != character
            ):
                matches = False
                break
            expected_position += 1
        if expected_position != len(expected):
            matches = False
        if not matches:
            return False, (
                f"{cue.id} differs from its authoritative reference range: "
                f"expected {expected!r}, got {actual_with_breaks!r}."
            )
        cursor = end
    if cursor != len(normalized_reference):
        return False, "Caption ranges do not cover the complete authoritative reference text."
    return True, "Every caption exactly preserves its authoritative reference range."


def build_reference_caption_document(
    alignment: ReferenceScriptAlignment,
    *,
    profile: CaptionProfile,
    language: str = "",
    source_name: str = "",
    sentence_end_min_fill: float = 0.60,
    comma_min_fill: float = 0.80,
    max_join_gap_ms: int = 1000,
) -> ReferenceCaptionBuild:
    profile.validate()
    if not alignment.can_generate_srt:
        return ReferenceCaptionBuild(
            alignment.readiness.value,
            alignment.reliability_message,
            None,
            (),
            False,
            0,
            0,
        )
    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.")
    if max_join_gap_ms < 0:
        raise ValueError("Maximum join gap must be non-negative.")
    if TAG_LIKE_RE.search(alignment.reference_text):
        finding = Finding(
            "ERROR",
            "UNSUPPORTED_REFERENCE_MARKUP",
            "The approved text contains angle-bracket text that an SRT importer may treat as hidden formatting. Remove or replace it before creating captions so the visible wording can be guaranteed.",
        )
        return ReferenceCaptionBuild(
            "UNSUPPORTED_REFERENCE_MARKUP",
            finding.message,
            None,
            (finding,),
            False,
            0,
            0,
        )
    tokens = _caption_tokens(alignment)
    exact_timing_count = sum(
        token.timing_source == WordTimingSource.TRANSCRIPT_EXACT for token in tokens
    )
    estimated_timing_count = len(tokens) - exact_timing_count
    can_store_word_timing = all(
        current.start_ms >= previous.end_ms
        or (
            previous.timing_source == WordTimingSource.SHARED_ANCHOR_WINDOW
            and current.timing_source == WordTimingSource.SHARED_ANCHOR_WINDOW
        )
        for previous, current in zip(tokens, tokens[1:])
    )
    findings: list[Finding] = []
    cues: list[Cue] = []
    for region_start, region_end in _regions(tokens, max_join_gap_ms=max_join_gap_ms):
        region_tokens = tokens[region_start:region_end]
        layout = _optimize_region(
            alignment,
            region_tokens,
            profile,
            sentence_end_min_fill=sentence_end_min_fill,
            comma_min_fill=comma_min_fill,
        )
        if layout is None:
            finding = Finding(
                "ERROR",
                "REFERENCE_LAYOUT_IMPOSSIBLE",
                "The approved text cannot be divided into valid captions with "
                "the current layout and duration settings without changing its "
                "wording or separating words that must remain together. Increase "
                "the maximum text length per line, number of lines, or maximum "
                "caption duration, then run the tool again.",
            )
            return ReferenceCaptionBuild(
                "REFERENCE_LAYOUT_IMPOSSIBLE",
                finding.message,
                None,
                (finding,),
                False,
                exact_timing_count,
                estimated_timing_count,
            )
        for start, end, wrapped, start_character, end_character, warnings in layout:
            selected = region_tokens[start:end]
            cue_id = f"cue_{len(cues) + 1:06d}"
            word_values = (
                [
                    TimedWord(
                        token.text,
                        token.start_ms,
                        token.end_ms,
                        timing_source=token.timing_source,
                    )
                    for token in selected
                ]
                if can_store_word_timing
                else []
            )
            cue = Cue(
                id=cue_id,
                srt_index=len(cues) + 1,
                start_ms=selected[0].start_ms,
                end_ms=selected[-1].end_ms,
                text=wrapped,
                words=word_values,
                metadata={
                    "referenceStartCharacter": start_character,
                    "referenceEndCharacter": end_character,
                    "referenceTokenStart": selected[0].reference_index,
                    "referenceTokenEnd": selected[-1].reference_index + 1,
                    "timingFidelity": sorted(
                        {token.timing_fidelity.value for token in selected}
                    ),
                    "timingSources": sorted(
                        {token.timing_source.value for token in selected}
                    ),
                    "matchKinds": sorted(
                        {
                            token.match_kind.value
                            if token.match_kind is not None
                            else "UNMATCHED_REFERENCE"
                            for token in selected
                        }
                    ),
                },
            )
            cues.append(cue)
            findings.extend(
                Finding("WARNING", "REFERENCE_LAYOUT_WARNING", warning, (cue_id,))
                for warning in warnings
            )
    timing_fidelity = (
        TimingFidelity.WORD_ALIGNED_EXACT
        if can_store_word_timing and estimated_timing_count == 0
        else TimingFidelity.WORD_ALIGNED_ESTIMATED
        if can_store_word_timing
        else TimingFidelity.ESTIMATED_WITHIN_GROUP
    )
    document = TimedTextDocument(
        language=language,
        source_kind="reference-script-active-sequence-transcripts",
        source_name=source_name,
        timing_fidelity=timing_fidelity,
        cues=cues,
        metadata={
            "authoritativeReferenceText": alignment.reference_text,
            "exactTimingTokenCount": exact_timing_count,
            "estimatedTimingTokenCount": estimated_timing_count,
        },
    )
    if can_store_word_timing:
        # The standard timing sidecar requires its word join to reconstruct each
        # cue. CJK and unusual reference spacing can be exact at the token level
        # while not satisfying that older Latin-oriented join convention.
        for cue in document.cues:
            joined = join_transcript_word_text(
                [{"text": word.text} for word in cue.words]
            )
            if normalize_reference_whitespace(joined) != normalize_reference_whitespace(
                cue.text
            ):
                can_store_word_timing = False
                break
        if not can_store_word_timing:
            document.timing_fidelity = TimingFidelity.ESTIMATED_WITHIN_GROUP
            for cue in document.cues:
                cue.words = []
            findings.append(
                Finding(
                    "INFO",
                    "WORD_TIMING_SIDECAR_NOT_AVAILABLE",
                    "Reference timing is retained in the comparison report, but the word-timing sidecar cannot reconstruct this script's spacing safely.",
                )
            )
    document.validate(allow_overlaps=False)
    invariant_ok, invariant_message = verify_reference_text_invariant(
        document, alignment.reference_text
    )
    if not invariant_ok:
        raise ValueError(
            "Reference caption text invariant failed; no SRT may be written. "
            + invariant_message
        )
    findings.append(
        Finding("INFO", "REFERENCE_TEXT_PRESERVED", invariant_message)
    )
    findings.extend(validate_document(document, profile))
    status = (
        "SUCCEEDED_WITH_WARNINGS"
        if any(finding.severity == "WARNING" for finding in findings)
        else "SUCCEEDED"
    )
    return ReferenceCaptionBuild(
        status,
        invariant_message,
        document,
        tuple(findings),
        True,
        exact_timing_count,
        estimated_timing_count,
    )


def reference_caption_build_report(build: ReferenceCaptionBuild) -> dict[str, Any]:
    return {
        "version": 1,
        "status": build.status,
        "statusMessage": build.status_message,
        "outputCreated": build.document is not None,
        "cueCount": len(build.document.cues) if build.document else 0,
        "timingFidelity": (
            build.document.timing_fidelity.value if build.document else None
        ),
        "textInvariantVerified": build.text_invariant_verified,
        "exactTimingTokenCount": build.exact_timing_token_count,
        "estimatedTimingTokenCount": build.estimated_timing_token_count,
        "authoritativeReferenceText": (
            build.document.metadata.get("authoritativeReferenceText")
            if build.document
            else None
        ),
        "cues": [
            {
                "id": cue.id,
                "srtIndex": cue.srt_index,
                "startMs": cue.start_ms,
                "endMs": cue.end_ms,
                "text": cue.text,
                "referenceRange": [
                    cue.metadata.get("referenceStartCharacter"),
                    cue.metadata.get("referenceEndCharacter"),
                ],
                "referenceTokenRange": [
                    cue.metadata.get("referenceTokenStart"),
                    cue.metadata.get("referenceTokenEnd"),
                ],
                "timingFidelity": cue.metadata.get("timingFidelity", []),
                "timingSources": cue.metadata.get("timingSources", []),
                "matchKinds": cue.metadata.get("matchKinds", []),
            }
            for cue in build.document.cues
        ]
        if build.document
        else [],
        "findings": [
            {
                "severity": finding.severity,
                "code": finding.code,
                "message": finding.message,
                "cueIds": list(finding.cue_ids),
            }
            for finding in build.findings
        ],
    }


def reference_caption_build_report_markdown(
    build: ReferenceCaptionBuild, *, heading_level: int = 1
) -> str:
    if heading_level < 1 or heading_level > 5:
        raise ValueError("Heading level must be between 1 and 5.")
    heading = "#" * heading_level
    section = "#" * (heading_level + 1)
    report = reference_caption_build_report(build)
    lines = [
        f"{heading} Reference Caption Build Report",
        "",
        f"- Result: `{build.status}`",
        f"- SRT captions created: {report['cueCount']}",
        (
            "- Approved reference text preserved exactly: yes"
            if build.text_invariant_verified
            else "- Approved reference text preserved exactly: not verified"
        ),
        f"- Direct transcript timing anchors: {build.exact_timing_token_count}",
        f"- Estimated or near-match timing tokens: {build.estimated_timing_token_count}",
        f"- Timing metadata quality: `{report['timingFidelity'] or 'none'}`",
        "",
    ]
    if build.document is None:
        lines.append("No SRT document was created.")
    else:
        lines.extend([f"{section} Created captions", ""])
        for cue in report["cues"]:
            lines.extend(
                [
                    f"- **{cue['id']}** — {cue['startMs']}–{cue['endMs']} ms",
                    f"  - Text: {cue['text'].replace(chr(10), ' / ')}",
                    f"  - Reference characters: {cue['referenceRange'][0]}–{cue['referenceRange'][1]}",
                    f"  - Timing source: {', '.join(cue['timingFidelity'])}",
                    f"  - Word timing provenance: {', '.join(cue['timingSources'])}",
                    f"  - Match types: {', '.join(cue['matchKinds'])}",
                ]
            )
    if build.findings:
        lines.extend(["", f"{section} Notes", ""])
        for finding in build.findings:
            lines.append(
                f"- **{finding.code.replace('_', ' ').title()}**: {finding.message}"
            )
    return "\n".join(lines).rstrip() + "\n"
