#!/usr/bin/env node
import { promises as fs } from "node:fs";
import path from "node:path";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
  CallToolRequestSchema,
  ListResourcesRequestSchema,
  ListToolsRequestSchema,
  ReadResourceRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";

const imageCache = new Map();

const TOOLS = [
  {
    name: "list_local_images",
    description: "List image files inside a local directory for debugging.",
    inputSchema: {
      type: "object",
      properties: {
        dirPath: {
          type: "string",
          description: "Absolute path of the directory to scan",
        },
      },
      required: ["dirPath"],
    },
  },
  {
    name: "load_local_image",
    description: "Load a local image file and return it as MCP image content so the agent can inspect it.",
    inputSchema: {
      type: "object",
      properties: {
        filePath: {
          type: "string",
          description: "Absolute path of the image file to load",
        },
      },
      required: ["filePath"],
    },
  },
];

function inferMimeType(filePath) {
  const ext = path.extname(filePath).toLowerCase();
  switch (ext) {
    case ".png":
      return "image/png";
    case ".jpg":
    case ".jpeg":
      return "image/jpeg";
    case ".webp":
      return "image/webp";
    case ".gif":
      return "image/gif";
    case ".bmp":
      return "image/bmp";
    case ".svg":
      return "image/svg+xml";
    default:
      throw new Error(`Unsupported image extension: ${ext || "<none>"}`);
  }
}

async function listLocalImages(dirPath) {
  const entries = await fs.readdir(dirPath, { withFileTypes: true });
  const files = entries
    .filter((entry) => entry.isFile())
    .map((entry) => path.join(dirPath, entry.name))
    .filter((filePath) => {
      try {
        inferMimeType(filePath);
        return true;
      } catch {
        return false;
      }
    })
    .sort();

  return {
    content: [
      {
        type: "text",
        text: files.length
          ? files.join("\n")
          : `No supported images found in ${dirPath}`,
      },
    ],
    isError: false,
  };
}

async function loadLocalImage(filePath) {
  const resolvedPath = path.resolve(filePath);
  const mimeType = inferMimeType(resolvedPath);
  const bytes = await fs.readFile(resolvedPath);
  const base64 = bytes.toString("base64");

  imageCache.set(resolvedPath, { mimeType, base64 });
  server.notification({ method: "notifications/resources/list_changed" });

  return {
    content: [
      {
        type: "text",
        text: `Loaded image ${resolvedPath}`,
      },
      {
        type: "image",
        data: base64,
        mimeType,
      },
    ],
    isError: false,
  };
}

const server = new Server(
  {
    name: "local-image-debug-server",
    version: "0.1.0",
  },
  {
    capabilities: {
      tools: {},
      resources: {},
    },
  }
);

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

server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const name = request.params.name;
  const args = request.params.arguments ?? {};

  try {
    switch (name) {
      case "list_local_images":
        return await listLocalImages(args.dirPath);
      case "load_local_image":
        return await loadLocalImage(args.filePath);
      default:
        return {
          content: [{ type: "text", text: `Unknown tool: ${name}` }],
          isError: true,
        };
    }
  } catch (error) {
    return {
      content: [
        {
          type: "text",
          text: `Tool failed: ${error instanceof Error ? error.message : String(error)}`,
        },
      ],
      isError: true,
    };
  }
});

server.setRequestHandler(ListResourcesRequestSchema, async () => ({
  resources: Array.from(imageCache.entries()).map(([filePath, image]) => ({
    uri: `debug-image://${encodeURIComponent(filePath)}`,
    name: path.basename(filePath),
    mimeType: image.mimeType,
  })),
}));

server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
  const uri = request.params.uri.toString();
  if (!uri.startsWith("debug-image://")) {
    throw new Error(`Unsupported resource URI: ${uri}`);
  }

  const filePath = decodeURIComponent(uri.replace("debug-image://", ""));
  const cached = imageCache.get(filePath);
  if (!cached) {
    throw new Error(`Resource not found: ${filePath}`);
  }

  return {
    contents: [
      {
        uri,
        mimeType: cached.mimeType,
        blob: cached.base64,
      },
    ],
  };
});

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