"""Ownership unit construction, filtering, and rebuild helpers."""

from __future__ import annotations

from collections.abc import Sequence
from typing import Literal, TypedDict, overload

from ...core.line_selection import LineRanges, LineSelection
from .display_lines import (
    OwnershipDisplayLine,
    build_display_lines_from_batch_source_lines,
)
from .model import (
    BatchOwnership,
)
from .claims import (
    LineRangeBuilder,
    parse_ownership_line_ranges,
)
from .unit_types import (
    OwnershipUnit as _UnitRecord,
    OwnershipUnitKind as _UnitKind,
)
from .replacement_units import normalize_replacement_units
from .references import BaselineReference


class _ClaimedDisplayRun(TypedDict):
    """Consecutive claimed lines in the ownership display."""

    display_ids: list[int]
    source_lines: list[int]
    next_index: int


class _DeletionDisplayRun(TypedDict):
    """Consecutive deleted lines in the ownership display."""

    display_ids: list[int]
    deletion_indices: list[int]
    next_index: int


def _presence_references_for_lines(
    ownership: BatchOwnership,
    source_lines: LineSelection,
) -> dict[int, BaselineReference]:
    """Return baseline references owned by one semantic unit."""
    return {
        line: reference
        for line, reference in ownership.presence_baseline_references().items()
        if line in source_lines
    }


def build_ownership_units_from_display_lines(
    ownership: BatchOwnership,
    display_lines: list[OwnershipDisplayLine],
) -> list[_UnitRecord]:
    """Build semantic ownership units from already reconstructed display lines.

    This is the fast path for callers that already need display lines for
    rendering.  It preserves the same grouping rules as
    build_ownership_units_from_batch_source_lines() without rebuilding the
    display model.
    """
    units, consumed_claimed_lines, consumed_deletion_indices = (
        _build_explicit_replacement_units_from_display_lines(
            ownership,
            display_lines,
        )
    )
    i = 0
    claimed_run: _ClaimedDisplayRun

    while i < len(display_lines):
        line = display_lines[i]
        if _display_line_is_consumed(
            line,
            consumed_claimed_lines,
            consumed_deletion_indices,
        ):
            i += 1
            continue

        if line["type"] == "deletion":
            # Collect consecutive deletion block
            deletion_run = _collect_display_run(
                display_lines,
                i,
                "deletion",
                consumed_claimed_lines,
                consumed_deletion_indices,
            )
            i = deletion_run["next_index"]

            # Check if immediately followed by claimed line (display adjacency)
            if (
                i < len(display_lines)
                and display_lines[i]["type"] == "claimed"
                and not _display_line_is_consumed(
                    display_lines[i],
                    consumed_claimed_lines,
                    consumed_deletion_indices,
                )
            ):
                # Collect single claimed line (to preserve fine-grained reset)
                claimed_display_id = _required_display_id(display_lines[i])
                claimed_source_line = display_lines[i]["source_line"]
                i += 1

                # Replacement unit: deletion block adjacent to single claimed line
                claimed_run = {
                    "display_ids": [claimed_display_id],
                    "source_lines": [claimed_source_line],
                    "next_index": i,
                }
                units.append(_build_replacement_unit(
                    ownership=ownership,
                    deletion_run=deletion_run,
                    claimed_run=claimed_run
                ))
            else:
                # Deletion-only unit: no adjacent claimed block
                units.append(_build_deletion_only_unit(
                    ownership=ownership,
                    deletion_run=deletion_run
                ))

        elif line["type"] == "claimed":
            # Collect single claimed line (not a block, to preserve fine-grained reset)
            claimed_display_id = _required_display_id(line)
            claimed_source_line = line["source_line"]
            i += 1

            # Check if immediately followed by deletion block (display adjacency)
            if (
                i < len(display_lines)
                and display_lines[i]["type"] == "deletion"
                and not _display_line_is_consumed(
                    display_lines[i],
                    consumed_claimed_lines,
                    consumed_deletion_indices,
                )
            ):
                # Collect consecutive deletion block
                deletion_run = _collect_display_run(
                    display_lines,
                    i,
                    "deletion",
                    consumed_claimed_lines,
                    consumed_deletion_indices,
                )
                i = deletion_run["next_index"]

                # Replacement unit: claimed line adjacent to deletion block
                claimed_run = {
                    "display_ids": [claimed_display_id],
                    "source_lines": [claimed_source_line],
                    "next_index": i,
                }
                units.append(_build_replacement_unit(
                    ownership=ownership,
                    deletion_run=deletion_run,
                    claimed_run=claimed_run
                ))
            else:
                # Presence-only unit: one claimed line without adjacent deletions
                # One unit per line allows independent reset
                units.append(_UnitRecord(
                    kind=_UnitKind.PRESENCE_ONLY,
                    claimed_source_lines=LineRanges.from_lines([claimed_source_line]),
                    deletion_claims=[],
                    display_line_ids=LineRanges.from_lines([claimed_display_id]),
                    baseline_references=_presence_references_for_lines(
                        ownership,
                        LineRanges.from_lines([claimed_source_line]),
                    ),
                    is_atomic=False,
                ))
        else:
            # Unknown type - skip
            i += 1

    return sorted(units, key=_ownership_unit_display_order_key)


def _ownership_unit_display_order_key(unit: _UnitRecord) -> int:
    """Return the first visible display line covered by a semantic unit."""
    return unit.display_line_ids.first() or 10**12


def _required_display_id(display_line: OwnershipDisplayLine) -> int:
    """Return the selectable ID guaranteed for claimed/deletion lines."""
    display_id = display_line.get("id")
    if display_id is None:
        raise ValueError("selectable ownership display line has no ID")
    return display_id


def _build_explicit_replacement_units_from_display_lines(
    ownership: BatchOwnership,
    display_lines: list[OwnershipDisplayLine],
) -> tuple[list[_UnitRecord], LineRanges, set[int]]:
    """Build units from persisted replacement metadata."""
    units: list[_UnitRecord] = []
    consumed_claimed_lines = LineRanges.empty()
    consumed_deletion_indices: set[int] = set()

    replacement_units = normalize_replacement_units(
        ownership.replacement_units,
        deletion_count=len(ownership.deletions),
    )
    if not replacement_units:
        return units, consumed_claimed_lines, consumed_deletion_indices

    for replacement_unit in replacement_units:
        claimed_source_lines = parse_ownership_line_ranges(
            replacement_unit.presence_lines
        )
        deletion_indices = set(replacement_unit.deletion_indices)
        claimed_display_id_builder = LineRangeBuilder()
        deletion_display_id_builder = LineRangeBuilder()

        for display_line in display_lines:
            display_id = display_line.get("id")
            if display_id is None:
                continue

            if (
                display_line["type"] == "claimed"
                and display_line["source_line"] in claimed_source_lines
            ):
                claimed_display_id_builder.add_line(display_id)
            elif (
                display_line["type"] == "deletion"
                and display_line["deletion_index"] in deletion_indices
            ):
                deletion_display_id_builder.add_line(display_id)

        claimed_display_ids = claimed_display_id_builder.finish()
        deletion_display_ids = deletion_display_id_builder.finish()
        if not claimed_display_ids or not deletion_display_ids:
            continue

        deletion_claims = [
            ownership.deletions[index]
            for index in sorted(deletion_indices)
        ]
        units.append(_UnitRecord(
            kind=_UnitKind.REPLACEMENT,
            claimed_source_lines=claimed_source_lines,
            deletion_claims=deletion_claims,
            display_line_ids=claimed_display_ids.union(deletion_display_ids),
            baseline_references=_presence_references_for_lines(
                ownership,
                claimed_source_lines,
            ),
            is_atomic=True,
            preserves_replacement_unit=True,
            replacement_origin=replacement_unit.origin,
        ))
        consumed_claimed_lines = consumed_claimed_lines.union(claimed_source_lines)
        consumed_deletion_indices.update(deletion_indices)

    return units, consumed_claimed_lines, consumed_deletion_indices


def _display_line_is_consumed(
    display_line: OwnershipDisplayLine,
    consumed_claimed_lines: LineSelection,
    consumed_deletion_indices: set[int],
) -> bool:
    """Return True when a display line is already covered by an explicit unit."""
    if display_line["type"] == "claimed":
        return display_line["source_line"] in consumed_claimed_lines
    if display_line["type"] == "deletion":
        return display_line["deletion_index"] in consumed_deletion_indices
    return False


@overload
def _collect_display_run(
    display_lines: list[OwnershipDisplayLine],
    start_index: int,
    expected_type: Literal["claimed"],
    consumed_claimed_lines: LineSelection,
    consumed_deletion_indices: set[int],
) -> _ClaimedDisplayRun: ...


@overload
def _collect_display_run(
    display_lines: list[OwnershipDisplayLine],
    start_index: int,
    expected_type: Literal["deletion"],
    consumed_claimed_lines: LineSelection,
    consumed_deletion_indices: set[int],
) -> _DeletionDisplayRun: ...


def _collect_display_run(
    display_lines: list[OwnershipDisplayLine],
    start_index: int,
    expected_type: Literal["claimed", "deletion"],
    consumed_claimed_lines: LineSelection,
    consumed_deletion_indices: set[int],
) -> _ClaimedDisplayRun | _DeletionDisplayRun:
    """Collect a consecutive run of display lines of the same type.

    Args:
        display_lines: List of display line dicts
        start_index: Starting index in display_lines
        expected_type: Expected line type ("deletion" or "claimed")

    Returns:
        Dict with:
        - display_ids: List of display IDs in the run
        - source_lines: List of source lines (for claimed) or None
        - deletion_indices: List of deletion indices (for deletion) or None
        - next_index: Index of first line after the run
    """
    display_ids: list[int] = []
    i = start_index
    if expected_type == "claimed":
        source_lines: list[int] = []
        while (
            i < len(display_lines)
            and display_lines[i]["type"] == expected_type
            and not _display_line_is_consumed(
                display_lines[i],
                consumed_claimed_lines,
                consumed_deletion_indices,
            )
        ):
            display_ids.append(_required_display_id(display_lines[i]))
            source_lines.append(display_lines[i]["source_line"])
            i += 1
        return {
            "display_ids": display_ids,
            "source_lines": source_lines,
            "next_index": i,
        }

    deletion_indices: list[int] = []
    while (
        i < len(display_lines)
        and display_lines[i]["type"] == expected_type
        and not _display_line_is_consumed(
            display_lines[i],
            consumed_claimed_lines,
            consumed_deletion_indices,
        )
    ):
        display_ids.append(_required_display_id(display_lines[i]))
        deletion_indices.append(display_lines[i]["deletion_index"])
        i += 1
    return {
        "display_ids": display_ids,
        "deletion_indices": deletion_indices,
        "next_index": i,
    }


def _build_replacement_unit(
    ownership: BatchOwnership,
    deletion_run: _DeletionDisplayRun,
    claimed_run: _ClaimedDisplayRun,
) -> _UnitRecord:
    """Build a REPLACEMENT unit from adjacent deletion and claimed runs.

    Args:
        ownership: BatchOwnership containing absence claims
        deletion_run: Dict from _collect_display_run for deletions
        claimed_run: Dict from _collect_display_run for claimed lines

    Returns:
        REPLACEMENT OwnershipUnit (atomic)
    """
    deletion_claims = [
        ownership.deletions[idx]
        for idx in set(deletion_run["deletion_indices"])
    ]

    claimed_source_lines = LineRanges.from_lines(claimed_run["source_lines"])
    return _UnitRecord(
        kind=_UnitKind.REPLACEMENT,
        claimed_source_lines=claimed_source_lines,
        deletion_claims=deletion_claims,
        display_line_ids=LineRanges.from_lines(
            [*deletion_run["display_ids"], *claimed_run["display_ids"]]
        ),
        baseline_references=_presence_references_for_lines(
            ownership,
            claimed_source_lines,
        ),
        is_atomic=True,
    )


def _build_deletion_only_unit(
    ownership: BatchOwnership,
    deletion_run: _DeletionDisplayRun,
) -> _UnitRecord:
    """Build a DELETION_ONLY unit from a deletion run with no adjacent claimed lines.

    Args:
        ownership: BatchOwnership containing absence claims
        deletion_run: Dict from _collect_display_run for deletions

    Returns:
        DELETION_ONLY OwnershipUnit (atomic)
    """
    deletion_claims = [
        ownership.deletions[idx]
        for idx in set(deletion_run["deletion_indices"])
    ]

    return _UnitRecord(
        kind=_UnitKind.DELETION_ONLY,
        claimed_source_lines=LineRanges.empty(),
        deletion_claims=deletion_claims,
        display_line_ids=LineRanges.from_lines(deletion_run["display_ids"]),
        is_atomic=True,
    )


def build_ownership_units_from_batch_source_lines(
    ownership: BatchOwnership,
    batch_source_lines: Sequence[bytes],
) -> list[_UnitRecord]:
    """Build semantic ownership units from indexed batch-source lines.

    Persisted replacement metadata is honored first, so captured replacements
    remain whole atomic units even if their lines are no longer display-adjacent.
    Remaining lines fall back to display-adjacency grouping in reconstructed
    display order, not source-line proximity. This reflects what the user
    actually sees in the batch display.
    """
    display_lines = build_display_lines_from_batch_source_lines(
        batch_source_lines,
        ownership,
    )
    return build_ownership_units_from_display_lines(ownership, display_lines)
