# system_tools.py
import os
import asyncssh
from mcp.server.fastmcp import FastMCP, Context
from mcp.server.session import ServerSession
from lifespan import SharedAppContext  # 导入共享上下文的类型定义


def create_system_mcp() -> FastMCP:
    """创建一个包含系统操作工具的 MCP 实例。"""

    system_mcp = FastMCP(
        name="System Tools",
        instructions="用于在远程服务器上进行基本文件和目录操作的工具集。",
        # 将 MCP 实例的 HTTP 路径设置为根，以便挂载在 /system 下 [1]
        streamable_http_path="/"
    )

    @system_mcp.tool()
    async def list_files(ctx: Context[ServerSession, dict], path: str = ".") -> str:
        """
        列出远程服务器安全沙箱路径下的文件和目录。
        Args:
            path: 要查看的相对路径，默认为当前沙箱目录。
        """
        try:
            # 从上下文中获取共享的 SSH 连接和沙箱路径 [1]
            shared_ctx: SharedAppContext = ctx.request_context.lifespan_context["shared_context"]

            conn = shared_ctx.ssh_connection
            sandbox_path = shared_ctx.sandbox_path

            # 安全检查：防止目录穿越攻击
            if ".." in path.split('/'):
                return "错误: 不允许使用 '..' 访问上级目录。"

            # 安全地拼接路径
            target_path = os.path.join(sandbox_path, path)
            await ctx.info(f"正在列出安全路径: {target_path}")

            # 使用 conn.run() 执行 ls -l 命令
            command_to_run = f'ls -la {conn.escape(target_path)}'
            result = await conn.run(command_to_run, check=True)
            return result.stdout.strip()
        except asyncssh.ProcessError as e:
            error_message = f"命令执行失败: {e.stderr.strip()}"
            await ctx.error(error_message)
            return error_message
        except Exception as e:
            error_message = f"执行 ls 命令时发生未知错误: {e}"
            await ctx.error(error_message)
            return error_message

    @system_mcp.tool()
    async def read_file(ctx: Context[ServerSession, dict], file_path: str) -> str:
        """
        读取远程服务器安全沙箱内指定文件的内容。
        Args:
            file_path: 相对于沙箱目录的文件路径。
        """
        try:
            shared_ctx: SharedAppContext = ctx.request_context.lifespan_context["shared_context"]

            conn = shared_ctx.ssh_connection
            sandbox_path = shared_ctx.sandbox_path
            # 安全检查：防止目录穿越
            if ".." in file_path.split('/'):
                return "错误: 不允许使用 '..' 访问上级目录。"

            target_path = os.path.join(sandbox_path, file_path)
            await ctx.info(f"正在读取远程文件: {target_path}")

            async with conn.start_sftp_client() as sftp:
                async with sftp.open(target_path, 'r') as f:
                    content = await f.read()
            return content
        except FileNotFoundError:
            error_message = f"读取文件失败: 文件 '{file_path}' 不存在。"
            await ctx.error(error_message)
            return error_message
        except Exception as e:
            await ctx.error(f"读取文件时发生未知错误: {e}")
            return f"错误: {e}"

    @system_mcp.tool()
    async def write_file(ctx: Context[ServerSession, dict], file_path: str, content: str) -> str:
        """
        向远程服务器安全沙箱内的文件写入内容。如果文件已存在，则会覆盖它。
        Args:
            file_path: 相对于沙箱目录的文件路径。
            content: 要写入文件的文本内容。
        """
        try:
            shared_ctx: SharedAppContext = ctx.request_context.lifespan_context["shared_context"]

            conn = shared_ctx.ssh_connection
            sandbox_path = shared_ctx.sandbox_path

            # 安全检查
            if ".." in file_path.split('/'):
                return "错误: 不允许使用 '..' 访问上级目录。"

            target_path = os.path.join(sandbox_path, file_path)
            await ctx.info(f"正在向远程文件写入: {target_path}")

            async with conn.start_sftp_client() as sftp:
                async with sftp.open(target_path, 'w') as f:
                    await f.write(content)
            return f"内容已成功写入到沙箱文件 '{file_path}'。"
        except Exception as e:
            await ctx.error(f"写入文件时发生未知错误: {e}")
            return f"错误: {e}"

    return system_mcp