"""Line-replacement support for discard commands."""

from __future__ import annotations

from collections.abc import Iterator
from contextlib import ExitStack, contextmanager
from dataclasses import dataclass
import os
from pathlib import Path

from ...batch.source.annotation import annotate_with_batch_source_working_lines
from ...batch.state.lifecycle import create_batch
from ...batch.ownership.metadata_loading import acquire_ownership_for_metadata_dict
from ...batch.ownership.merging import merge_batch_ownership
from ...batch.ownership.remapping import remap_batch_ownership_with_lineage
from ...batch.ownership.translation import translate_lines_to_batch_ownership
from ...batch.ownership_update import acquire_batch_ownership_update_for_selection
from ...batch.merge.baseline_reference_translation import (
    translate_ownership_baseline_references,
)
from ...batch.state.query import read_batch_metadata
from ...batch.state.metadata_types import BatchFileMetadataDict
from ...batch.selection import (
    parse_command_line_selection,
    require_line_selection_in_view,
)
from ...batch.source.advancement import (
    advance_source_lines_preserving_existing_presence,
)
from ...batch.source.selected_line_refresh import (
    refresh_selected_lines_against_source_lines,
)
from ...batch.text_file_storage import add_file_to_batch
from ...batch.state.batch_names import batch_exists
from ...core.buffer import LineBuffer, buffer_ends_with_lf
from ...core.models import LineEntry, LineLevelChange
from ...core.replacement import ReplacementPayload, coerce_replacement_payload
from ...batch.ownership.model import BatchOwnership
from ...batch.source.cache import (
    load_session_batch_sources,
    save_session_batch_sources,
)
from ...batch.source.snapshots import create_batch_source_commit
from ...data.file_modes import detect_file_mode
from ...data.file_hunk_display import build_file_hunk_from_buffer
from ...data.line_state import load_line_changes_from_state
from ...utils.repository_buffers import (
    read_git_object_buffer_or_none,
    read_git_object_buffer_or_empty,
    load_working_tree_file_as_buffer,
)
from ...data.session import snapshot_file_if_untracked
from ...exceptions import exit_with_error
from ...i18n import _
from ...staging.content_buffers import (
    build_target_working_tree_buffer_from_lines,
    build_target_working_tree_buffer_with_replaced_lines,
)
from ...utils.git_repository import get_git_repository_root_path
from . import replacement_selection


@dataclass(frozen=True)
class DiscardLineReplacementSelection:
    """Prepared replacement selection for discard-to-batch."""

    line_changes: LineLevelChange
    file_path: str
    working_file_path: Path
    selected_lines: list[LineEntry]
    rewritten_line_changes: LineLevelChange
    rewritten_selected_lines: list[LineEntry]
    rewritten_working_lines: LineBuffer


@contextmanager
def prepare_discard_line_replacement_selection(
    line_id_specification: str,
    replacement_text: str | ReplacementPayload,
    *,
    no_edge_overlap: bool = False,
) -> Iterator[DiscardLineReplacementSelection]:
    """Prepare rewritten line selection state for discard-to-batch."""
    line_changes = load_line_changes_from_state()
    if line_changes is None:
        exit_with_error(_("No selected hunk. Run 'start' first."))
    requested_ids = set(parse_command_line_selection(line_id_specification))
    require_line_selection_in_view(
        line_changes,
        requested_ids,
        line_id_specification=line_id_specification,
    )
    effective_ids = replacement_selection.expand_replacement_selection_ids(
        line_changes,
        requested_ids,
    )

    selected_lines = [line for line in line_changes.lines if line.id in effective_ids]
    if not selected_lines:
        exit_with_error(
            _("No matching lines found for selection: {ids}").format(
                ids=line_id_specification
            )
        )

    working_file_path = get_git_repository_root_path() / line_changes.path
    if not os.path.lexists(working_file_path):
        exit_with_error(
            _("File not found in working tree: {file}").format(file=line_changes.path)
        )

    replacement_payload = coerce_replacement_payload(replacement_text)
    try:
        with load_working_tree_file_as_buffer(line_changes.path) as working_lines:
            rewritten_working_buffer = (
                build_target_working_tree_buffer_with_replaced_lines(
                    line_changes,
                    effective_ids,
                    replacement_payload,
                    working_lines,
                    working_has_trailing_newline=buffer_ends_with_lf(working_lines),
                    trim_unchanged_edge_anchors=not no_edge_overlap,
                )
            )
    except ValueError as error:
        exit_with_error(str(error))

    with rewritten_working_buffer as rewritten_working_lines:
        rewritten_cached_lines = build_file_hunk_from_buffer(
            line_changes.path,
            rewritten_working_lines,
        )
        if rewritten_cached_lines is None:
            exit_with_error(
                _("No changes in file '{file}'.").format(file=line_changes.path)
            )
        rewritten_line_changes = annotate_with_batch_source_working_lines(
            line_changes.path,
            rewritten_cached_lines,
            rewritten_working_lines,
        )
        rewritten_selected_lines = _select_rewritten_replacement_lines(
            selected_lines,
            rewritten_line_changes,
        )
        yield DiscardLineReplacementSelection(
            line_changes=line_changes,
            file_path=line_changes.path,
            working_file_path=working_file_path,
            selected_lines=selected_lines,
            rewritten_line_changes=rewritten_line_changes,
            rewritten_selected_lines=rewritten_selected_lines,
            rewritten_working_lines=rewritten_working_lines,
        )


def build_discard_line_replacement_target_buffer(
    selection: DiscardLineReplacementSelection,
) -> LineBuffer:
    """Return the worktree buffer after removing rewritten replacement lines."""
    rewritten_selected_ids = {
        line.id for line in selection.rewritten_selected_lines if line.id is not None
    }
    return build_target_working_tree_buffer_from_lines(
        selection.rewritten_line_changes,
        rewritten_selected_ids,
        selection.rewritten_working_lines,
    )


def add_discard_line_replacement_to_batch(
    batch_name: str,
    selection: DiscardLineReplacementSelection,
) -> None:
    """Persist a rewritten discard replacement selection to a batch."""
    if not batch_exists(batch_name):
        create_batch(batch_name, "Auto-created")

    metadata = read_batch_metadata(batch_name)
    file_metadata = metadata.get("files", {}).get(selection.file_path)

    with ExitStack() as ownership_stack:
        batch_source_commit: str | None
        try:
            if file_metadata is None:
                batch_source_commit = create_batch_source_commit(
                    selection.file_path,
                    file_buffer_override=selection.rewritten_working_lines,
                )
                _record_session_batch_source(
                    selection.file_path,
                    batch_source_commit,
                )
                prepared_selected_lines = (
                    refresh_selected_lines_against_source_lines(
                        selection.rewritten_selected_lines,
                        source_lines=selection.rewritten_working_lines,
                        working_lines=selection.rewritten_working_lines,
                        coordinate_lines=selection.rewritten_line_changes.lines,
                    )
                )
                reference_source_lines = ownership_stack.enter_context(
                    read_git_object_buffer_or_empty(
                        f"HEAD:{selection.file_path}"
                    )
                )
                update = ownership_stack.enter_context(
                    acquire_batch_ownership_update_for_selection(
                        batch_name=batch_name,
                        file_path=selection.file_path,
                        file_metadata=None,
                        selected_lines=prepared_selected_lines,
                        initial_batch_source_commit=batch_source_commit,
                        reference_source_lines=reference_source_lines,
                        batch_baseline_commit=metadata.get("baseline"),
                    )
                )
                ownership = update.ownership_after
                batch_source_commit = update.batch_source_commit
            else:
                ownership, batch_source_commit = _merge_replacement_with_batch(
                    selection,
                    file_metadata=file_metadata,
                    batch_baseline_commit=metadata.get("baseline"),
                    ownership_stack=ownership_stack,
                )
        except ValueError as e:
            exit_with_error(
                _(
                    "Cannot discard lines to batch: batch source is stale and remapping failed.\n"
                    "File: {file}\n"
                    "Batch: {batch}\n"
                    "Error: {error}"
                ).format(file=selection.file_path, batch=batch_name, error=str(e))
            )

        snapshot_file_if_untracked(selection.file_path)
        add_file_to_batch(
            batch_name,
            selection.file_path,
            ownership,
            detect_file_mode(selection.file_path),
            batch_source_commit=batch_source_commit,
        )


def _merge_replacement_with_batch(
    selection: DiscardLineReplacementSelection,
    *,
    file_metadata: BatchFileMetadataDict,
    batch_baseline_commit: str | None,
    ownership_stack: ExitStack,
) -> tuple[BatchOwnership, str]:
    if not isinstance(batch_baseline_commit, str) or not batch_baseline_commit:
        raise ValueError("replacement update requires a batch baseline commit")

    current_batch_source = file_metadata.get("batch_source_commit")
    existing_ownership = ownership_stack.enter_context(
        acquire_ownership_for_metadata_dict(file_metadata)
    )
    old_source_buffer = read_git_object_buffer_or_none(
        f"{current_batch_source}:{selection.file_path}"
    )
    if old_source_buffer is None:
        exit_with_error(
            _(
                "Cannot discard lines to batch: failed to read batch source for '{file}'."
            ).format(file=selection.file_path)
        )

    with (
        old_source_buffer as old_source_lines,
        read_git_object_buffer_or_empty(
            f"HEAD:{selection.file_path}"
        ) as reference_source_lines,
        read_git_object_buffer_or_empty(
            f"{batch_baseline_commit}:{selection.file_path}"
        ) as reference_target_lines,
        advance_source_lines_preserving_existing_presence(
            old_lines=old_source_lines,
            working_lines=selection.rewritten_working_lines,
            ownership=existing_ownership,
        ) as source_with_provenance,
    ):
        remapped_existing_ownership = remap_batch_ownership_with_lineage(
            ownership=existing_ownership,
            lineage=source_with_provenance.lineage,
        )
        refreshed_selected_lines = refresh_selected_lines_against_source_lines(
            selection.rewritten_selected_lines,
            source_lines=source_with_provenance.source_buffer,
            working_lines=(),
            lineage=source_with_provenance.lineage,
            coordinate_lines=selection.rewritten_line_changes.lines,
        )
        new_ownership = translate_lines_to_batch_ownership(refreshed_selected_lines)
        translate_ownership_baseline_references(
            new_ownership,
            reference_source_lines,
            reference_target_lines,
        )
        batch_source_commit = create_batch_source_commit(
            selection.file_path,
            file_buffer_override=source_with_provenance.source_buffer,
        )
        _record_session_batch_source(selection.file_path, batch_source_commit)
        return (
            merge_batch_ownership(remapped_existing_ownership, new_ownership),
            batch_source_commit,
        )


def _record_session_batch_source(file_path: str, batch_source_commit: str) -> None:
    batch_sources = load_session_batch_sources()
    batch_sources[file_path] = batch_source_commit
    save_session_batch_sources(batch_sources)


def _select_rewritten_replacement_lines(
    original_selected_lines: list[LineEntry],
    rewritten_line_changes: LineLevelChange,
) -> list[LineEntry]:
    """Find the rewritten changed span that overlaps the original selection."""
    original_old_lines = {
        line.old_line_number
        for line in original_selected_lines
        if line.old_line_number is not None
    }
    original_new_lines = {
        line.new_line_number
        for line in original_selected_lines
        if line.new_line_number is not None
    }

    matching_indices = [
        index
        for index, line in enumerate(rewritten_line_changes.lines)
        if line.kind != " "
        and (
            line.old_line_number in original_old_lines
            or line.new_line_number in original_new_lines
        )
    ]
    if matching_indices:
        start_index = min(matching_indices)
        end_index = max(matching_indices)
        return [
            line
            for line in rewritten_line_changes.lines[start_index : end_index + 1]
            if line.kind != " "
        ]

    exit_with_error(
        _("Replacement selection could not be located after rewriting the file.")
    )
