import http from "node:http";
import fs from "node:fs";
import path from "node:path";
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import {
  ListToolsRequestSchema,
  CallToolRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { vectorTools, handleVectorTool } from "./tools/vectorTool.js";
import { curateTools, handleCurateTool } from "./tools/curateTool.js";
import { routerTools, handleRouterTool } from "./tools/routerTool.js";

const PORT = Number(process.env.PORT || 3000);
const HOST = process.env.HOST || "0.0.0.0";
const DATA_DIR = process.env.DATA_DIR || "/data";
const CONFIG_PATH = path.join(DATA_DIR, "config.json");

function loadConfig() {
  try {
    if (!fs.existsSync(CONFIG_PATH)) {
      console.error(`Config file not found at ${CONFIG_PATH}`);
      process.exit(1);
    }
    const config = JSON.parse(fs.readFileSync(CONFIG_PATH, "utf8"));
    return config;
  } catch (error) {
    console.error(`Failed to load config from ${CONFIG_PATH}:`, error.message);
    process.exit(1);
  }
}

const config = loadConfig();

function createMcpServer() {
  const server = new Server(
    {
      name: "mcponly-server",
      version: "1.0.0",
    },
    {
      capabilities: {
        tools: {},
      },
    },
  );

  server.setRequestHandler(ListToolsRequestSchema, async () => ({
    tools: [...routerTools, ...curateTools, ...vectorTools],
  }));

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

    const routerResult = await handleRouterTool(name, args, config);
    if (routerResult) return routerResult;

    const curateResult = await handleCurateTool(name, args, config);
    if (curateResult) return curateResult;

    const vectorResult = await handleVectorTool(name, args);
    if (vectorResult) return vectorResult;

    return {
      content: [{ type: "text", text: `Unknown tool: ${name}` }],
      isError: true,
    };
  });

  return server;
}

function sendJson(res, statusCode, payload) {
  res.writeHead(statusCode, { "Content-Type": "application/json" });
  res.end(JSON.stringify(payload));
}

function readBody(req) {
  return new Promise((resolve, reject) => {
    const chunks = [];

    req.on("data", (chunk) => chunks.push(chunk));
    req.on("end", () => {
      if (chunks.length === 0) {
        resolve(undefined);
        return;
      }

      const text = Buffer.concat(chunks).toString("utf8");

      try {
        resolve(JSON.parse(text));
      } catch (error) {
        reject(new Error("Invalid JSON body"));
      }
    });
    req.on("error", reject);
  });
}

async function main() {
  const transports = new Map();

  const httpServer = http.createServer(async (req, res) => {
    if (!req.url) {
      sendJson(res, 400, {
        jsonrpc: "2.0",
        error: { code: -32600, message: "Missing request URL" },
        id: null,
      });
      return;
    }

    const hostHeader = req.headers.host || `${HOST}:${PORT}`;
    const url = new URL(req.url, `http://${hostHeader}`);

    if (req.method === "GET" && url.pathname === "/health") {
      sendJson(res, 200, { status: "ok" });
      return;
    }

    if (req.method === "GET" && url.pathname === "/sse") {
      const transport = new SSEServerTransport("/messages", res);
      transports.set(transport.sessionId, transport);

      transport.onclose = () => {
        transports.delete(transport.sessionId);
      };

      res.on("close", () => {
        transports.delete(transport.sessionId);
      });

      await createMcpServer().connect(transport);
      return;
    }

    if (req.method === "POST" && url.pathname === "/messages") {
      const sessionId = url.searchParams.get("sessionId");

      if (!sessionId) {
        sendJson(res, 400, {
          jsonrpc: "2.0",
          error: { code: -32600, message: "Missing sessionId query parameter" },
          id: null,
        });
        return;
      }

      const transport = transports.get(sessionId);

      if (!transport) {
        sendJson(res, 404, {
          jsonrpc: "2.0",
          error: { code: -32000, message: "Session not found" },
          id: null,
        });
        return;
      }

      const body = await readBody(req);
      await transport.handlePostMessage(req, res, body);
      return;
    }

    sendJson(res, 404, {
      jsonrpc: "2.0",
      error: { code: -32601, message: "Not found" },
      id: null,
    });
  });

  httpServer.listen(PORT, HOST, () => {
    console.log(`MCP SSE server listening on http://${HOST}:${PORT}`);
    console.log(`SSE endpoint: http://${HOST}:${PORT}/sse`);
    console.log(`POST endpoint: http://${HOST}:${PORT}/messages`);
  });

  process.on("SIGINT", () => {
    httpServer.close(() => process.exit(0));
  });
}

main().catch((error) => {
  console.error("MCP server failed to start:", error);
  process.exit(1);
});