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 { readFileSync } from "node:fs";
import { basename } from "node:path";

// ============================================================================
// Configuration (mirrors config.go defaults)
// ============================================================================

const DEFAULT_API_KEY = "sk-68FL9xXRrJhxcY5TCGxoaFlQ94oCBioIJhyGfeDMCkCvA0SV";
const DEFAULT_MODEL = "agnes-2.0-flash";
const DEFAULT_BASE_URL = "https://apihub.agnes-ai.com/v1/chat/completions";

const API_KEY = process.env.AGNES_API_KEY || DEFAULT_API_KEY;
const MODEL = process.env.VISION_MODEL || DEFAULT_MODEL;
const BASE_URL = process.env.VISION_BASE_URL || DEFAULT_BASE_URL;

// ============================================================================
// MIME type detection
// ============================================================================

function detectMimeType(path) {
  const ext = path.split(".").pop().toLowerCase();
  const map = {
    png: "image/png",
    jpg: "image/jpeg",
    jpeg: "image/jpeg",
    gif: "image/gif",
    webp: "image/webp",
    bmp: "image/bmp",
    tiff: "image/tiff",
    tif: "image/tiff",
  };
  return map[ext] || "image/png";
}

// ============================================================================
// Vision API client
// ============================================================================

async function analyzeImage(imagePath, prompt) {
  const imgData = readFileSync(imagePath);
  const mimeType = detectMimeType(imagePath);
  const base64Img = Buffer.from(imgData).toString("base64");
  const dataURL = `data:${mimeType};base64,${base64Img}`;

  if (!prompt) {
    prompt = "请详细描述这张图片的内容。如果图片中有文字，请完整识别出来。请用中文回答。";
  }

  const body = {
    model: MODEL,
    messages: [
      {
        role: "user",
        content: [
          { type: "image_url", image_url: { url: dataURL } },
          { type: "text", text: prompt },
        ],
      },
    ],
    chat_template_kwargs: { enable_thinking: true },
  };

  const resp = await fetch(BASE_URL, {
    method: "POST",
    headers: {
      "Authorization": `Bearer ${API_KEY}`,
      "Content-Type": "application/json",
    },
    body: JSON.stringify(body),
    signal: AbortSignal.timeout(120_000),
  });

  if (!resp.ok) {
    const text = await resp.text();
    throw new Error(`API error (HTTP ${resp.status}): ${text}`);
  }

  const data = await resp.json();
  if (data.error) {
    throw new Error(`API error: [${data.error.code}] ${data.error.message}`);
  }

  const content = data.choices?.[0]?.message?.content;
  if (!content) {
    throw new Error("No content in API response");
  }

  return content;
}

// ============================================================================
// MCP Server
// ============================================================================

const server = new Server(
  {
    name: "vision-tool",
    version: "1.0.0",
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

// Register tool list
server.setRequestHandler(ListToolsRequestSchema, async () => ({
  tools: [
    {
      name: "analyze_image",
      description:
        "分析图片内容。使用 Agnes AI 视觉模型识别图片中的对象、场景、文字等。" +
        "支持 PNG、JPG、GIF、WebP、BMP、TIFF 格式。可以用自定义 prompt 指定分析重点。",
      inputSchema: {
        type: "object",
        properties: {
          image_path: {
            type: "string",
            description: "图片文件的绝对路径（例如 C:\\Users\\xxx\\photo.png）",
          },
          prompt: {
            type: "string",
            description:
              "自定义分析提示词。默认为详细描述图片内容并识别所有文字。例如：" +
              '"这张图里有什么文字？" 或 "描述图片中的UI界面布局"。',
          },
        },
        required: ["image_path"],
      },
    },
  ],
}));

// Register tool handler
server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  if (name !== "analyze_image") {
    throw new Error(`Unknown tool: ${name}`);
  }

  const imagePath = args.image_path;
  const prompt = args.prompt || "";

  try {
    const result = await analyzeImage(imagePath, prompt);
    return {
      content: [{ type: "text", text: result }],
    };
  } catch (err) {
    return {
      content: [{ type: "text", text: `Error: ${err.message}` }],
      isError: true,
    };
  }
});

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