"""Multi-file discard-to-batch command support."""

from __future__ import annotations

from collections.abc import Iterator, Sequence
from contextlib import ExitStack
from dataclasses import dataclass
import os

from ...batch.source.annotation import annotate_with_batch_source
from ...batch.file_state import BatchMetadataRevision, SourceBoundOwnership
from ...batch.ownership import insertion_references as _insertion_references
from ...batch.state.metadata_types import BatchMetadataDict
from ...batch.state.lifecycle import create_batch
from ...batch.ownership_update import acquire_batch_ownership_update_for_selection
from ...batch.state.query import read_batch_metadata
from ...batch.text_file_storage import BatchFileUpdate, add_files_to_batch
from ...batch.state.batch_names import batch_exists
from ...core.buffer import LineBuffer
from ...core.diff_parser import (
    acquire_unified_diff,
    build_line_changes_from_patch_lines,
    patch_requires_unidiff_zero,
)
from ...core.hashing import compute_stable_hunk_hash_from_lines
from ...core.models import (
    BinaryFileChange,
    FileModeChange,
    GitlinkChange,
    LineEntry,
    RenameChange,
    TextFileDeletionChange,
)
from ...utils.context_lines import get_context_lines
from ...data.file_modes import detect_file_mode_from_root
from ...data.file_tracking import auto_add_untracked_files
from ...data.live_diff import stream_live_git_diff
from ...data.progress import record_hunks_discarded
from ...data.session import snapshot_files_if_untracked
from ...exceptions import exit_with_error
from ...git_paths import display_path
from ...i18n import _
from ...utils.file_io import read_text_file_line_set
from ...utils.git_worktree import (
    git_apply_to_worktree,
)
from ...utils.git_repository import (
    get_git_repository_root_path,
    require_git_repository,
)
from ...utils.journal import log_journal
from ...utils.paths import (
    ensure_state_directory_exists,
    get_block_list_file_path,
)
from ...utils.repository_buffers import read_git_object_buffer_or_empty
from ...utils.session_start_point import session_comparison_base
from ..selection.action_completion import finish_selected_change_action
from ..index_cleanup import remove_path_from_index
from .discard_file_to_batch import discard_file_to_batch


@dataclass(frozen=True)
class _PreparedPatchDiscard:
    """One patch scheduled for reverse application."""

    patch_lines: Sequence[bytes]
    patch_hash: str


@dataclass
class _TextFileDiscardInput:
    """Collected text hunks for one file-scope discard."""

    file_path: str
    file_mode: str
    comparison_base: str
    all_lines_to_batch: list[LineEntry]
    patches_to_discard: list[_PreparedPatchDiscard]


@dataclass(frozen=True)
class _CollectedTextFileDiscards:
    """Collected text-file discard inputs from one live diff."""

    inputs_by_file: dict[str, _TextFileDiscardInput]
    files_with_text_patches: set[str]


@dataclass(frozen=True)
class _PreparedTextFileDiscardToBatch:
    """Prepared text-file discard that can be published atomically."""

    file_path: str
    file_mode: str
    bound_ownership: SourceBoundOwnership
    batch_source_commit: str
    expected_metadata_revision: BatchMetadataRevision
    patches_to_discard: list[_PreparedPatchDiscard]


@dataclass(frozen=True)
class DiscardFilesToBatchResult:
    """Aggregate result for multi-file discard-to-batch actions."""

    discarded_hunks: int
    discarded_files: list[str]


class _DiscardFilesToBatchSession:
    """Mutable publication state for one multi-file discard-to-batch action."""

    def __init__(self, batch_name: str, metadata: BatchMetadataDict) -> None:
        self._batch_name = batch_name
        self._metadata = metadata
        self._ownership_stack = ExitStack()
        self._prepared_discards: list[_PreparedTextFileDiscardToBatch] = []
        self._discarded_hunks = 0
        self._discarded_files: list[str] = []

    @property
    def discarded_hunks(self) -> int:
        """Return the total discarded hunk count."""
        return self._discarded_hunks

    @property
    def discarded_files(self) -> list[str]:
        """Return files that produced discarded hunks."""
        return self._discarded_files

    def close(self) -> None:
        """Release any pending batch ownership update contexts."""
        self._ownership_stack.close()

    def prepare_text_file(self, discard_input: _TextFileDiscardInput) -> bool:
        """Stage one text file for a later batch publication flush."""
        prepared = _prepare_text_file_discard_to_batch(
            self._batch_name,
            discard_input,
            metadata=self._metadata,
            ownership_stack=self._ownership_stack,
        )
        if prepared is None:
            return False
        self._prepared_discards.append(prepared)
        return True

    def flush(self) -> None:
        """Publish any pending prepared text file discards."""
        ownership_stack = self._ownership_stack
        prepared_discards = self._prepared_discards
        self._ownership_stack = ExitStack()
        self._prepared_discards = []

        with ownership_stack:
            result = _discard_prepared_text_files_to_batch(
                self._batch_name,
                prepared_discards,
            )
        self._record_result(result)

    def record_single_file_discard(self, file_path: str, discarded_hunks: int) -> bool:
        """Record fallback single-file command output."""
        if discarded_hunks <= 0:
            return False
        self._record_result(
            DiscardFilesToBatchResult(
                discarded_hunks=discarded_hunks,
                discarded_files=[file_path],
            )
        )
        return True

    def _record_result(self, result: DiscardFilesToBatchResult) -> None:
        if result.discarded_hunks <= 0:
            return
        self._discarded_hunks += result.discarded_hunks
        self._discarded_files.extend(result.discarded_files)
        self._metadata = read_batch_metadata(self._batch_name)


def _prepare_text_file_discard_to_batch(
    batch_name: str,
    discard_input: _TextFileDiscardInput,
    *,
    metadata: BatchMetadataDict,
    ownership_stack: ExitStack,
) -> _PreparedTextFileDiscardToBatch | None:
    """Prepare one normal text file discard without publishing batch state."""
    if not discard_input.all_lines_to_batch:
        return None

    file_path = discard_input.file_path
    file_metadata = metadata.get("files", {}).get(file_path)
    metadata_revision = BatchMetadataRevision.from_metadata(metadata)

    try:
        with read_git_object_buffer_or_empty(
            f"{discard_input.comparison_base}:{file_path}"
        ) as reference_source_lines:
            update = ownership_stack.enter_context(
                acquire_batch_ownership_update_for_selection(
                    batch_name=batch_name,
                    file_path=file_path,
                    file_metadata=file_metadata,
                    metadata_revision=metadata_revision,
                    selected_lines=discard_input.all_lines_to_batch,
                    reference_source_lines=reference_source_lines,
                    batch_baseline_commit=metadata.get("baseline"),
                )
            )
    except ValueError as e:
        exit_with_error(
            _(
                "Cannot discard file to batch: batch source is stale and remapping failed.\n"
                "File: {file}\n"
                "Batch: {batch}\n"
                "Error: {error}"
            ).format(
                file=display_path(file_path),
                batch=batch_name,
                error=str(e),
            )
        )

    return _PreparedTextFileDiscardToBatch(
        file_path=file_path,
        file_mode=discard_input.file_mode,
        bound_ownership=update.bound_ownership,
        batch_source_commit=update.batch_source_commit,
        expected_metadata_revision=update.expected_metadata_revision,
        patches_to_discard=discard_input.patches_to_discard,
    )


def _collect_text_file_discard_inputs(
    files: list[str],
    *,
    blocked_hashes: set[str],
    patch_stack: ExitStack,
) -> _CollectedTextFileDiscards:
    """Collect normal text file discard inputs from one Git diff."""
    if not files:
        return _CollectedTextFileDiscards(
            inputs_by_file={},
            files_with_text_patches=set(),
        )

    repo_root = get_git_repository_root_path()
    inputs_by_file: dict[str, _TextFileDiscardInput] = {}
    files_with_text_patches: set[str] = set()
    comparison_base = session_comparison_base()

    with acquire_unified_diff(
        stream_live_git_diff(
            base=comparison_base,
            context_lines=get_context_lines(),
            paths=files,
        )
    ) as patches:
        for patch in patches:
            if isinstance(patch, FileModeChange):
                continue
            if isinstance(patch, RenameChange):
                continue

            if isinstance(patch, TextFileDeletionChange):
                continue

            if isinstance(patch, GitlinkChange):
                exit_with_error(
                    _("Discarding submodule pointer changes to a batch is not supported yet.")
                )

            if isinstance(patch, BinaryFileChange):
                continue

            file_path = patch.path()
            files_with_text_patches.add(file_path)

            patch_hash = compute_stable_hunk_hash_from_lines(patch.lines)
            if patch_hash in blocked_hashes:
                continue

            hunk_lines = build_line_changes_from_patch_lines(
                patch.lines,
                annotator=annotate_with_batch_source,
            )
            _insertion_references.record_baseline_references_for_additions(
                hunk_lines,
            )
            discard_input = inputs_by_file.get(file_path)
            if discard_input is None:
                discard_input = _TextFileDiscardInput(
                    file_path=file_path,
                    file_mode=detect_file_mode_from_root(repo_root, file_path),
                    comparison_base=comparison_base,
                    all_lines_to_batch=[],
                    patches_to_discard=[],
                )
                inputs_by_file[file_path] = discard_input
            discard_input.all_lines_to_batch.extend(hunk_lines.lines)
            discard_input.patches_to_discard.append(
                _PreparedPatchDiscard(
                    patch_lines=patch_stack.enter_context(
                        LineBuffer.from_chunks(patch.lines)
                    ),
                    patch_hash=patch_hash,
                )
            )
            blocked_hashes.add(patch_hash)

    return _CollectedTextFileDiscards(
        inputs_by_file=inputs_by_file,
        files_with_text_patches=files_with_text_patches,
    )


def _run_reverse_apply_for_prepared_discards(
    prepared_discards: list[_PreparedTextFileDiscardToBatch],
    *,
    check_only: bool = False,
) -> None:
    apply_groups: dict[bool, list[_PreparedTextFileDiscardToBatch]] = {
        False: [],
        True: [],
    }
    for prepared in prepared_discards:
        requirements = {
            patch_requires_unidiff_zero(patch.patch_lines)
            for patch in prepared.patches_to_discard
        }
        if len(requirements) > 1:
            # One Git diff acquisition uses one context setting.  A normal
            # contextual diff can lack context only when its sole hunk spans
            # the complete file, so generated per-file patches are
            # homogeneous.  Refuse unexpected mixed input instead of
            # weakening Git's placement rules for its contextual hunks.
            exit_with_error(
                _("Failed to discard changes from file: {err}").format(
                    err=(
                        "mixed zero-context and contextual patches for "
                        f"{display_path(prepared.file_path)}"
                    )
                )
            )
        if requirements:
            apply_groups[requirements.pop()].append(prepared)

    for unidiff_zero, group in apply_groups.items():
        if not group:
            continue

        def patch_chunks() -> Iterator[bytes]:
            for prepared in group:
                for patch in prepared.patches_to_discard:
                    yield from patch.patch_lines

        apply_result = git_apply_to_worktree(
            patch_chunks(),
            reverse=True,
            unidiff_zero=unidiff_zero,
            check_only=check_only,
            check=False,
        )

        if apply_result.returncode != 0:
            exit_with_error(
                _("Failed to discard changes from file: {err}").format(
                    err=apply_result.stderr
                )
            )


def _discard_prepared_text_files_to_batch(
    batch_name: str,
    prepared_discards: list[_PreparedTextFileDiscardToBatch],
) -> DiscardFilesToBatchResult:
    """Publish prepared text file discards once, then update the worktree."""
    if not prepared_discards:
        return DiscardFilesToBatchResult(discarded_hunks=0, discarded_files=[])

    snapshot_files_if_untracked([prepared.file_path for prepared in prepared_discards])

    _run_reverse_apply_for_prepared_discards(prepared_discards, check_only=True)
    add_files_to_batch(
        batch_name,
        [
            BatchFileUpdate(
                file_path=prepared.file_path,
                batch_source_commit=prepared.batch_source_commit,
                bound_ownership=prepared.bound_ownership,
                expected_metadata_revision=prepared.expected_metadata_revision,
                file_mode=prepared.file_mode,
            )
            for prepared in prepared_discards
        ],
    )
    _run_reverse_apply_for_prepared_discards(prepared_discards)

    repo_root = get_git_repository_root_path()
    for prepared in prepared_discards:
        full_path = repo_root / prepared.file_path
        if not os.path.lexists(full_path):
            remove_path_from_index(prepared.file_path)

    hunk_hashes = [
        patch.patch_hash
        for prepared in prepared_discards
        for patch in prepared.patches_to_discard
    ]
    record_hunks_discarded(hunk_hashes)

    return DiscardFilesToBatchResult(
        discarded_hunks=len(hunk_hashes),
        discarded_files=[
            prepared.file_path
            for prepared in prepared_discards
            if prepared.patches_to_discard
        ],
    )


def discard_files_to_batch(
    batch_name: str,
    files: list[str],
    *,
    quiet: bool = False,
    advance: bool = True,
    auto_advance: bool | None = None,
) -> DiscardFilesToBatchResult:
    """Save resolved text files to a batch with one batch publication."""
    require_git_repository()
    ensure_state_directory_exists()

    if not files:
        return DiscardFilesToBatchResult(discarded_hunks=0, discarded_files=[])
    auto_add_untracked_files(files)
    if not batch_exists(batch_name):
        create_batch(batch_name, "Auto-created")

    blocklist_path = get_block_list_file_path()
    blocked_hashes = read_text_file_line_set(blocklist_path)
    session = _DiscardFilesToBatchSession(
        batch_name=batch_name,
        metadata=read_batch_metadata(batch_name),
    )
    patch_stack = ExitStack()

    try:
        collected_discards = _collect_text_file_discard_inputs(
            files,
            blocked_hashes=blocked_hashes,
            patch_stack=patch_stack,
        )

        for file_path in files:
            log_journal(
                "discard_files_to_batch_file_start",
                batch_name=batch_name,
                file_path=file_path,
                quiet=quiet,
            )
            outcome = "failed"
            try:
                discard_input = collected_discards.inputs_by_file.get(file_path)
                if (
                    discard_input is None
                    and file_path in collected_discards.files_with_text_patches
                ):
                    outcome = "skipped"
                    continue

                if (
                    discard_input is None
                    or not session.prepare_text_file(discard_input)
                ):
                    session.flush()
                    discarded_hunks = discard_file_to_batch(
                        batch_name,
                        file_path,
                        quiet=True,
                        advance=False,
                        auto_advance=auto_advance,
                    )
                    if session.record_single_file_discard(
                        file_path,
                        discarded_hunks,
                    ):
                        blocked_hashes = read_text_file_line_set(blocklist_path)
                    outcome = "fallback"
                    continue

                outcome = "prepared"
            finally:
                log_journal(
                    "discard_files_to_batch_file_end",
                    batch_name=batch_name,
                    file_path=file_path,
                    outcome=outcome,
                )

        session.flush()
    finally:
        session.close()
        patch_stack.close()

    if advance:
        finish_selected_change_action(quiet=quiet, auto_advance=auto_advance)

    return DiscardFilesToBatchResult(
        discarded_hunks=session.discarded_hunks,
        discarded_files=session.discarded_files,
    )
