# -*- coding: utf-8 -*-
"""
终端命令执行工具
执行终端命令，带安全限制和失败重试
"""
import os
import re
import sys
import time
import signal
import threading
import subprocess
from typing import Optional
from .base import Tool, coerce_int


class ShellExecTool(Tool):
    """执行终端命令的工具，带安全限制。失败时自动重试，避免临时网络/IO 抖动。"""

    # 最大重试次数（不含首次执行）。例如 max_retries=2 表示最多执行 3 次
    MAX_RETRIES = 2

    # 重试退避基数（秒），实际等待 = BASE_BACKOFF * 2^attempt
    BASE_BACKOFF = 0.5

    DANGEROUS_PATTERNS = [
        "rm -rf /", "rm -rf /*", "> /dev/sda", "dd if=/dev/zero",
        "mkfs.", "chmod -R 777 /", ":(){ :|:& };:", "> /etc/",
        "del /f /q", "fdisk"
    ]

    # 正则匹配的危险模式：format 是 Windows 格式化磁盘命令（format C: /Q），
    # 但要避免误伤 PowerShell 的 "Get-Date -Format" 这类合法用法。
    DANGEROUS_REGEX = [
        re.compile(r'(^|[\s;&|])format\s+[a-zA-Z]:', re.IGNORECASE),
    ]

    # 自杀式进程杀检测：会杀掉 abyss 自身（python 进程）的命令。
    # 痛点：LLM 高频生成 `Stop-Process -Name python -Force` 来清理子进程，
    # 但 abyss 自己也是 python 进程，会被一起杀掉导致 CLI 退出。
    # 检测范围：Stop-Process -Name python* / taskkill /IM python.exe / killall python
    # 例外：用 -Id 精确指定 PID 的命令不拦截（合法用法）
    _SUICIDE_PATTERNS = [
        # PowerShell: Stop-Process ... -Name python/pythonw/python*（参数顺序无关）
        # 用 [^\n;]* 限制在同一语句内，避免跨语句误匹配
        re.compile(r'\bStop-Process\b[^\n;]*-Name\s+python\w*', re.IGNORECASE),
        # cmd: taskkill /F /IM python.exe / IM pythonw.exe
        re.compile(r'\btaskkill\b[^\n;]*\/IM\s+python\w*\.exe', re.IGNORECASE),
        # Unix: killall python / pkill -f python
        re.compile(r'\b(?:killall|pkill)\s+(?:-\w+\s+)*python\b', re.IGNORECASE),
    ]

    # 读取命令的正则模式：用于检测读敏感文件
    # 覆盖：cat / type / Get-Content / gc / Get-Item / gi
    # 注意：仅检测显式路径参数，不检测变量、管道输入
    _READ_CMD_PATTERNS = [
        # Unix: cat <path>  /  head -n N <path>  /  tail <path>  /  less <path>
        re.compile(r'\b(?:cat|head|tail|less|more)\s+["\']?([^\s"\';&|]+)["\']?'),
        # Windows cmd: type <path>
        re.compile(r'\btype\s+["\']?([^\s"\';&|]+)["\']?', re.IGNORECASE),
        # PowerShell: Get-Content <path> / gc <path>
        re.compile(r'\b(?:Get-Content|gc)\s+["\']?([^\s"\';&|]+)["\']?', re.IGNORECASE),
        # PowerShell: Get-Item <path> / gi <path>  （包含属性访问会泄露内容）
        re.compile(r'\b(?:Get-Item|gi)\s+["\']?([^\s"\';&|]+)["\']?', re.IGNORECASE),
    ]

    @property
    def name(self) -> str:
        return "shell_exec"

    @property
    def description(self) -> str:
        return "执行终端命令，支持设置工作目录和超时。失败时自动重试。"

    @property
    def parameters(self) -> dict:
        return {
            "type": "object",
            "properties": {
                "command": {"type": "string", "description": "要执行的命令"},
                "cwd": {"type": "string", "description": "工作目录"},
                "timeout": {"type": "integer", "description": "超时秒数，默认60"}
            },
            "required": ["command"]
        }

    def _is_safe(self, command: str) -> bool:
        """检查命令是否安全"""
        lowered = command.lower()
        for pattern in self.DANGEROUS_PATTERNS:
            if pattern in lowered:
                return False
        for pattern in self.DANGEROUS_REGEX:
            if pattern.search(command):
                return False
        return True

    @staticmethod
    def _detect_suicide(command: str) -> Optional[str]:
        """检测自杀式进程杀命令。命中返回原因，否则返回 None。

        痛点：LLM 生成 `Stop-Process -Name python -Force` 时会杀掉 abyss 自身
        （abyss 也是 python 进程），导致 CLI 退出。

        合法替代：用 `Start-Process ... -PassThru` 拿到进程对象，再
        `Stop-Process -Id $proc.Id` 精确杀指定 PID。
        """
        for pattern in ShellExecTool._SUICIDE_PATTERNS:
            if pattern.search(command):
                return (
                    "检测到自杀式进程杀命令（按名称杀 python 进程会同时杀掉 abyss 自身）。"
                    "请改用 -PassThru 拿到进程对象后按 -Id 精确杀指定 PID：\n"
                    "  $p = Start-Process python -ArgumentList xxx.py -PassThru\n"
                    "  Stop-Process -Id $p.Id -Force"
                )
        return None

    def _check_read_paths(self, command: str) -> Optional[str]:
        """检测读取命令的目标路径是否在工作目录外。

        防止 AI 通过 cat/Get-Content 等命令绕过 PathSecurityService 读取
        file_read 工具拒绝访问的敏感文件（如系统 hosts、其他项目源码）。

        Returns:
            拒绝原因字符串，None 表示放行。
        """
        from abyss.security.path_security import get_default_service

        svc = get_default_service()
        for pattern in self._READ_CMD_PATTERNS:
            for match in pattern.finditer(command):
                path = match.group(1)
                # 跳过 PowerShell 变量（$variable）和管道符
                if path.startswith("$") or path.startswith("|"):
                    continue
                # 跳过选项参数（-Force / -LiteralPath）
                if path.startswith("-"):
                    continue
                # 调用 PathSecurityService 校验
                check = svc.check_read(path)
                if not check.allowed and not check.needs_confirmation:
                    return f"shell_exec 读取路径被拒绝（{path}）：{check.reason}"
        return None

    @staticmethod
    def _enrich_timeout_error(command: str, timeout: int, result: dict) -> dict:
        """为超时错误附加针对性建议。目标：让 AI 看到错误后能立刻换策略。"""
        cmd_lower = command.lower()
        suggestions = []
        if re.search(r"\b(pip|pip3|python\s+-m\s+pip)\b", cmd_lower) and "install" in cmd_lower:
            suggestions.extend([
                "pip install 在国内常因 pypi 源慢而超时。",
                "建议加 -i https://pypi.tuna.tsinghua.edu.cn/simple（清华源）",
                "或阿里源 https://mirrors.aliyun.com/pypi/simple/",
                "并把 timeout 加大到 300+ 秒。",
            ])
        elif re.search(r"\bgit\s+clone\b", cmd_lower):
            suggestions.extend([
                "git clone 慢可换镜像：git clone https://mirror.ghproxy.com/<原github url>",
                "或 https://gitclone.com/<原 url>",
            ])
        elif re.search(r"\b(curl|wget|invoke-webrequest)\b", cmd_lower):
            suggestions.extend([
                "下载慢可换国内镜像或在浏览器中下载。",
            ])
        else:
            suggestions.append(
                f"已强制终止。建议把 timeout 调大（当前 {timeout}s），或拆成更小的步骤。"
            )
        if suggestions:
            new_err = result.get("error", "") + "\n[建议] " + " ".join(suggestions)
            result = dict(result)
            result["error"] = new_err
        return result

    def _wrap_powershell(self, command: str) -> str:
        """包装用户命令，强制 PowerShell 用 UTF-8 输出（避免中文乱码）。
        同时把 PYTHONIOENCODING=utf-8 写到子进程环境，让子 Python 进程也用 UTF-8 输出，
        避免 print('✓') 这类 Unicode 字符在 Windows GBK 终端上崩溃。

        关键：不能用 chcp 65001 或 [Console]::OutputEncoding = UTF8。
        两者都调用 SetConsoleOutputCP，跨进程改父控制台代码页，触发 conhost 清屏重绘
        —— 在 cmd.exe（默认 936）下表现为"第一个 shell 命令后前边记录消失"。
        改用 StreamWriter 包裹 stdout/stderr 管道写 UTF-8，完全不碰控制台 CP。
        """
        prefix = (
            "$OutputEncoding = [System.Text.Encoding]::UTF8; "
            "$u = New-Object Text.UTF8Encoding($false); "
            "$o = New-Object IO.StreamWriter([Console]::OpenStandardOutput(), $u); "
            "$o.AutoFlush = $true; [Console]::SetOut($o); "
            "$e = New-Object IO.StreamWriter([Console]::OpenStandardError(), $u); "
            "$e.AutoFlush = $true; [Console]::SetError($e); "
        )
        return prefix + command

    def _build_env(self) -> dict:
        """构造子进程环境：保留系统环境，并注入 PYTHONIOENCODING=utf-8"""
        import os
        env = os.environ.copy()
        env["PYTHONIOENCODING"] = "utf-8"
        env["PYTHONUTF8"] = "1"
        return env

    def _start_proc(self, command: str, cwd, env) -> subprocess.Popen:
        """启动子进程并返回 Popen 对象。Windows 走 PowerShell，其他平台走 sh。"""
        if sys.platform == "win32":
            return subprocess.Popen(
                ["powershell.exe", "-NoProfile", "-Command", self._wrap_powershell(command)],
                cwd=cwd,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                encoding="utf-8",
                errors="replace",
                env=env,
            )
        return subprocess.Popen(
            command,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            encoding="utf-8",
            errors="replace",
            env=env,
        )

    @staticmethod
    def _start_reader(stream, sink: list) -> threading.Thread:
        """启动后台线程，把 stream 读到 sink 列表中。读到 EOF 自然退出。

        用底层 buffer.read1() 而非 TextIOWrapper.read()：read() 会试图填满
        内部缓冲区，孙进程持有管道时它会阻塞在「再读一点」上，导致已缓冲的
        PassThru 输出也读不出来。read1() 有多少读多少，立刻返回。
        """
        def _read_stream():
            try:
                raw = getattr(stream, "buffer", stream)
                while True:
                    chunk = raw.read1(4096) if hasattr(raw, "read1") else raw.read(4096)
                    if not chunk:
                        break
                    if isinstance(chunk, bytes):
                        sink.append(chunk.decode("utf-8", errors="replace"))
                    else:
                        sink.append(chunk)
            except Exception:
                pass
        thread = threading.Thread(target=_read_stream, daemon=True)
        thread.start()
        return thread

    @staticmethod
    def _join_readers(*threads, drain_timeout: float = 2.0):
        """等待读取线程结束，最多 drain_timeout 秒。超时后线程留在后台。

        关键：孙进程可能继承管道写端（Start-Process -NoNewWindow 场景），
        此时读取线程永远读不到 EOF。我们不能等它 — 直接返回已读到的内容，
        让孙进程在后台继续运行（这正是用户启动后台服务的本意）。
        """
        deadline = time.monotonic() + drain_timeout
        for thread in threads:
            remaining = max(0.0, deadline - time.monotonic())
            thread.join(timeout=remaining)

    @staticmethod
    def _kill_tree(pid: int):
        """杀掉进程及其所有子孙进程。Windows 用 taskkill /T，POSIX 用进程组。"""
        if sys.platform == "win32":
            subprocess.run(
                f"taskkill /F /T /PID {pid}",
                shell=True,
                capture_output=True,
            )
        else:
            try:
                os.killpg(os.getpgid(pid), signal.SIGKILL)
            except (ProcessLookupError, PermissionError):
                pass

    def _run_once(self, command: str, cwd, timeout) -> dict:
        """单次执行命令（不重试）。

        用 Popen + 后台线程读取，而不是 subprocess.run(capture_output=True)。
        原因：subprocess.run 内部调用 communicate() 等待 stdout/stderr 管道 EOF，
        当用户执行 `Start-Process -NoNewWindow -PassThru <长进程>` 时，
        孙进程会继承 PowerShell 的管道写端，PowerShell 退出后管道仍被持有，
        communicate() 永远读不到 EOF → 工具卡死。

        新方案：用 proc.wait() 只等 PowerShell 进程退出（与管道 EOF 解耦），
        再用读取线程带超时地拉取已缓冲的输出。
        """
        env = self._build_env()
        proc = self._start_proc(command, cwd, env)
        stdout_chunks: list = []
        stderr_chunks: list = []
        stdout_thread = self._start_reader(proc.stdout, stdout_chunks)
        stderr_thread = self._start_reader(proc.stderr, stderr_chunks)

        try:
            proc.wait(timeout=timeout)
        except subprocess.TimeoutExpired:
            # 父进程仍在运行 — 真正的超时，杀整棵进程树（含孙进程）
            self._kill_tree(proc.pid)
            try:
                proc.wait(timeout=5)
            except subprocess.TimeoutExpired:
                pass
            self._join_readers(stdout_thread, stderr_thread, drain_timeout=1.0)
            return {
                "command": command,
                "success": False,
                "returncode": proc.returncode,
                "stdout": "".join(stdout_chunks),
                "stderr": "".join(stderr_chunks),
                "error": f"命令执行超时（{timeout}秒），已强制终止",
            }

        # 父进程已退出 — 给读取线程一小段时间把缓冲读干。
        # 若孙进程仍持有管道（Start-Process 场景），线程不会结束，
        # 此时直接返回已读内容，让孙进程继续在后台运行。
        self._join_readers(stdout_thread, stderr_thread, drain_timeout=2.0)
        return {
            "command": command,
            "success": proc.returncode == 0,
            "returncode": proc.returncode,
            "stdout": "".join(stdout_chunks),
            "stderr": "".join(stderr_chunks),
        }

    def execute(self, command: str, cwd: str = None, timeout: int = 60) -> dict:
        """执行命令，失败时按指数退避自动重试。超时场景不重试（重试也没用）"""
        # 兼容 LLM 将 timeout 序列化为字符串的情形
        timeout = coerce_int(timeout, default=60)

        # 同步 LLM 传入的 cwd 到 PathSecurityService（修复 file_write 边界不一致）
        # 场景：用户在 A 目录启动 abyss，/quit 后在 B 目录重启并 /session 恢复会话。
        # AI 从历史里"记住"原工作目录 A，通过 shell_exec(cwd=A) 执行命令成功，
        # 但 file_write 无 cwd 参数，PathSecurityService 仍用 os.getcwd()=B 做边界，
        # 导致写入 A 目录下的文件被误判为 "path outside workspace boundary"。
        # 这里让 file_write/file_read 的边界跟随 LLM 最近一次 shell_exec 的 cwd。
        if cwd:
            from abyss.security.path_security import get_default_service as _get_path_sec
            _get_path_sec().set_working_dir(cwd)

        if not self._is_safe(command):
            return {"success": False, "error": "命令被安全策略拦截"}

        # 自杀式进程杀检测：拦截 Stop-Process -Name python 等会杀掉 abyss 自身的命令
        suicide_reason = self._detect_suicide(command)
        if suicide_reason is not None:
            return {"success": False, "error": suicide_reason}

        # ToolSafetyGuard 多层纵深防护（Phase 8 CE）
        # 与 _is_safe 互补：_is_safe 偏 Windows（format C:），
        # ToolSafetyGuard 偏 Unix（fork bomb / mkfs / curl|sh / shutdown 等 12 类）
        from ..security.tool_safety_guard import ToolSafetyGuard
        guard = ToolSafetyGuard(working_dir=cwd or os.getcwd())
        guard_result = guard.check_bash_command(command)
        if guard_result.should_block:
            return {
                "success": False,
                "error": f"命令被 ToolSafetyGuard 拦截：{guard_result.reason}",
            }

        # 路径安全校验：检测 cat/type/Get-Content 等读取命令是否在工作目录外
        # 防止 AI 绕过 file_read 工具的安全防护读取敏感文件
        read_block_reason = self._check_read_paths(command)
        if read_block_reason is not None:
            return {"success": False, "error": read_block_reason}

        last_result = None
        for attempt in range(self.MAX_RETRIES + 1):
            try:
                last_result = self._run_once(command, cwd, timeout)
            except subprocess.TimeoutExpired:
                # 超时不重试：第一次就附加建议，避免无意义的重试浪费时间
                last_result = {
                    "command": command,
                    "success": False,
                    "error": f"命令执行超时（{timeout}秒），已强制终止",
                }
                return self._enrich_timeout_error(command, timeout, last_result)
            except Exception as e:
                last_result = {"command": command, "success": False, "error": str(e)}

            # 如果是超时（不抛异常时由 _run_once 返回 success=False 且 error 含"超时"），
            # 不再重试，附加建议后直接返回。
            err_msg = (last_result or {}).get("error", "") or ""
            if not last_result.get("success") and "超时" in err_msg:
                return self._enrich_timeout_error(command, timeout, last_result)

            if last_result.get("success"):
                return last_result

            # 失败：除最后一次外都退避
            if attempt < self.MAX_RETRIES:
                time.sleep(self.BASE_BACKOFF * (2 ** attempt))

        return last_result
