# 中文注释：
# 文件：echobot/agent/tools/shell.py
# 说明：智能体工具层，实现文件、Shell、网络与消息等工具调用能力。

"""Shell Execution Tool - Shell 执行工具
=================================================

此模块提供 Shell 命令执行工具。

功能说明：
- 执行任意 Shell 命令
- 返回命令输出（stdout 和 stderr）
- 支持超时控制
- 支持工作目录设置

安全注意事项：
- Agent 可以执行任意命令，请谨慎使用
- 建议在工作空间目录下执行
- 建议设置超时防止恶意命令

使用示例：
    result = await registry.execute("exec", {
        "command": "ls -la",
        "working_dir": "/tmp"
    })
"""

import asyncio
import os
from pathlib import Path
from typing import Any

from echobot.agent.tools.base import Tool


# ==================== Runtime Directory Registry ====================
# 运行时允许的目录（可通过 dir_manager 工具动态添加）
_runtime_allowed_dirs: set[Path] = set()


def get_runtime_allowed_dirs() -> list[Path]:
    """获取运行时允许的目录列表"""
    return sorted(_runtime_allowed_dirs)


def add_runtime_allowed_dir(path: str | Path) -> bool:
    """添加运行时允许的目录"""
    p = Path(path).expanduser().resolve()
    _runtime_allowed_dirs.add(p)
    return True


def remove_runtime_allowed_dir(path: str | Path) -> bool:
    """移除运行时允许的目录"""
    p = Path(path).expanduser().resolve()
    if p in _runtime_allowed_dirs:
        _runtime_allowed_dirs.discard(p)
        return True
    return False


class ExecTool(Tool):
    """
    Shell 命令执行工具
    
    允许 Agent 执行任意 Shell 命令并获取输出。
    
    工具元数据：
    - 名称：exec
    - 描述：Execute a shell command and return its output.
    - 参数：command（命令）、working_dir（可选工作目录）
    
    特性：
    - 超时控制（默认 60 秒）
    - 输出截断（超过 10000 字符）
    - 错误码返回
    
    使用示例：
        await registry.execute("exec", {"command": "python script.py"})
    
    注意事项：
    - 使用 asyncio.create_subprocess_shell 实现
    - 长时间运行的命令会被终止
    - 非常长的输出会被截断
    """
    
    def __init__(
        self,
        timeout: int = 60,
        working_dir: str | None = None,
        root_dir: str | Path | None = None,
        extra_allowed_dirs: list[Path] | None = None,
    ):
        """
        初始化执行工具
        
        Args:
            timeout: 超时时间（秒）
            working_dir: 默认工作目录
            root_dir: 可选工作区根目录。设置后会限制命令工作目录不可越界。
            extra_allowed_dirs: 额外的允许目录列表（用于扩展 root_dir 限制）
        """
        self.timeout = timeout
        self.working_dir = Path(working_dir).expanduser().resolve() if working_dir else None
        # root_dir 用于限制 working_dir 的边界，避免命令在工作区外执行。
        self.root_dir = Path(root_dir).expanduser().resolve() if root_dir else None
        # 额外的允许目录
        self.extra_allowed_dirs = extra_allowed_dirs or []
    
    @property
    def name(self) -> str:
        return "exec"
    
    @property
    def description(self) -> str:
        return "Execute a shell command and return its output. Use with caution."
    
    @property
    def parameters(self) -> dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "command": {
                    "type": "string",
                    "description": "The shell command to execute"
                },
                "working_dir": {
                    "type": "string",
                    "description": "Optional working directory for the command"
                }
            },
            "required": ["command"]
        }
    
    async def execute(self, command: str, working_dir: str | None = None, **kwargs: Any) -> str:
        """
        执行 Shell 命令
        
        Args:
            command: 要执行的命令
            working_dir: 可选工作目录
        
        Returns:
            str: 命令输出（stdout、stderr、退出码）
        """
        # 统一解析工作目录来源：调用参数 > 默认配置 > 当前进程目录。
        cwd_candidate = working_dir or self.working_dir or os.getcwd()
        cwd_path = Path(cwd_candidate).expanduser()

        # 若传入相对路径且配置了 root_dir，强制相对 root_dir 解析，保证路径稳定。
        if self.root_dir and not cwd_path.is_absolute():
            cwd_path = self.root_dir / cwd_path

        cwd = cwd_path.resolve()

        # 安全边界：禁止在 root_dir 外执行命令，降低误操作/越界风险。
        # 同时检查 extra_allowed_dirs 和运行时添加的目录
        allowed = False
        if self.root_dir and cwd.is_relative_to(self.root_dir):
            allowed = True
        elif self.extra_allowed_dirs:
            for extra_dir in self.extra_allowed_dirs:
                if cwd.is_relative_to(extra_dir.resolve()):
                    allowed = True
                    break
        # 检查运行时添加的目录
        if not allowed and _runtime_allowed_dirs:
            for extra_dir in _runtime_allowed_dirs:
                if cwd.is_relative_to(extra_dir):
                    allowed = True
                    break

        if not allowed:
            return f"Error: Working directory is outside allowed boundary: {cwd}"
        
        try:
            # 创建子进程
            process = await asyncio.create_subprocess_shell(
                command,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=str(cwd),
            )
            
            try:
                stdout, stderr = await asyncio.wait_for(
                    process.communicate(),
                    timeout=self.timeout
                )
            except asyncio.TimeoutError:
                process.kill()
                return f"Error: Command timed out after {self.timeout} seconds"
            
            output_parts = []
            
            # 收集 stdout
            if stdout:
                output_parts.append(stdout.decode("utf-8", errors="replace"))
            
            # 收集 stderr
            if stderr:
                stderr_text = stderr.decode("utf-8", errors="replace")
                if stderr_text.strip():
                    output_parts.append(f"STDERR:\n{stderr_text}")
            
            # 如果命令失败，添加退出码
            if process.returncode != 0:
                output_parts.append(f"\nExit code: {process.returncode}")
            
            result = "\n".join(output_parts) if output_parts else "(no output)"
            
            # 截断非常长的输出
            max_len = 10000
            if len(result) > max_len:
                result = result[:max_len] + f"\n... (truncated, {len(result) - max_len} more chars)"
            
            return result
            
        except Exception as e:
            return f"Error executing command: {str(e)}"
