"""
进程/子进程清理回归测试

目标：确保 aacode 在执行任务（工具调用）过程中 spawn 的子进程与后台 asyncio
任务，在任务完成/超时后全部退出，不残留孤儿或僵尸进程，也不随调用次数累积。

覆盖的 spawn 点：
- atomic_tools.run_shell          -> asyncio.create_subprocess_shell + 两个 reader task
- code_tools.execute_python       -> asyncio.create_subprocess_exec
- context_manager.get_context     -> asyncio.create_subprocess_shell (find)

依赖 psutil（dev 依赖）；未安装时自动跳过。
"""
import asyncio
import sys
import tempfile
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

import pytest

psutil = pytest.importorskip("psutil")

from aacode.tools.atomic_tools import AtomicTools
from aacode.tools.code_tools import CodeTools
from aacode.utils.safety import SafetyGuard
from aacode.utils.context_manager import ContextManager


def _live_children(proc):
    """当前进程的存活（非僵尸）子进程 pid 集合"""
    alive = set()
    for child in proc.children(recursive=True):
        try:
            if child.status() != psutil.STATUS_ZOMBIE and child.is_running():
                alive.add(child.pid)
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            pass
    return alive


def _leaked_children(proc, baseline, timeout=5.0):
    """轮询等待，返回相对 baseline 新增且仍存活的子进程 pid（应为空）"""
    deadline = time.time() + timeout
    leaked = _live_children(proc) - baseline
    while leaked and time.time() < deadline:
        time.sleep(0.1)
        leaked = _live_children(proc) - baseline
    return leaked


def _extra_tasks():
    """当前事件循环中除自身外仍未完成的任务"""
    current = asyncio.current_task()
    return [t for t in asyncio.all_tasks() if t is not current and not t.done()]


@pytest.fixture
def tools():
    with tempfile.TemporaryDirectory() as td:
        project_path = Path(td).resolve()
        (project_path / "a.py").write_text("print('x')\n", encoding="utf-8")
        (project_path / "b.md").write_text("# doc\n", encoding="utf-8")
        sg = SafetyGuard(project_path)
        yield {
            "path": project_path,
            "atomic": AtomicTools(project_path, sg),
            "code": CodeTools(project_path, sg),
            "context": ContextManager(project_path),
        }


class TestSubprocessCleanup:
    @pytest.mark.asyncio
    async def test_run_shell_no_leftover_child(self, tools):
        proc = psutil.Process()
        baseline = _live_children(proc)
        result = await tools["atomic"].run_shell("echo hello", timeout=10)
        assert result["success"] is True
        assert _leaked_children(proc, baseline) == set(), "run_shell 后残留子进程"

    @pytest.mark.asyncio
    async def test_run_shell_timeout_kills_child(self, tools):
        """超时命令必须被 terminate/kill，不能留下存活的 sleep 进程"""
        proc = psutil.Process()
        baseline = _live_children(proc)
        result = await tools["atomic"].run_shell("sleep 30", timeout=1)
        assert result.get("timeout") is True
        assert _leaked_children(proc, baseline) == set(), "超时后 sleep 子进程未被清理"

    @pytest.mark.asyncio
    async def test_run_shell_no_leftover_reader_tasks(self, tools):
        """run_shell 的 stdout/stderr reader task 结束后不应残留"""
        before = len(_extra_tasks())
        await tools["atomic"].run_shell("echo hi", timeout=10)
        await asyncio.sleep(0)
        after = len(_extra_tasks())
        assert after <= before, "run_shell 残留了后台 reader task"

    @pytest.mark.asyncio
    async def test_execute_python_no_leftover_child(self, tools):
        proc = psutil.Process()
        baseline = _live_children(proc)
        result = await tools["code"].execute_python("print(1 + 1)", timeout=10)
        assert result.get("success") is True
        assert _leaked_children(proc, baseline) == set(), "execute_python 后残留子进程"

    @pytest.mark.asyncio
    async def test_execute_python_timeout_no_live_child(self, tools):
        """执行超时后不应有存活的 python 子进程"""
        proc = psutil.Process()
        baseline = _live_children(proc)
        result = await tools["code"].execute_python(
            "import time; time.sleep(30)", timeout=1
        )
        assert "error" in result
        assert _leaked_children(proc, baseline) == set(), "超时后 python 子进程未被清理"

    @pytest.mark.asyncio
    async def test_get_context_no_leftover_child(self, tools):
        """context 收集（find 子进程）结束后不应残留"""
        proc = psutil.Process()
        baseline = _live_children(proc)
        await tools["context"].get_context()
        assert _leaked_children(proc, baseline) == set(), "get_context 后残留 find 子进程"

    @pytest.mark.asyncio
    async def test_repeated_calls_do_not_accumulate(self, tools):
        """连续多次工具调用后，子进程数应回到基线（无累积泄漏）"""
        proc = psutil.Process()
        baseline = _live_children(proc)
        for i in range(8):
            await tools["atomic"].run_shell(f"echo {i}", timeout=10)
            await tools["code"].execute_python(f"print({i})", timeout=10)
            await tools["context"].get_context()
        assert _leaked_children(proc, baseline) == set(), "多次调用后子进程累积"

    @pytest.mark.asyncio
    async def test_mixed_normal_and_timeout(self, tools):
        """正常与超时命令交替执行后仍无残留"""
        proc = psutil.Process()
        baseline = _live_children(proc)
        await tools["atomic"].run_shell("echo ok", timeout=10)
        await tools["atomic"].run_shell("sleep 20", timeout=1)
        await tools["atomic"].run_shell("echo again", timeout=10)
        assert _leaked_children(proc, baseline) == set()


class TestFullTaskCleanup:
    """端到端：一次 AICoder 任务结束后不残留子进程"""

    @pytest.mark.asyncio
    async def test_aicoder_run_leaves_no_children(self, monkeypatch):
        from aacode import main as main_mod

        proc = psutil.Process()
        baseline = _live_children(proc)

        with tempfile.TemporaryDirectory() as td:
            project_path = Path(td)

            coder = main_mod.AICoder(project_path=str(project_path))
            atomic = AtomicTools(coder.target_project, coder.safety_guard)

            async def fake_execute(*a, **kw):
                # 模拟任务过程中真实 spawn 子进程后结束
                await atomic.run_shell("echo work", timeout=10)
                await coder.context_manager.get_context()
                return {"success": True, "result": "done"}

            monkeypatch.setattr(coder.main_agent, "execute", fake_execute)

            await coder.run("dummy task")

        assert _leaked_children(proc, baseline) == set(), "任务结束后仍有子进程存活"
