import os
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("tree")

def generate_tree(dir_path: str, prefix: str = "", max_depth: int = 3, current_depth: int = 1) -> list:
    if current_depth > max_depth:
        return [prefix + "└── ..."]
        
    try:
        items = sorted(os.listdir(dir_path))
    except PermissionError:
        return [prefix + "└── [Permission Denied]"]
        
    lines = []
    # 分离文件夹和文件，并排序
    dirs = [i for i in items if os.path.isdir(os.path.join(dir_path, i))]
    files = [i for i in items if not os.path.isdir(os.path.join(dir_path, i))]
    sorted_items = dirs + files

    for idx, item in enumerate(sorted_items):
        path = os.path.join(dir_path, item)
        is_last = (idx == len(sorted_items) - 1)
        connector = "└── " if is_last else "├── "
        lines.append(prefix + connector + item)
        
        if os.path.isdir(path):
            extension = "    " if is_last else "│   "
            lines.extend(generate_tree(path, prefix + extension, max_depth, current_depth + 1))
            
    return lines

@mcp.tool()
def show_tree(directory: str, max_depth: int = 2) -> str:
    """以树状结构展示目录内容。"""
    if not os.path.isdir(directory):
        return f"目录不存在或不是一个目录: {directory}"
        
    lines = [directory]
    lines.extend(generate_tree(directory, "", max_depth=max_depth))
    return "\n".join(lines)

def main():
    mcp.run()

if __name__ == "__main__":
    main()
