#!/usr/bin/env bun

import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
import { z } from "zod";
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
  Tool,
} from "@modelcontextprotocol/sdk/types.js";
import { createServer, IncomingMessage, ServerResponse } from "node:http";
import { URL } from "node:url";
import { searchDatasets, fetchDatasetAggregate, textContent } from "./utils.ts";

const SearchSchema = z.object({
  query: z
    .string()
    .describe(
      "Search query for datasets (e.g., 'nlp sentiment analysis', 'computer vision', 'question answering')",
    ),
});

const FetchSchema = z.object({
  id: z
    .string()
    .describe(
      "Unique dataset identifier (e.g., 'squad', 'imdb', 'cfahlgren1/hub-stats')",
    ),
});

const TOOL_DEFS: Tool[] = [
  {
    name: "search",
    description: "Search for datasets on Hugging Face Hub",
    inputSchema: {
      type: "object",
      properties: {
        query: {
          type: "string",
          description:
            "Search query for datasets (e.g., 'nlp sentiment analysis', 'computer vision', 'question answering')",
        },
      },
      required: ["query"],
    },
  },
  {
    name: "fetch",
    description:
      "Retrieve full information and sample data for a specific dataset",
    inputSchema: {
      type: "object",
      properties: {
        id: {
          type: "string",
          description:
            "Unique dataset identifier (e.g., 'squad', 'imdb', 'cfahlgren1/hub-stats')",
        },
      },
      required: ["id"],
    },
  },
];

const server = new Server(
  { name: "dataset-viewer-mcp", version: "1.0.0" },
  { capabilities: { tools: {} } }
);

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

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

  try {
    switch (name) {
      case "search": {
        const validatedArgs = SearchSchema.parse(args);
        const { query } = validatedArgs;

        const results = await searchDatasets(query);
        return textContent({ results });
      }

      case "fetch": {
        const validatedArgs = FetchSchema.parse(args);
        const { id: datasetId } = validatedArgs;

        const fetchResult = await fetchDatasetAggregate(datasetId);
        return textContent(fetchResult);
      }
      
      default:
        throw new Error(`Unknown tool: ${name}`);
    }
  } catch (error) {
    return {
      content: [{ type: "text", text: `Error: ${error instanceof Error ? error.message : String(error)}` }],
      isError: true,
    };
  }
});

async function main() {
  const port = Number(process.env.PORT) || 3000;
  const transports = new Map<string, SSEServerTransport>();
  
  const httpServer = createServer(async (req: IncomingMessage, res: ServerResponse) => {
    res.setHeader("Access-Control-Allow-Origin", "*");
    res.setHeader("Access-Control-Allow-Headers", "Content-Type");
    res.setHeader("Access-Control-Allow-Methods", "GET, POST, OPTIONS");
    
    const url = new URL(req.url || "", `http://localhost:${port}`);
    
    if (req.method === "OPTIONS") {
      res.writeHead(204);
      res.end();
      return;
    }
    
    if (url.pathname === "/sse") {
      if (req.method === "GET") {
        const transport = new SSEServerTransport("/sse", res);
        transport.onclose = () => transports.delete(transport.sessionId);
        await server.connect(transport);
        transports.set(transport.sessionId, transport);
        return;
      }
      
      if (req.method === "POST") {
        const sessionId = url.searchParams.get("sessionId");
        if (!sessionId) {
          res.writeHead(400).end("Missing sessionId");
          return;
        }
        
        const transport = transports.get(sessionId);
        if (!transport) {
          res.writeHead(404).end("Session not found");
          return;
        }
        
        try {
          await transport.handlePostMessage(req, res);
        } catch (error) {
          console.error("Error handling POST message:", error);
          if (!res.headersSent) {
            res.writeHead(500).end("Internal Server Error");
          }
        }
        return;
      }
    }
    
    res.writeHead(200, { "Content-Type": "application/json" });
    res.end(JSON.stringify({
      name: "Dataset Viewer MCP Server",
      version: "1.0.0",
      transport: "SSE",
      endpoint: "/sse",
      usage: "GET /sse for SSE connection, POST /sse?sessionId=<id> for messages"
    }));
  });
  
  httpServer.listen(port, () => {
    console.log(`Dataset Viewer MCP server running on http://localhost:${port}/sse`);
  });
}

main().catch((error) => {
  console.error("Server error:", error);
  process.exit(1);
});