"""Batch ownership remapping across source-line spaces."""

from __future__ import annotations

from typing import Protocol

from ...core.line_selection import LineRanges, LineSelection
from ...core.coordinates import BatchSourceSpace
from ..line_matching.lineage import BatchSourceLineage
from ..line_matching.transforms import BatchSourceExactTransform
from .model import BatchOwnership
from .absence_claims import AbsenceClaim
from .claims import (
    format_ownership_line_set,
    parse_ownership_line_ranges,
    presence_claims_from_source_lines,
)
from .replacement_units import (
    NoReplacementUnitOrigin,
    ReplacementUnit,
    normalize_replacement_units,
)


class _SourceLineTranslator(Protocol):
    def first_unmapped_source_line(
        self,
        selection: LineSelection,
    ) -> int | None: ...

    def translate_source_selection(
        self,
        selection: LineSelection,
    ) -> LineRanges: ...

    def translate_source_line(self, line_number: int) -> int | None: ...


def _first_unmapped_line(
    line_selection: LineSelection,
    lineage: _SourceLineTranslator,
) -> int | None:
    return lineage.first_unmapped_source_line(line_selection)


def _remap_replacement_units_with_lineage(
    replacement_units: list[ReplacementUnit],
    *,
    lineage: _SourceLineTranslator,
    deletion_count: int,
) -> list[ReplacementUnit]:
    """Remap replacement-unit presence lines with refreshed source lineage."""
    remapped_units: list[ReplacementUnit] = []

    for unit in replacement_units:
        old_presence_lines = parse_ownership_line_ranges(unit.presence_lines)
        first_unmapped = _first_unmapped_line(old_presence_lines, lineage)
        if first_unmapped is not None:
            raise ValueError(
                f"Cannot remap replacement unit presence line {first_unmapped} "
                f"from old source to new source: no unique mapping found."
            )

        origin_evidence = unit.origin_evidence
        origin = unit.origin
        if origin is not None:
            old_origin_lines = LineRanges.from_ranges(
                ((origin.new_start, origin.new_end),)
            )
            first_unmapped_origin = _first_unmapped_line(
                old_origin_lines,
                lineage,
            )
            if first_unmapped_origin is None:
                translated_origin_lines = lineage.translate_source_selection(
                    old_origin_lines
                )
                translated_ranges = translated_origin_lines.ranges()
                if len(translated_ranges) == 1:
                    translated_start, translated_end = translated_ranges[0]
                    origin_evidence = unit.with_origin(
                        origin.with_new_lines(translated_start, translated_end)
                    ).origin_evidence
                else:
                    origin_evidence = NoReplacementUnitOrigin()
            else:
                origin_evidence = NoReplacementUnitOrigin()

        remapped_units.append(ReplacementUnit(
            presence_lines=format_ownership_line_set(
                lineage.translate_source_selection(old_presence_lines)
            ),
            deletion_indices=unit.deletion_indices,
            origin_evidence=origin_evidence,
        ))

    return normalize_replacement_units(
        remapped_units,
        deletion_count=deletion_count,
    )


def remap_batch_ownership_with_lineage(
    ownership: BatchOwnership,
    lineage: BatchSourceLineage,
) -> BatchOwnership:
    """Remap ownership using provenance from source refresh construction."""
    return _remap_batch_ownership(ownership, lineage)


def remap_batch_ownership_with_transform(
    ownership: BatchOwnership,
    transform: BatchSourceExactTransform[
        BatchSourceSpace,
        BatchSourceSpace,
    ],
) -> BatchOwnership:
    """Remap ownership through snapshot-bound exact source lineage."""
    return _remap_batch_ownership(ownership, transform)


def _remap_batch_ownership(
    ownership: BatchOwnership,
    lineage: _SourceLineTranslator,
) -> BatchOwnership:
    old_presence = ownership.presence_line_set()
    first_unmapped = _first_unmapped_line(old_presence, lineage)
    if first_unmapped is not None:
        raise ValueError(
            f"Cannot remap presence line {first_unmapped} from old source to new source: "
            f"no preserved source lineage found."
        )
    new_presence = lineage.translate_source_selection(old_presence)

    new_deletions = []
    for deletion in ownership.deletions:
        if deletion.anchor_line is None:
            new_deletions.append(AbsenceClaim(
                anchor_line=None,
                content_lines=deletion.content_lines,
                baseline_reference=deletion.baseline_reference,
                source_alternative=deletion.source_alternative,
            ))
            continue

        new_anchor = lineage.translate_source_line(deletion.anchor_line)
        if new_anchor is None:
            raise ValueError(
                f"Cannot remap deletion anchor line {deletion.anchor_line} from old source "
                f"to new source: no preserved source lineage found."
            )
        new_deletions.append(AbsenceClaim(
            anchor_line=new_anchor,
            content_lines=deletion.content_lines,
            baseline_reference=deletion.baseline_reference,
            source_alternative=deletion.source_alternative,
        ))

    new_replacement_units = _remap_replacement_units_with_lineage(
        ownership.replacement_units,
        lineage=lineage,
        deletion_count=len(new_deletions),
    )

    new_presence_baseline_references = {}
    for old_line_num, reference in ownership.presence_baseline_references().items():
        new_line_num = lineage.translate_source_line(old_line_num)
        if new_line_num is not None:
            new_presence_baseline_references[new_line_num] = reference

    return BatchOwnership(
        presence_claims=presence_claims_from_source_lines(
            new_presence,
            new_presence_baseline_references,
        ),
        deletions=new_deletions,
        replacement_units=new_replacement_units,
    )
