"""Isolated workspaces and the shared warm cache for RLM children (R12)."""

from __future__ import annotations

import asyncio
import importlib
import os
import shutil
import subprocess
import tempfile
import unittest
from contextlib import contextmanager
from collections.abc import Iterator
from pathlib import Path
from unittest.mock import AsyncMock, patch

bash_module = importlib.import_module("rlm.bash")
rlm_module = importlib.import_module("rlm")

_HAS_GIT = shutil.which("git") is not None
_HAS_JJ = shutil.which("jj") is not None

_SPAWN_PAYLOAD = {
    "rlm_child_id": "sub-a1b2c3d4",
    "name": "worker-1",
    "session_dir": "/tmp/children/sub-a1b2c3d4",
    "model": "test/model",
}

_SUBAGENT_PAYLOAD = {
    "rlm_child_id": "sub-a1b2c3d4",
    "active_session_id": "child-active-1",
    "session_id": "child-session-1",
    "session_name": "worker-1",
    "session_dir": "/tmp/children/sub-a1b2c3d4",
    "status": "running",
}


def _run(command: list[str], cwd: Path) -> subprocess.CompletedProcess[str]:
    return subprocess.run(command, cwd=str(cwd), capture_output=True, text=True)


@contextmanager
def _without_launcher_cache_env() -> Iterator[None]:
    """The launcher may already run with a shared cache; these tests assert the
    runtime's own decision, so the ambient variables are hidden for the body."""
    keys = ("CARGO_TARGET_DIR", "PRIME_AGENT_SHARED_CACHE", "PRIME_AGENT_SHARED_CACHE_DIR")
    saved = {key: os.environ.pop(key) for key in keys if key in os.environ}
    try:
        yield
    finally:
        os.environ.update(saved)


class RepoFixture:
    """A throwaway repository plus the working directory the runtime would sit in."""

    def __init__(self, vcs: str, *, cargo: bool = False) -> None:
        self._tmp = tempfile.TemporaryDirectory(prefix="prime-agent-r12-")
        self._previous = Path.cwd()
        self.root = Path(self._tmp.name).resolve()
        self.vcs = vcs
        (self.root / "a.txt").write_text("a\n", encoding="utf-8")
        if cargo:
            (self.root / "Cargo.toml").write_text(
                '[package]\nname = "fixture"\nversion = "0.1.0"\n', encoding="utf-8"
            )
        if vcs == "jj":
            _run(["jj", "git", "init"], self.root)
        else:
            _run(["git", "init", "-q", "-b", "main"], self.root)
            _run(["git", "config", "user.email", "test@example.com"], self.root)
            _run(["git", "config", "user.name", "test"], self.root)
            _run(["git", "add", "a.txt"], self.root)
            if cargo:
                _run(["git", "add", "Cargo.toml"], self.root)
            _run(["git", "commit", "-q", "-m", "init"], self.root)

    def __enter__(self) -> "RepoFixture":
        os.chdir(self.root)
        return self

    def __exit__(self, *exc: object) -> None:
        os.chdir(self._previous)
        self._tmp.cleanup()


@unittest.skipUnless(_HAS_GIT, "git is not installed")
class GitWorkspaceTests(unittest.TestCase):
    def test_adds_a_linked_checkout_inside_the_repository(self) -> None:
        with RepoFixture("git") as repo:
            workspace = bash_module.create_workspace("worker-1")

            self.assertEqual(workspace.vcs, "git")
            self.assertEqual(workspace.path, repo.root / ".worktrees" / "worker-1")
            self.assertTrue(workspace.path.is_dir())
            self.assertIn(repo.root, workspace.path.parents)
            # A linked checkout, not a copy: it points at the repository it came from.
            pointer = (workspace.path / ".git").read_text(encoding="utf-8")
            self.assertTrue(pointer.startswith("gitdir:"))
            self.assertIn(str(repo.root), pointer)
            self.assertIn(str(workspace.path), _run(["git", "worktree", "list"], repo.root).stdout)
            # The checkout cannot show up as an untracked tree in its own repository.
            exclude = (repo.root / ".git" / "info" / "exclude").read_text(encoding="utf-8")
            self.assertIn(".worktrees/", exclude)
            self.assertEqual(_run(["git", "status", "--porcelain"], repo.root).stdout.strip(), "")

    def test_prunes_the_checkout_and_its_worktree_entry(self) -> None:
        with RepoFixture("git") as repo:
            workspace = bash_module.create_workspace("worker-1")

            self.assertIsNone(workspace.prune())

            self.assertFalse(workspace.path.exists())
            self.assertNotIn("worker-1", _run(["git", "worktree", "list"], repo.root).stdout)
            self.assertEqual(_run(["git", "status", "--porcelain"], repo.root).stdout.strip(), "")
            # A repeat prune reports instead of raising: disposal may race.
            self.assertIsInstance(workspace.prune(), str)

    def test_prune_survives_a_failing_vcs_command(self) -> None:
        with RepoFixture("git"):
            workspace = bash_module.create_workspace("worker-1")

            with patch.object(bash_module, "_run_vcs", side_effect=RuntimeError("no git")):
                note = workspace.prune()

            self.assertIsInstance(note, str)
            self.assertFalse(workspace.path.exists())

    def test_rejects_a_name_that_leaves_the_repository(self) -> None:
        with RepoFixture("git") as repo:
            with self.assertRaises(RuntimeError):
                bash_module.create_workspace("../../outside")
            with self.assertRaises(ValueError):
                bash_module.create_workspace("   ")
            self.assertFalse((repo.root.parent / "outside").exists())

    def test_refuses_a_directory_that_is_not_a_repository(self) -> None:
        with tempfile.TemporaryDirectory(prefix="prime-agent-r12-plain-") as tmp:
            previous = Path.cwd()
            os.chdir(tmp)
            try:
                with self.assertRaises(RuntimeError):
                    bash_module.create_workspace("worker-1")
            finally:
                os.chdir(previous)


@unittest.skipUnless(_HAS_JJ, "jj is not installed")
class JjWorkspaceTests(unittest.TestCase):
    def test_adds_and_prunes_a_jj_workspace(self) -> None:
        with RepoFixture("jj") as repo:
            workspace = bash_module.create_workspace("worker-1")

            self.assertEqual(workspace.vcs, "jj")
            self.assertTrue(workspace.path.is_dir())
            self.assertEqual(workspace.path, repo.root / ".worktrees" / "worker-1")
            # A linked workspace, not a copy: `.jj/repo` names the main store.
            self.assertTrue((workspace.path / ".jj" / "repo").is_file())
            self.assertIn("worker-1", _run(["jj", "workspace", "list"], repo.root).stdout)
            # Both the checkout and the main working copy resolve to one repository.
            os.chdir(workspace.path)
            self.assertEqual(bash_module._repo_root(), workspace.path)
            self.assertEqual(bash_module._shared_repo_root(workspace.path), repo.root)
            os.chdir(repo.root)

            self.assertIsNone(workspace.prune())

            self.assertFalse(workspace.path.exists())
            self.assertNotIn("worker-1", _run(["jj", "workspace", "list"], repo.root).stdout)


@unittest.skipUnless(_HAS_GIT, "git is not installed")
class SharedCacheTests(unittest.TestCase):
    def test_every_session_of_a_cargo_repository_builds_into_one_cache(self) -> None:
        with _without_launcher_cache_env(), RepoFixture("git", cargo=True) as repo:
            cache = repo.root / "target" / "subagents-shared"

            self.assertEqual(bash_module._child_env()["CARGO_TARGET_DIR"], str(cache))
            self.assertEqual((cache / ".gitignore").read_text(encoding="utf-8"), "*\n")
            self.assertEqual(_run(["git", "status", "--porcelain"], repo.root).stdout.strip(), "")

            workspace = bash_module.create_workspace("worker-1")
            os.chdir(workspace.path)
            try:
                self.assertEqual(bash_module._child_env()["CARGO_TARGET_DIR"], str(cache))
            finally:
                os.chdir(repo.root)

    def test_a_repository_that_does_not_build_with_cargo_is_left_alone(self) -> None:
        with _without_launcher_cache_env(), RepoFixture("git") as repo:
            self.assertNotIn("CARGO_TARGET_DIR", bash_module._child_env())
            self.assertFalse((repo.root / "target").exists())

    def test_opt_out_leaves_the_toolchain_untouched(self) -> None:
        with _without_launcher_cache_env(), RepoFixture("git", cargo=True) as repo:
            with patch.dict(os.environ, {"PRIME_AGENT_SHARED_CACHE": "0"}):
                self.assertNotIn("CARGO_TARGET_DIR", bash_module._child_env())

            self.assertFalse((repo.root / "target").exists())

    def test_a_value_from_the_launcher_wins(self) -> None:
        with RepoFixture("git", cargo=True):
            with patch.dict(os.environ, {"CARGO_TARGET_DIR": "/opt/launcher-target"}):
                self.assertEqual(bash_module._child_env()["CARGO_TARGET_DIR"], "/opt/launcher-target")

    def test_an_explicit_cache_directory_wins(self) -> None:
        with _without_launcher_cache_env(), RepoFixture("git", cargo=True) as repo:
            with patch.dict(os.environ, {"PRIME_AGENT_SHARED_CACHE_DIR": str(repo.root / "elsewhere")}):
                self.assertEqual(bash_module._child_env()["CARGO_TARGET_DIR"], str(repo.root / "elsewhere"))


@unittest.skipUnless(_HAS_GIT, "git is not installed")
class SpawnWorkspaceTests(unittest.TestCase):
    def tearDown(self) -> None:
        rlm_module._WORKSPACES.clear()

    def test_spawn_starts_the_child_in_its_own_checkout(self) -> None:
        with RepoFixture("git") as repo:
            host_request = AsyncMock(return_value=dict(_SPAWN_PAYLOAD))

            with patch.object(rlm_module, "host_request", host_request):
                handle = asyncio.run(
                    rlm_module.spawn("task", name="worker-1", model="test/model", workspace="isolated")
                )

            expected = repo.root / ".worktrees" / "worker-1"
            self.assertEqual(handle.workspace, expected)
            self.assertTrue(expected.is_dir())
            request = host_request.await_args.args[1]
            self.assertEqual(request["kwargs"]["cwd"], str(expected))
            self.assertEqual(request["kwargs"]["name"], "worker-1")

    def test_spawn_without_isolation_reuses_the_parent_working_copy(self) -> None:
        with RepoFixture("git") as repo:
            host_request = AsyncMock(return_value=dict(_SPAWN_PAYLOAD))

            with patch.object(rlm_module, "host_request", host_request):
                handle = asyncio.run(rlm_module.spawn("task", name="worker-1", model="test/model"))

            self.assertIsNone(handle.workspace)
            self.assertNotIn("cwd", host_request.await_args.args[1]["kwargs"])
            self.assertFalse((repo.root / ".worktrees").exists())

    def test_deleting_a_child_prunes_its_checkout(self) -> None:
        with RepoFixture("git") as repo:
            with patch.object(rlm_module, "host_request", AsyncMock(return_value=dict(_SPAWN_PAYLOAD))):
                asyncio.run(
                    rlm_module.spawn("task", name="worker-1", model="test/model", workspace="isolated")
                )
            checkout = repo.root / ".worktrees" / "worker-1"
            self.assertTrue(checkout.is_dir())

            with patch.object(
                rlm_module, "host_request", AsyncMock(return_value={"subagent": dict(_SUBAGENT_PAYLOAD)})
            ):
                asyncio.run(rlm_module.delete_subagent("sub-a1b2c3d4"))

            self.assertFalse(checkout.exists())
            self.assertEqual(rlm_module._WORKSPACES, {})

    def test_a_rejected_spawn_prunes_its_checkout(self) -> None:
        with RepoFixture("git") as repo:
            host_request = AsyncMock(side_effect=RuntimeError("Unsupported rlm.spawn kwargs: cwd"))

            with patch.object(rlm_module, "host_request", host_request):
                with self.assertRaises(RuntimeError):
                    asyncio.run(
                        rlm_module.spawn("task", name="worker-1", model="test/model", workspace="isolated")
                    )

            self.assertFalse((repo.root / ".worktrees" / "worker-1").exists())

    def test_isolated_workspace_refuses_other_values_and_remote_hosts(self) -> None:
        with RepoFixture("git"):
            with self.assertRaises(TypeError):
                asyncio.run(rlm_module.spawn("task", name="worker-1", model="test/model", workspace="shared"))
            with self.assertRaises(ValueError):
                asyncio.run(
                    rlm_module.spawn(
                        "task", name="worker-1", model="test/model", workspace="isolated", host="chernobog"
                    )
                )