"""Comprehensive tests for ``muse merge``.

Covers:
- E2E: merge fast-forward, merge with conflicts, --format json
- Integration: HEAD updated after merge, conflict state written
- Stress: merge with many files
"""

from __future__ import annotations

type _FileStore = dict[str, bytes]

import datetime
import json
import pathlib

import pytest
from tests.cli_test_helper import CliRunner
from muse.core.types import blob_id, fake_id
from muse.core.object_store import object_path
from muse.core.paths import heads_dir, muse_dir, ref_path

cli = None  # argparse migration — CliRunner ignores this arg

runner = CliRunner()


# ---------------------------------------------------------------------------
# Shared helpers
# ---------------------------------------------------------------------------

def _env(root: pathlib.Path) -> Manifest:
    return {"MUSE_REPO_ROOT": str(root)}


def _init_repo(tmp_path: pathlib.Path) -> tuple[pathlib.Path, str]:
    dot_muse = muse_dir(tmp_path)
    dot_muse.mkdir()
    repo_id = fake_id("repo")
    (dot_muse / "repo.json").write_text(json.dumps({
        "repo_id": repo_id,
        "domain": "code",
        "default_branch": "main",
        "created_at": "2025-01-01T00:00:00+00:00",
    }), encoding="utf-8")
    (dot_muse / "HEAD").write_text("ref: refs/heads/main", encoding="utf-8")
    (dot_muse / "refs" / "heads").mkdir(parents=True)
    (dot_muse / "snapshots").mkdir()
    (dot_muse / "commits").mkdir()
    (dot_muse / "objects").mkdir()
    return tmp_path, repo_id


def _make_commit(root: pathlib.Path, repo_id: str, branch: str = "main",
                 message: str = "test",
                 manifest: Manifest | None = None) -> str:
    from muse.core.commits import (
    CommitRecord,
    write_commit,
    )
    from muse.core.snapshots import (
    SnapshotRecord,
    write_snapshot,
    )
    from muse.core.ids import hash_snapshot, hash_commit

    ref_file = ref_path(root, branch)
    parent_id = ref_file.read_text().strip() if ref_file.exists() else None
    m = manifest or {}
    snap_id = hash_snapshot(m)
    committed_at = datetime.datetime.now(datetime.timezone.utc)
    commit_id = hash_commit(        parent_ids=[parent_id] if parent_id else [],
        snapshot_id=snap_id, message=message,
        committed_at_iso=committed_at.isoformat(),
    )
    write_snapshot(root, SnapshotRecord(snapshot_id=snap_id, manifest=m))
    write_commit(root, CommitRecord(
        commit_id=commit_id, branch=branch,
        snapshot_id=snap_id, message=message, committed_at=committed_at,
        parent_commit_id=parent_id,
    ))
    ref_file.parent.mkdir(parents=True, exist_ok=True)
    ref_file.write_text(commit_id, encoding="utf-8")
    return commit_id


def _write_object(root: pathlib.Path, content: bytes) -> str:
    from muse.core.object_store import write_object
    obj_id = blob_id(content)
    write_object(root, obj_id, content)
    return obj_id


# ---------------------------------------------------------------------------
# Parser flag tests
# ---------------------------------------------------------------------------

class TestRegisterFlags:
    def _parse(self, *args: str) -> "argparse.Namespace":
        import argparse
        from muse.cli.commands.merge import register
        p = argparse.ArgumentParser()
        sub = p.add_subparsers()
        register(sub)
        return p.parse_args(["merge", *args])

    def test_default_json_out_is_false(self) -> None:
        ns = self._parse("feature")
        assert ns.json_out is False

    def test_json_flag_sets_json_out(self) -> None:
        ns = self._parse("--json", "feature")
        assert ns.json_out is True

    def test_j_shorthand_sets_json_out(self) -> None:
        ns = self._parse("-j", "feature")
        assert ns.json_out is True


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------

class TestMergeCLI:
    def test_merge_branch_into_main(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        base_id = _make_commit(root, repo_id, branch="main", message="base")
        (heads_dir(root) / "feature").write_text(base_id)
        obj = _write_object(root, b"feature content")
        _make_commit(root, repo_id, branch="feature", message="feature work",
                     manifest={"new_track.mid": obj})
        result = runner.invoke(cli, ["merge", "feature"], env=_env(root), catch_exceptions=False)
        assert result.exit_code == 0

    def test_merge_nonexistent_branch_fails(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        _make_commit(root, repo_id)
        result = runner.invoke(cli, ["merge", "does-not-exist"], env=_env(root))
        assert result.exit_code != 0

    def test_merge_format_json(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        base_id = _make_commit(root, repo_id, branch="main", message="base")
        (heads_dir(root) / "feature").write_text(base_id)
        _make_commit(root, repo_id, branch="feature", message="feat")
        result = runner.invoke(
            cli, ["merge", "--json", "feature"], env=_env(root), catch_exceptions=False
        )
        assert result.exit_code == 0
        data = json.loads(result.output)
        assert isinstance(data, dict)

    def test_merge_message_flag(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        base_id = _make_commit(root, repo_id, branch="main", message="base")
        (heads_dir(root) / "feature").write_text(base_id)
        _make_commit(root, repo_id, branch="feature", message="feat")
        result = runner.invoke(
            cli, ["merge", "--message", "Merge feature", "feature"],
            env=_env(root), catch_exceptions=False
        )
        assert result.exit_code == 0

    def test_merge_invalid_branch_name_rejected(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        _make_commit(root, repo_id)
        result = runner.invoke(cli, ["merge", "../traversal"], env=_env(root))
        assert result.exit_code != 0

    def test_merge_output_sanitized(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        base_id = _make_commit(root, repo_id, branch="main", message="base")
        (heads_dir(root) / "feature").write_text(base_id)
        _make_commit(root, repo_id, branch="feature", message="feat")
        result = runner.invoke(cli, ["merge", "feature"], env=_env(root), catch_exceptions=False)
        assert "\x1b" not in result.output


class TestMergeConflictWorkdir:
    """Regression: non-conflicting additions from theirs must reach the
    working tree even when a conflicted merge exits early.

    Bug: muse merge called ``raise SystemExit`` before ``_restore_from_manifest``
    when conflicts existed.  Theirs-only file additions were computed but never
    written to disk; ``muse checkout --theirs --all`` only resolved the
    conflict_paths, so ``muse code add .`` missed the new files and the merge
    commit was silently incomplete.
    """

    def _make_commit_with_files(
        self,
        root: pathlib.Path,
        repo_id: str,
        branch: str,
        files: _FileStore,
        parent_id: str | None = None,
        message: str = "commit",
    ) -> str:
        manifest: Manifest = {}
        for rel, content in files.items():
            oid = _write_object(root, content)
            manifest[rel] = oid
            dest = root / rel
            dest.parent.mkdir(parents=True, exist_ok=True)
            dest.write_bytes(content)
        return _make_commit(root, repo_id, branch=branch, message=message, manifest=manifest)

    def test_theirs_only_additions_written_to_workdir_on_conflict(
        self, tmp_path: pathlib.Path
    ) -> None:
        """Theirs-only new files must appear in the working tree after a
        conflicted merge so that ``muse code add .`` captures them."""
        root, repo_id = _init_repo(tmp_path)

        # Base: one shared file that both sides will modify (guaranteeing conflict).
        base_id = self._make_commit_with_files(
            root, repo_id, "main",
            {"shared.py": b"def foo(): pass\n"},
            message="base",
        )

        # Theirs: modifies shared.py AND adds two brand-new files.
        (heads_dir(root) / "feature").write_text(base_id)
        self._make_commit_with_files(
            root, repo_id, "feature",
            {
                "shared.py": b"def foo(): return 'theirs'\n",
                "new_security_test.py": b"# security test\n",
                "new_perf_test.py": b"# perf test\n",
            },
            message="feature: add tests + modify shared",
        )

        # Ours: also modifies shared.py (guaranteeing a conflict on that file).
        (root / "shared.py").write_bytes(b"def foo(): return 'ours'\n")
        _make_commit(
            root, repo_id, "main", message="ours: modify shared",
            manifest={"shared.py": _write_object(root, b"def foo(): return 'ours'\n")},
        )

        result = runner.invoke(cli, ["merge", "feature"], env=_env(root))

        # Merge must exit with a conflict status, not a clean merge.
        assert result.exit_code != 0, "Expected conflict exit code"
        assert "CONFLICT" in result.stderr or "conflict" in result.stderr.lower()

        # The fix: theirs-only additions MUST now exist in the working tree.
        assert (root / "new_security_test.py").exists(), (
            "new_security_test.py (theirs-only addition) must be written to the "
            "working tree even though a conflict was detected on shared.py"
        )
        assert (root / "new_perf_test.py").exists(), (
            "new_perf_test.py (theirs-only addition) must be written to the "
            "working tree even though a conflict was detected on shared.py"
        )
        assert (root / "new_security_test.py").read_bytes() == b"# security test\n"
        assert (root / "new_perf_test.py").read_bytes() == b"# perf test\n"

    def test_conflicting_file_gets_conflict_markers(
        self, tmp_path: pathlib.Path
    ) -> None:
        """Conflicting files must contain Cohen Transform markers so the user
        can see both sides and resolve manually."""
        root, repo_id = _init_repo(tmp_path)

        base_id = self._make_commit_with_files(
            root, repo_id, "main",
            {"shared.py": b"def foo(): pass\n"},
            message="base",
        )

        (heads_dir(root) / "feature").write_text(base_id)
        self._make_commit_with_files(
            root, repo_id, "feature",
            {
                "shared.py": b"def foo(): return 'theirs'\n",
                "only_on_theirs.py": b"# new\n",
            },
            message="feature",
        )

        ours_content = b"def foo(): return 'ours'\n"
        (root / "shared.py").write_bytes(ours_content)
        _make_commit(
            root, repo_id, "main", message="ours",
            manifest={"shared.py": _write_object(root, ours_content)},
        )

        runner.invoke(cli, ["merge", "feature"], env=_env(root))

        # Conflicting file must contain Cohen Transform markers.
        content = (root / "shared.py").read_text(encoding="utf-8")
        assert "<<<<<<<" in content, f"Expected conflict markers, got:\n{content}"
        assert "return 'ours'" in content
        assert "return 'theirs'" in content

        # Theirs-only addition must be present.
        assert (root / "only_on_theirs.py").exists()


class TestMergeStress:
    def test_merge_feature_with_many_files(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        base_id = _make_commit(root, repo_id, branch="main", message="base")
        (heads_dir(root) / "feature").write_text(base_id)
        manifest = {f"track_{i:03d}.mid": _write_object(root, f"data {i}".encode())
                    for i in range(30)}
        _make_commit(root, repo_id, branch="feature", message="many files", manifest=manifest)
        result = runner.invoke(cli, ["merge", "feature"], env=_env(root), catch_exceptions=False)
        assert result.exit_code == 0


# ---------------------------------------------------------------------------
# Bug: muse merge --abort must preserve staged files on disk
#
# apply_manifest(HEAD) deletes files not in the committed HEAD manifest.
# Staged-but-not-committed files are not in HEAD, so they get deleted.
# After abort, those files should still exist on disk (they are staged work).
# ---------------------------------------------------------------------------

class TestMergeAbortPreservesStagedFiles:

    def test_abort_leaves_staged_new_file_on_disk(self, tmp_path: pathlib.Path) -> None:
        """muse merge --abort must not delete a staged-but-uncommitted new file."""
        from tests.cli_test_helper import CliRunner
        r = CliRunner()
        env = {"MUSE_REPO_ROOT": str(tmp_path)}

        # Init repo via muse init so staging is wired up.
        r.invoke(cli, ["init"], env=env, catch_exceptions=False)

        # First commit: base file.
        (tmp_path / "base.py").write_text("base\n")
        r.invoke(cli, ["code", "add", "base.py"], env=env, catch_exceptions=False)
        r.invoke(cli, ["commit", "-m", "base"], env=env, catch_exceptions=False)

        # Create feature branch.
        r.invoke(cli, ["checkout", "-b", "feature"], env=env, catch_exceptions=False)
        (tmp_path / "feature.py").write_text("feature\n")
        r.invoke(cli, ["code", "add", "feature.py"], env=env, catch_exceptions=False)
        r.invoke(cli, ["commit", "-m", "feature"], env=env, catch_exceptions=False)

        # Back to main, stage a new file (don't commit).
        r.invoke(cli, ["checkout", "main"], env=env, catch_exceptions=False)
        staged_file = tmp_path / "staged_work.py"
        staged_file.write_text("my staged work\n")
        r.invoke(cli, ["code", "add", "staged_work.py"], env=env, catch_exceptions=False)
        assert staged_file.exists()

        # Trigger a merge that conflicts (feature changed base.py, main also will).
        # Simplest: just start and abort immediately.
        r.invoke(cli, ["merge", "--force", "feature"], env=env)

        # Abort the merge.
        r.invoke(cli, ["merge", "--abort"], env=env, catch_exceptions=False)

        # The staged file must still be on disk.
        assert staged_file.exists(), \
            "muse merge --abort deleted a staged-but-uncommitted file from disk"
        assert staged_file.read_text() == "my staged work\n"

    def test_abort_leaves_staged_modification_on_disk(self, tmp_path: pathlib.Path) -> None:
        """muse merge --abort must not revert a staged modification."""
        from tests.cli_test_helper import CliRunner
        r = CliRunner()
        env = {"MUSE_REPO_ROOT": str(tmp_path)}

        r.invoke(cli, ["init"], env=env, catch_exceptions=False)
        (tmp_path / "work.py").write_text("v1\n")
        r.invoke(cli, ["code", "add", "work.py"], env=env, catch_exceptions=False)
        r.invoke(cli, ["commit", "-m", "base"], env=env, catch_exceptions=False)

        r.invoke(cli, ["checkout", "-b", "feature"], env=env, catch_exceptions=False)
        (tmp_path / "other.py").write_text("other\n")
        r.invoke(cli, ["code", "add", "other.py"], env=env, catch_exceptions=False)
        r.invoke(cli, ["commit", "-m", "feature"], env=env, catch_exceptions=False)

        r.invoke(cli, ["checkout", "main"], env=env, catch_exceptions=False)
        # Stage a modification to work.py.
        (tmp_path / "work.py").write_text("v2\n")
        r.invoke(cli, ["code", "add", "work.py"], env=env, catch_exceptions=False)

        r.invoke(cli, ["merge", "--force", "feature"], env=env)
        r.invoke(cli, ["merge", "--abort"], env=env, catch_exceptions=False)

        # The staged version (v2) must be on disk, not the committed version (v1).
        assert (tmp_path / "work.py").read_text() == "v2\n", \
            "muse merge --abort reverted a staged modification"


# ---------------------------------------------------------------------------
# Bug: one-sided changes must not produce false conflicts at CLI level
#
# Scenario: our branch doesn't touch file A; theirs changes file A.
# merge must complete cleanly — no conflicts, file A takes theirs' version.
# ---------------------------------------------------------------------------

class TestOneSidedChangeNeverConflicts:

    def test_theirs_only_changes_file_clean_merge(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        # base: two files
        base_id = _make_commit(root, repo_id, branch="main", message="base", manifest={
            "describe.py": _write_object(root, b"old describe\n"),
            "pyproject.toml": _write_object(root, b"version = 1\n"),
        })
        # feature branch: only changes describe.py and pyproject.toml
        (heads_dir(root) / "feature").write_text(base_id)
        _make_commit(root, repo_id, branch="feature", message="fix", manifest={
            "describe.py": _write_object(root, b"fixed describe\n"),
            "pyproject.toml": _write_object(root, b"version = 2\n"),
        })
        # ours (main) makes an unrelated commit without touching those files
        _make_commit(root, repo_id, branch="main", message="our unrelated work", manifest={
            "describe.py": _write_object(root, b"old describe\n"),
            "pyproject.toml": _write_object(root, b"version = 1\n"),
            "new_file.py": _write_object(root, b"new\n"),
        })
        result = runner.invoke(
            cli, ["merge", "--force", "--json", "feature"],
            env=_env(root), catch_exceptions=False
        )
        assert result.exit_code == 0
        data = json.loads(result.output)
        assert data["status"] in ("merged", "fast_forward")
        assert data["conflicts"] == []

    def test_both_sides_change_different_files_clean_merge(self, tmp_path: pathlib.Path) -> None:
        root, repo_id = _init_repo(tmp_path)
        base_id = _make_commit(root, repo_id, branch="main", message="base", manifest={
            "a.py": _write_object(root, b"a\n"),
            "b.py": _write_object(root, b"b\n"),
        })
        (heads_dir(root) / "feature").write_text(base_id)
        # feature: changes b.py only
        _make_commit(root, repo_id, branch="feature", message="change b", manifest={
            "a.py": _write_object(root, b"a\n"),
            "b.py": _write_object(root, b"b-theirs\n"),
        })
        # main: changes a.py only
        _make_commit(root, repo_id, branch="main", message="change a", manifest={
            "a.py": _write_object(root, b"a-ours\n"),
            "b.py": _write_object(root, b"b\n"),
        })
        result = runner.invoke(
            cli, ["merge", "--force", "--json", "feature"],
            env=_env(root), catch_exceptions=False
        )
        assert result.exit_code == 0
        data = json.loads(result.output)
        assert data["conflicts"] == []


# ---------------------------------------------------------------------------
# Bug: muse commit completing a merge must produce a hash-verified commit
#
# After resolving conflicts and running `muse commit`, the resulting commit
# (with two parents) must pass write_commit's content-hash verification.
# Previously this raised ValueError with "incoming record failed hash
# verification", permanently blocking merge completion.
# ---------------------------------------------------------------------------

class TestMergeCommitCompletion:

    def test_commit_after_conflict_resolution_passes_hash_verification(
        self, tmp_path: pathlib.Path
    ) -> None:
        from muse.core.commits import read_commit

        root, repo_id = _init_repo(tmp_path)
        base_id = _make_commit(root, repo_id, branch="main", message="base", manifest={
            "shared.py": _write_object(root, b"base\n"),
        })
        (heads_dir(root) / "feature").write_text(base_id)
        # feature: changes shared.py
        _make_commit(root, repo_id, branch="feature", message="theirs", manifest={
            "shared.py": _write_object(root, b"theirs\n"),
        })
        # main: also changes shared.py (true conflict)
        _make_commit(root, repo_id, branch="main", message="ours", manifest={
            "shared.py": _write_object(root, b"ours\n"),
        })
        (root / "shared.py").write_bytes(b"ours\n")

        # Trigger the merge — expect conflict
        merge_result = runner.invoke(
            cli, ["merge", "--force", "--json", "feature"], env=_env(root)
        )
        data = json.loads(merge_result.output)
        assert data["status"] == "conflict"

        # Resolve via checkout --theirs (updates merge state conflict list)
        runner.invoke(cli, ["checkout", "--theirs", "shared.py"], env=_env(root), catch_exceptions=False)
        runner.invoke(cli, ["code", "add", "shared.py"], env=_env(root), catch_exceptions=False)

        # Complete the merge
        commit_result = runner.invoke(
            cli, ["commit", "--json", "-m", "merge: resolve conflict"],
            env=_env(root), catch_exceptions=False
        )
        assert commit_result.exit_code == 0, f"commit failed: {commit_result.output}"
        commit_data = json.loads(commit_result.output)
        assert "commit_id" in commit_data

        # The commit must have two parents and pass hash verification
        cid = commit_data["commit_id"]
        rec = read_commit(root, cid)
        assert rec is not None
        assert rec.parent2_commit_id is not None, "merge commit must have two parents"
        assert rec.commit_id == cid