#!/usr/bin/env node
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { CallToolRequestSchema, ListToolsRequestSchema } from "@modelcontextprotocol/sdk/types.js";
import { execSync } from "child_process";

function runS3(args) {
  const cmd = `aws s3 ${args}`;
  return execSync(cmd, { encoding: "utf8", maxBuffer: 10*1024*1024 });
}

const tools = [
  {
    name: "s3_list_buckets",
    description: "列出所有S3存储桶",
    inputSchema: { type: "object", properties: {} }
  },
  {
    name: "s3_list_files",
    description: "列出S3存储桶中的文件",
    inputSchema: {
      type: "object",
      properties: {
        bucket: { type: "string", description: "桶名称" },
        prefix: { type: "string", description: "路径前缀", default: "" },
        maxItems: { type: "number", description: "最大数量", default: 50 }
      },
      required: ["bucket"]
    }
  },
  {
    name: "s3_read_file",
    description: "读取S3中文本文件的内容",
    inputSchema: {
      type: "object",
      properties: {
        bucket: { type: "string", description: "桶名称" },
        key: { type: "string", description: "S3中的文件路径" },
        maxLines: { type: "number", description: "最大行数", default: 100 }
      },
      required: ["bucket", "key"]
    }
  },
  {
    name: "s3_upload_file",
    description: "上传文件到S3",
    inputSchema: {
      type: "object",
      properties: {
        localPath: { type: "string", description: "本地文件路径" },
        bucket: { type: "string", description: "目标桶" },
        key: { type: "string", description: "S3中的路径/文件名" }
      },
      required: ["localPath", "bucket", "key"]
    }
  },
  {
    name: "s3_download_file",
    description: "从S3下载文件",
    inputSchema: {
      type: "object",
      properties: {
        bucket: { type: "string", description: "桶名称" },
        key: { type: "string", description: "S3中的文件路径" },
        localPath: { type: "string", description: "本地保存路径" }
      },
      required: ["bucket", "key", "localPath"]
    }
  }
];

const server = new Server({ name: "s3-mcp", version: "1.0.1" }, {
  capabilities: { tools: {} }
});

server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools }));

server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;
  
  try {
    switch (name) {
      case "s3_list_buckets": {
        const result = runS3("ls").trim();
        return { content: [{ type: "text", text: result || "(空)" }] };
      }
      case "s3_list_files": {
        const prefix = args.prefix ? args.prefix + "/" : "";
        const max = args.maxItems || 50;
        const all = runS3(`ls s3://${args.bucket}/${prefix} --recursive`);
        const lines = all.trim().split('\n').slice(0, max);
        return { content: [{ type: "text", text: lines.join('\n') || "(空)" }] };
      }
      case "s3_read_file": {
        const tmpFile = `/tmp/s3_read_${Date.now()}.txt`;
        runS3(`cp s3://${args.bucket}/${args.key} ${tmpFile}`);
        const content = execSync(`head -n ${args.maxLines || 100} ${tmpFile}`, { encoding: "utf8" });
        execSync(`rm -f ${tmpFile}`);
        return { content: [{ type: "text", text: content || "(空)" }] };
      }
      case "s3_upload_file": {
        runS3(`cp "${args.localPath}" s3://${args.bucket}/${args.key}`);
        return { content: [{ type: "text", text: `✅ 已上传到 s3://${args.bucket}/${args.key}` }] };
      }
      case "s3_download_file": {
        runS3(`cp s3://${args.bucket}/${args.key} "${args.localPath}"`);
        return { content: [{ type: "text", text: `✅ 已下载到 ${args.localPath}` }] };
      }
      default:
        throw new Error(`未知工具: ${name}`);
    }
  } catch (e) {
    return { content: [{ type: "text", text: `❌ 错误: ${e.message}` }], isError: true };
  }
});

const transport = new StdioServerTransport();
server.connect(transport);
