"""DockerExecutor 单元测试。

mock subprocess(`docker exec` 命令的实际跑由部署机 smoke 验,RUN.md 有 5 条命令)。
覆盖关键路径:
- 信任域 dispatch:host 工具直通 / container 工具走 docker exec
- argv 形态:--user / --workdir / setsid / bash -c / python <script>
- tmp .py:写到 host 侧 `.zcbot_tmp/<task_id>/`,执行完 unlink,无残留
- timeout / cancel:Popen.kill() 兜底
- schemas() / has_tool() 透传 host
"""
from __future__ import annotations

import sys
import tempfile
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch
from uuid import uuid4

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from core.executor import ExecCtx, ToolResult
from core.executor_docker import DockerExecutor, TMP_SUBDIR
from core.executor_host import HostExecutor


class FakePool:
    """SandboxPool 替身:ensure 返固定容器名,mark_active 记录调用。"""

    def __init__(self):
        self.ensure_calls = []
        self.mark_active_calls = []

    def ensure(self, user_id):
        name = f"zcbot-sandbox-{user_id}"
        self.ensure_calls.append(user_id)
        return name

    def mark_active(self, user_id):
        self.mark_active_calls.append(user_id)


class FakeTool:
    """tools.base.Tool 替身:execute 返串,schema 暴露 name + 空 parameters。"""

    def __init__(self, name, output="ok"):
        self.name = name
        self._output = output
        self.execute_calls = []

    @property
    def schema(self):
        return {"type": "function", "function": {"name": self.name}}

    def execute(self, **kwargs):
        self.execute_calls.append(kwargs)
        return self._output


def make_executor(tools_dict=None):
    """构造 DockerExecutor + FakePool + tmp user_root。返回 (executor, pool, tmp_dir)。"""
    tmp = tempfile.mkdtemp()
    user_root = Path(tmp) / "users" / "u1"
    user_root.mkdir(parents=True)
    working_dir = user_root / "demo"
    working_dir.mkdir()

    if tools_dict is None:
        tools_dict = {
            "read": FakeTool("read", "READ_OUT"),
            "shell": FakeTool("shell"),  # host shell 不应被调用
            "run_python": FakeTool("run_python"),
        }
    host = HostExecutor(tools_dict)
    pool = FakePool()
    executor = DockerExecutor(
        host=host,
        pool=pool,
        user_id=uuid4(),
        user_root=user_root,
        working_dir=working_dir,
    )
    return executor, pool, Path(tmp)


def make_ctx(executor):
    return ExecCtx(
        user_id=executor.user_id,
        task_id=uuid4(),
        working_dir=executor.working_dir,
        cancel_check=None,
    )


class TestHostPassthrough(unittest.TestCase):
    """非 container tool 直通 host backend,不调 pool / subprocess。"""

    def test_read_passthrough_to_host(self):
        executor, pool, _ = make_executor()
        ctx = make_ctx(executor)
        result = executor.call_tool("read", {"file": "x"}, ctx)
        self.assertEqual(result.content, "READ_OUT")
        self.assertEqual(result.exit_code, 0)
        self.assertEqual(pool.ensure_calls, [])
        self.assertEqual(pool.mark_active_calls, [])

    def test_schemas_and_has_tool_from_host(self):
        executor, _, _ = make_executor()
        names = [s["function"]["name"] for s in executor.schemas()]
        self.assertIn("read", names)
        self.assertIn("shell", names)
        self.assertTrue(executor.has_tool("shell"))
        self.assertFalse(executor.has_tool("nope"))


class TestShellExec(unittest.TestCase):
    """shell 调用走 docker exec subprocess,argv 形态正确。"""

    def test_shell_invokes_docker_exec(self):
        executor, pool, _ = make_executor()
        ctx = make_ctx(executor)

        proc = MagicMock()
        proc.communicate.return_value = ("hello\n", "")
        proc.returncode = 0

        with patch("core.executor_docker.subprocess.Popen", return_value=proc) as popen:
            result = executor.call_tool("shell", {"command": "echo hello"}, ctx)

        self.assertIn("[stdout]\nhello", result.content)
        self.assertIn("[exit 0]", result.content)
        self.assertEqual(result.exit_code, 0)

        argv = popen.call_args[0][0]
        self.assertEqual(argv[:2], ["docker", "exec"])
        self.assertIn("--user", argv)
        self.assertIn("--workdir", argv)
        # workdir 应是 /workspace/demo(working_dir 相对 user_root)
        self.assertEqual(argv[argv.index("--workdir") + 1], "/workspace/demo")
        # container name = zcbot-sandbox-<uid>
        container_idx = argv.index(f"zcbot-sandbox-{executor.user_id}")
        # setsid bash -c 必须出现且紧跟 container 之后
        self.assertEqual(argv[container_idx + 1:], ["setsid", "bash", "-c", "echo hello"])

        self.assertEqual(pool.ensure_calls, [executor.user_id])
        self.assertEqual(pool.mark_active_calls, [executor.user_id])

    def test_shell_bad_args(self):
        executor, _, _ = make_executor()
        ctx = make_ctx(executor)
        result = executor.call_tool("shell", {"command": ""}, ctx)
        self.assertIn("[Error]", result.content)
        self.assertEqual(result.exit_code, 2)

    def test_shell_timeout(self):
        executor, pool, _ = make_executor()
        ctx = make_ctx(executor)

        import subprocess as real_subprocess
        proc = MagicMock()
        # 第一次 communicate 抛 TimeoutExpired,第二次(kill 后)返空
        proc.communicate.side_effect = [
            real_subprocess.TimeoutExpired(cmd="docker", timeout=0.5),
            ("", "killed\n"),
        ]
        proc.returncode = -9

        with patch("core.executor_docker.subprocess.Popen", return_value=proc), \
             patch("core.executor_docker.time.monotonic", side_effect=[0, 100]):
            result = executor.call_tool("shell", {"command": "sleep 9999", "timeout": 1}, ctx)

        self.assertIn("timed out after 1s", result.content)
        self.assertEqual(result.exit_code, 124)
        proc.kill.assert_called_once()

    def test_shell_cancel(self):
        executor, _, _ = make_executor()
        ctx = ExecCtx(
            user_id=executor.user_id,
            task_id=uuid4(),
            working_dir=executor.working_dir,
            cancel_check=lambda: True,  # 立即 cancel
        )

        import subprocess as real_subprocess
        proc = MagicMock()
        proc.communicate.side_effect = [
            real_subprocess.TimeoutExpired(cmd="docker", timeout=0.5),
            ("", ""),
        ]
        proc.returncode = -15

        with patch("core.executor_docker.subprocess.Popen", return_value=proc):
            result = executor.call_tool("shell", {"command": "sleep 9999"}, ctx)

        self.assertIn("cancelled by user", result.content)
        self.assertEqual(result.exit_code, 130)
        proc.kill.assert_called_once()


class TestRunPython(unittest.TestCase):
    """run_python:tmp .py 落 user_root/.zcbot_tmp/<task_id>/,跑完 unlink。"""

    def test_run_python_tmp_script(self):
        executor, pool, tmp_root = make_executor()
        ctx = make_ctx(executor)

        proc = MagicMock()
        proc.communicate.return_value = ("42\n", "")
        proc.returncode = 0

        captured_argv = []

        def _popen(argv, **kwargs):
            captured_argv.append(argv)
            return proc

        with patch("core.executor_docker.subprocess.Popen", side_effect=_popen):
            result = executor.call_tool(
                "run_python", {"code": "print(42)"}, ctx
            )

        self.assertIn("[stdout]\n42", result.content)
        self.assertEqual(result.exit_code, 0)

        argv = captured_argv[0]
        # 末尾形态:setsid python /workspace/.zcbot_tmp/<task_id>/<rand>.py
        self.assertEqual(argv[-3], "setsid")
        self.assertEqual(argv[-2], "python")
        self.assertTrue(argv[-1].startswith(f"/workspace/{TMP_SUBDIR}/{ctx.task_id}/"))
        self.assertTrue(argv[-1].endswith(".py"))
        # PYTHONIOENCODING / PYTHONPATH 注入
        env_kvs = [argv[i + 1] for i, a in enumerate(argv) if a == "-e"]
        self.assertIn("PYTHONIOENCODING=utf-8", env_kvs)
        self.assertIn("PYTHONPATH=/workspace", env_kvs)

        # host 侧 tmp 已 unlink(目录可能仍在,无所谓 —— ensure 容器时会重新 mkdir)
        tmp_subroot = executor.user_root / TMP_SUBDIR / str(ctx.task_id)
        leftover = list(tmp_subroot.glob("*.py")) if tmp_subroot.exists() else []
        self.assertEqual(leftover, [], f"tmp .py not cleaned up: {leftover}")

    def test_run_python_bad_code_type(self):
        executor, _, _ = make_executor()
        ctx = make_ctx(executor)
        result = executor.call_tool("run_python", {"code": 123}, ctx)
        self.assertIn("[Error]", result.content)
        self.assertEqual(result.exit_code, 2)

    def test_run_python_cleans_tmp_on_exception(self):
        """Popen 抛异常时 tmp .py 仍要被清理(finally 兜底)。"""
        executor, _, _ = make_executor()
        ctx = make_ctx(executor)

        with patch(
            "core.executor_docker.subprocess.Popen",
            side_effect=RuntimeError("boom"),
        ):
            result = executor.call_tool("run_python", {"code": "x"}, ctx)

        self.assertIn("[Error executing run_python via docker]", result.content)
        self.assertEqual(result.exit_code, 1)
        tmp_subroot = executor.user_root / TMP_SUBDIR / str(ctx.task_id)
        leftover = list(tmp_subroot.glob("*.py")) if tmp_subroot.exists() else []
        self.assertEqual(leftover, [])


class TestUnknownTool(unittest.TestCase):
    def test_unknown_tool_goes_to_host(self):
        executor, _, _ = make_executor(tools_dict={})  # 空 host → 啥都没
        ctx = make_ctx(executor)
        result = executor.call_tool("nope", {}, ctx)
        self.assertIn("unknown tool", result.content)
        self.assertEqual(result.exit_code, 2)

    def test_container_tool_not_registered_on_host(self):
        """caps.enable_run_python=False:host 没装 run_python,docker 也应拒。"""
        executor, _, _ = make_executor(tools_dict={"read": FakeTool("read")})
        ctx = make_ctx(executor)
        result = executor.call_tool("run_python", {"code": "x"}, ctx)
        self.assertIn("unknown tool", result.content)
        self.assertEqual(result.exit_code, 2)


if __name__ == "__main__":
    unittest.main()