"""Tests for checkout data-integrity guarantees.

Coverage tiers
--------------
Unit         — read_checkout_head / checkout_head_path helpers.
Integration  — pre-flight aborts cleanly (zero mutations), CHECKOUT_HEAD
               marker lifecycle, status detection of interrupted checkout.
End-to-end   — full CLI round-trips: missing object, simulated mid-flight
               kill (marker left behind), recovery via retry checkout.
Stress       — rapid branch switching never leaves a stale marker.

Key invariants under test
-------------------------
1. Pre-flight: if any restore object is absent from the store, checkout
   prints an error and does NOT modify any file on disk.
2. CHECKOUT_HEAD written before first mutation; removed on success.
3. ``muse status`` (text + JSON) detects a stale marker and warns loudly.
4. A successful checkout on retry clears the stale marker.
5. Interrupted checkout leaves HEAD pointing to the *old* branch.
"""

from __future__ import annotations

import json
import os
import pathlib
import shutil

import pytest

from tests.cli_test_helper import CliRunner, InvokeResult
from muse.core.types import MsgpackDict
from muse.core.refs import (
    get_head_commit_id,
    read_current_branch,
)
from muse.core.object_store import has_object
from muse.cli.commands.checkout import checkout_head_path, read_checkout_head

runner = CliRunner()

# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────


def _invoke(repo: pathlib.Path, args: list[str]) -> InvokeResult:
    saved = os.getcwd()
    try:
        os.chdir(repo)
        return runner.invoke(None, args)
    finally:
        os.chdir(saved)


def _commit(repo: pathlib.Path, msg: str = "commit") -> InvokeResult:
    _invoke(repo, ["code", "add", "."])
    return _invoke(repo, ["commit", "-m", msg])


def _checkout(repo: pathlib.Path, branch: str, *flags: str) -> InvokeResult:
    return _invoke(repo, ["checkout", *flags, branch])


def _status_json(repo: pathlib.Path) -> MsgpackDict:
    r = _invoke(repo, ["status", "--json"])
    return json.loads(r.output)


def _status_text(repo: pathlib.Path) -> tuple[str, str]:
    """Return (stdout, stderr) of muse status."""
    r = _invoke(repo, ["status"])
    return r.output, (r.stderr or "")


# ──────────────────────────────────────────────────────────────────────────────
# Fixtures
# ──────────────────────────────────────────────────────────────────────────────


@pytest.fixture()
def two_branch_repo(tmp_path: pathlib.Path) -> pathlib.Path:
    """Repo with main and feat branches, each with distinct files.

    main: a.py
    feat: a.py (unchanged) + b.py (feat-only)
    """
    saved = os.getcwd()
    try:
        os.chdir(tmp_path)
        runner.invoke(None, ["init"])
    finally:
        os.chdir(saved)
    (tmp_path / "a.py").write_text("x = 1\n")
    _commit(tmp_path, "main initial")
    _invoke(tmp_path, ["checkout", "-b", "feat"])
    (tmp_path / "b.py").write_text("y = 2\n")
    _commit(tmp_path, "feat adds b.py")
    _checkout(tmp_path, "main")
    return tmp_path


# ──────────────────────────────────────────────────────────────────────────────
# Unit: helper functions
# ──────────────────────────────────────────────────────────────────────────────


class TestCheckoutHeadHelpers:
    def test_checkout_head_path_points_into_muse_dir(self, two_branch_repo: pathlib.Path) -> None:
        p = checkout_head_path(two_branch_repo)
        assert p.parent == two_branch_repo / ".muse"
        assert p.name == "CHECKOUT_HEAD"

    def test_read_checkout_head_returns_none_when_absent(self, two_branch_repo: pathlib.Path) -> None:
        assert read_checkout_head(two_branch_repo) is None

    def test_read_checkout_head_returns_content_when_present(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n", encoding="utf-8")
        assert read_checkout_head(two_branch_repo) == "feat"
        marker.unlink()

    def test_read_checkout_head_strips_trailing_newline(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n\n", encoding="utf-8")
        assert read_checkout_head(two_branch_repo) == "feat"
        marker.unlink()

    def test_read_checkout_head_empty_file_returns_none(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("", encoding="utf-8")
        assert read_checkout_head(two_branch_repo) is None
        marker.unlink()


# ──────────────────────────────────────────────────────────────────────────────
# Integration: pre-flight object existence check
# ──────────────────────────────────────────────────────────────────────────────


class TestPreflightObjectCheck:
    """Pre-flight: abort before any mutation if a restore object is missing."""

    def _sabotage_object(self, repo: pathlib.Path, rel_path: str) -> pathlib.Path:
        """Remove the object backing *rel_path* from the store; return its path."""
        from muse.core.refs import get_head_commit_id
        from muse.core.commits import read_commit
        from muse.core.snapshots import read_snapshot
        from muse.core.object_store import _object_path_with_fallback

        # Switch to feat first so its snapshot is current, then read its manifest
        _checkout(repo, "feat")
        branch = read_current_branch(repo)
        commit_id = get_head_commit_id(repo, branch)
        assert commit_id
        commit = read_commit(repo, commit_id)
        assert commit
        snap = read_snapshot(repo, commit.snapshot_id)
        assert snap
        obj_id = snap.manifest[rel_path]
        obj_path = _object_path_with_fallback(repo, obj_id)
        assert obj_path.exists(), f"Object for {rel_path} not in store"
        # Back on main before sabotaging
        _checkout(repo, "main")
        # Now remove the object
        obj_path.unlink()
        return obj_path

    def test_missing_object_exits_nonzero(self, two_branch_repo: pathlib.Path) -> None:
        self._sabotage_object(two_branch_repo, "b.py")
        result = _checkout(two_branch_repo, "feat")
        assert result.exit_code != 0

    def test_missing_object_error_on_stderr(self, two_branch_repo: pathlib.Path) -> None:
        self._sabotage_object(two_branch_repo, "b.py")
        result = _checkout(two_branch_repo, "feat")
        err = result.stderr or result.output
        assert "missing" in err.lower() or "object" in err.lower()

    def test_missing_object_working_tree_unchanged(self, two_branch_repo: pathlib.Path) -> None:
        """Zero mutations — b.py must not appear on main after a failed checkout."""
        self._sabotage_object(two_branch_repo, "b.py")
        _checkout(two_branch_repo, "feat")
        # Working tree must not contain b.py — no partial mutations
        assert not (two_branch_repo / "b.py").exists()

    def test_missing_object_a_py_unchanged(self, two_branch_repo: pathlib.Path) -> None:
        """Files that would be kept unchanged must also remain untouched."""
        self._sabotage_object(two_branch_repo, "b.py")
        _checkout(two_branch_repo, "feat")
        # a.py was already on main and unchanged; must still be present
        assert (two_branch_repo / "a.py").exists()

    def test_missing_object_head_stays_on_old_branch(self, two_branch_repo: pathlib.Path) -> None:
        """HEAD must not advance when pre-flight aborts."""
        self._sabotage_object(two_branch_repo, "b.py")
        _checkout(two_branch_repo, "feat")
        assert read_current_branch(two_branch_repo) == "main"

    def test_missing_object_no_checkout_head_marker(self, two_branch_repo: pathlib.Path) -> None:
        """Pre-flight abort fires before any marker is written."""
        self._sabotage_object(two_branch_repo, "b.py")
        _checkout(two_branch_repo, "feat")
        # Marker must NOT be present — pre-flight aborted before any mutation
        assert read_checkout_head(two_branch_repo) is None

    def test_stderr_names_the_missing_file(self, two_branch_repo: pathlib.Path) -> None:
        self._sabotage_object(two_branch_repo, "b.py")
        result = _checkout(two_branch_repo, "feat")
        err = result.stderr or result.output
        assert "b.py" in err

    def test_stderr_says_working_tree_not_modified(self, two_branch_repo: pathlib.Path) -> None:
        self._sabotage_object(two_branch_repo, "b.py")
        result = _checkout(two_branch_repo, "feat")
        err = result.stderr or result.output
        assert "NOT modified" in err or "not modified" in err.lower()


# ──────────────────────────────────────────────────────────────────────────────
# Integration: CHECKOUT_HEAD marker lifecycle
# ──────────────────────────────────────────────────────────────────────────────


class TestCheckoutHeadMarker:
    """CHECKOUT_HEAD written before mutations, cleared after success."""

    def test_marker_absent_after_clean_checkout(self, two_branch_repo: pathlib.Path) -> None:
        _checkout(two_branch_repo, "feat")
        assert read_checkout_head(two_branch_repo) is None

    def test_marker_absent_after_round_trip(self, two_branch_repo: pathlib.Path) -> None:
        _checkout(two_branch_repo, "feat")
        _checkout(two_branch_repo, "main")
        assert read_checkout_head(two_branch_repo) is None

    def test_stale_marker_survives_failed_checkout(self, two_branch_repo: pathlib.Path) -> None:
        """Manually plant a stale marker; it must persist (not auto-cleared)."""
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n", encoding="utf-8")
        # Successful checkout to main should clear it
        _checkout(two_branch_repo, "main")
        # main is already current, but checkout still fires the snapshot path
        # which should clear the marker on success
        # (Already on main — 'already_on' path does NOT call _checkout_snapshot,
        # so the marker is unaffected.  Plant while NOT on main.)
        # Reset: switch to feat first, plant marker, switch back
        _checkout(two_branch_repo, "feat")
        marker.write_text("main\n", encoding="utf-8")
        _checkout(two_branch_repo, "main")
        assert read_checkout_head(two_branch_repo) is None

    def test_simulated_interrupted_marker_persists(self, two_branch_repo: pathlib.Path) -> None:
        """Simulate kill mid-checkout: plant marker, verify it stays."""
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n", encoding="utf-8")
        # Do not run checkout — marker lingers
        assert read_checkout_head(two_branch_repo) == "feat"
        marker.unlink()

    def test_retry_checkout_clears_stale_marker(self, two_branch_repo: pathlib.Path) -> None:
        """Retry of the interrupted checkout must clear the marker."""
        # Simulate interrupted checkout of feat (marker left behind,
        # b.py not restored)
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n", encoding="utf-8")
        # Retry — should succeed and clear marker
        result = _checkout(two_branch_repo, "feat")
        assert result.exit_code == 0
        assert read_checkout_head(two_branch_repo) is None

    def test_marker_records_target_branch_name(self, two_branch_repo: pathlib.Path) -> None:
        """After a successful checkout the marker is gone; its content was the target."""
        # We can't easily inspect the marker mid-flight without mocking,
        # but we can write it ourselves and verify read_checkout_head returns it.
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")
        assert read_checkout_head(two_branch_repo) == "feat"
        marker.unlink()


# ──────────────────────────────────────────────────────────────────────────────
# Integration: muse status detects interrupted checkout
# ──────────────────────────────────────────────────────────────────────────────


class TestStatusDetectsInterruptedCheckout:
    """muse status warns loudly when CHECKOUT_HEAD exists."""

    def test_json_checkout_interrupted_false_normally(self, two_branch_repo: pathlib.Path) -> None:
        data = _status_json(two_branch_repo)
        assert data["checkout_interrupted"] is False

    def test_json_checkout_target_null_normally(self, two_branch_repo: pathlib.Path) -> None:
        data = _status_json(two_branch_repo)
        assert data["checkout_target"] is None

    def test_json_checkout_interrupted_true_with_marker(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")
        try:
            data = _status_json(two_branch_repo)
            assert data["checkout_interrupted"] is True
        finally:
            marker.unlink(missing_ok=True)

    def test_json_checkout_target_set_with_marker(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")
        try:
            data = _status_json(two_branch_repo)
            assert data["checkout_target"] == "feat"
        finally:
            marker.unlink(missing_ok=True)

    def test_json_keys_always_present(self, two_branch_repo: pathlib.Path) -> None:
        data = _status_json(two_branch_repo)
        assert "checkout_interrupted" in data
        assert "checkout_target" in data

    def test_text_status_warns_on_interrupted_checkout(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")
        try:
            _, stderr = _status_text(two_branch_repo)
            assert "CHECKOUT INTERRUPTED" in stderr or "CHECKOUT INTERRUPTED" in stderr.upper()
        finally:
            marker.unlink(missing_ok=True)

    def test_text_status_names_target_branch(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")
        try:
            _, stderr = _status_text(two_branch_repo)
            assert "feat" in stderr
        finally:
            marker.unlink(missing_ok=True)

    def test_text_status_mentions_retry_command(self, two_branch_repo: pathlib.Path) -> None:
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")
        try:
            _, stderr = _status_text(two_branch_repo)
            assert "muse checkout" in stderr
        finally:
            marker.unlink(missing_ok=True)

    def test_text_status_clean_no_warning(self, two_branch_repo: pathlib.Path) -> None:
        _, stderr = _status_text(two_branch_repo)
        assert "CHECKOUT INTERRUPTED" not in stderr


# ──────────────────────────────────────────────────────────────────────────────
# End-to-end: full recovery flow
# ──────────────────────────────────────────────────────────────────────────────


class TestCheckoutIntegrityE2E:
    """Full end-to-end scenarios for the checkout integrity system."""

    def test_successful_checkout_leaves_clean_status(self, two_branch_repo: pathlib.Path) -> None:
        _checkout(two_branch_repo, "feat")
        data = _status_json(two_branch_repo)
        assert data["checkout_interrupted"] is False
        assert data["checkout_target"] is None
        assert data["clean"] is True

    def test_interrupted_checkout_then_status_then_retry(self, two_branch_repo: pathlib.Path) -> None:
        """Full recovery flow: interrupt → status warns → retry → clean."""
        # Simulate interruption: plant marker and delete b.py as if
        # the checkout partially ran (deleted files but didn't restore yet)
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")

        # status should warn
        data = _status_json(two_branch_repo)
        assert data["checkout_interrupted"] is True

        # retry the checkout — should succeed and clean up
        result = _checkout(two_branch_repo, "feat")
        assert result.exit_code == 0

        # marker gone, status clean
        assert read_checkout_head(two_branch_repo) is None
        data2 = _status_json(two_branch_repo)
        assert data2["checkout_interrupted"] is False
        assert data2["checkout_target"] is None

    def test_interrupted_checkout_head_unchanged(self, two_branch_repo: pathlib.Path) -> None:
        """HEAD must still point to the old branch after a simulated interruption."""
        # On main, plant a marker pretending we were switching to feat
        marker = checkout_head_path(two_branch_repo)
        marker.write_text("feat\n")
        # HEAD has not changed — we only wrote the marker, didn't run checkout
        assert read_current_branch(two_branch_repo) == "main"
        marker.unlink()

    def test_b_py_present_after_recovery_checkout(self, two_branch_repo: pathlib.Path) -> None:
        """After successful retry checkout to feat, feat-only file must exist."""
        # First do a clean checkout to feat to confirm it works
        result = _checkout(two_branch_repo, "feat")
        assert result.exit_code == 0
        assert (two_branch_repo / "b.py").exists()

    def test_b_py_absent_after_checkout_back_to_main(self, two_branch_repo: pathlib.Path) -> None:
        """After checking back out to main, feat-only file must be gone."""
        _checkout(two_branch_repo, "feat")
        result = _checkout(two_branch_repo, "main")
        assert result.exit_code == 0
        assert not (two_branch_repo / "b.py").exists()

    def test_preflight_then_retry_with_restored_object(
        self, two_branch_repo: pathlib.Path
    ) -> None:
        """Remove an object, verify clean abort, restore object, verify checkout succeeds."""
        from muse.core.refs import get_head_commit_id
        from muse.core.commits import read_commit
        from muse.core.snapshots import read_snapshot
        from muse.core.object_store import _object_path_with_fallback

        # Read the object ID for b.py from the feat snapshot
        _checkout(two_branch_repo, "feat")
        commit_id = get_head_commit_id(two_branch_repo, "feat")
        commit = read_commit(two_branch_repo, commit_id)
        snap = read_snapshot(two_branch_repo, commit.snapshot_id)
        obj_id = snap.manifest["b.py"]
        obj_path = _object_path_with_fallback(two_branch_repo, obj_id)
        # Save contents before sabotage
        saved = obj_path.read_bytes()

        _checkout(two_branch_repo, "main")
        obj_path.unlink()

        # Pre-flight fails cleanly
        result = _checkout(two_branch_repo, "feat")
        assert result.exit_code != 0
        assert not (two_branch_repo / "b.py").exists()
        assert read_current_branch(two_branch_repo) == "main"

        # Restore the object (simulating a fetch)
        obj_path.write_bytes(saved)

        # Retry succeeds
        result2 = _checkout(two_branch_repo, "feat")
        assert result2.exit_code == 0
        assert (two_branch_repo / "b.py").exists()
        assert read_current_branch(two_branch_repo) == "feat"


# ──────────────────────────────────────────────────────────────────────────────
# Stress: rapid switching never leaves a stale marker
# ──────────────────────────────────────────────────────────────────────────────


class TestCheckoutIntegrityStress:
    def test_rapid_switching_no_stale_marker(self, two_branch_repo: pathlib.Path) -> None:
        """Switching branches 50 times must never leave CHECKOUT_HEAD behind."""
        for i in range(25):
            r1 = _checkout(two_branch_repo, "feat")
            assert r1.exit_code == 0, f"iter {i}: checkout feat failed"
            assert read_checkout_head(two_branch_repo) is None, f"iter {i}: marker after feat checkout"
            r2 = _checkout(two_branch_repo, "main")
            assert r2.exit_code == 0, f"iter {i}: checkout main failed"
            assert read_checkout_head(two_branch_repo) is None, f"iter {i}: marker after main checkout"

    def test_status_json_schema_stable_across_branches(self, two_branch_repo: pathlib.Path) -> None:
        """checkout_interrupted and checkout_target always present in JSON output."""
        required = {"checkout_interrupted", "checkout_target"}
        for branch in ("feat", "main", "feat", "main"):
            _checkout(two_branch_repo, branch)
            data = _status_json(two_branch_repo)
            missing = required - data.keys()
            assert not missing, f"Missing keys after checkout to {branch}: {missing}"


# ──────────────────────────────────────────────────────────────────────────────
# Regression: checkout must update HEAD to the target branch
# ──────────────────────────────────────────────────────────────────────────────


class TestCheckoutHeadUpdate:
    """Regression tests for the _ref_path-not-imported bug.

    Previously, checkout exited with code 1 (NameError on _ref_path) and left
    HEAD pointing at the old branch.  Every subsequent commit then landed on
    the wrong branch.  These tests pin the invariant that a successful checkout
    always updates HEAD.
    """

    def test_checkout_updates_head_to_target_branch(
        self, two_branch_repo: pathlib.Path
    ) -> None:
        """Switching to feat must update HEAD; switching back must restore main."""
        assert read_current_branch(two_branch_repo) == "main"

        result = _checkout(two_branch_repo, "feat")
        assert result.exit_code == 0, f"checkout feat failed: {result.output}"
        assert read_current_branch(two_branch_repo) == "feat", (
            "HEAD must point to feat after successful checkout"
        )

        result = _checkout(two_branch_repo, "main")
        assert result.exit_code == 0, f"checkout main failed: {result.output}"
        assert read_current_branch(two_branch_repo) == "main", (
            "HEAD must point to main after switching back"
        )

    def test_commit_after_checkout_lands_on_correct_branch(
        self, two_branch_repo: pathlib.Path
    ) -> None:
        """A commit made after checkout must advance the target branch tip, not main."""
        from muse.core.refs import get_head_commit_id

        tip_main_before = get_head_commit_id(two_branch_repo, "main")

        result = _checkout(two_branch_repo, "feat")
        assert result.exit_code == 0

        # Write a new file and commit on feat
        (two_branch_repo / "on_feat.py").write_text("x = 1\n")
        _invoke(two_branch_repo, ["code", "add", "on_feat.py"])
        _invoke(two_branch_repo, ["commit", "-m", "test: commit on feat"])

        # main tip must be unchanged
        tip_main_after = get_head_commit_id(two_branch_repo, "main")
        assert tip_main_before == tip_main_after, (
            "Commit after checkout feat must not advance main"
        )

        # feat tip must have advanced
        tip_feat = get_head_commit_id(two_branch_repo, "feat")
        assert tip_feat != tip_main_after, (
            "Commit after checkout feat must advance feat, not main"
        )