import { randomUUID } from "node:crypto"
import express from "express"
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js"
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js"
import { isInitializeRequest } from "@modelcontextprotocol/sdk/types.js"
import { z } from "zod"

const PORT = Number(process.env.MCP_PORT ?? 3100)
const API_BASE = (process.env.ULTID_API_URL ?? "http://localhost:8080/api/v1").replace(/\/$/, "")

const TOOL_GROUPS: Record<string, string[]> = {
  mail: [
    "mail_search",
    "mail_read",
    "mail_list",
    "mail_send",
    "mail_update_labels",
    "mail_update_flags",
    "mail_delete",
  ],
  drive: [
    "drive_list",
    "drive_file_info",
    "drive_create_file",
    "drive_create_folder",
    "drive_delete",
    "drive_move",
    "drive_rename",
    "drive_share_create",
    "drive_share_update",
    "drive_share_delete",
  ],
  contacts: [
    "contacts_search",
    "contacts_list_books",
    "contacts_list",
    "contacts_get",
    "contacts_create",
    "contacts_update",
    "contacts_delete",
  ],
  search: ["suite_search"],
  docs: ["docs_read", "docs_save"],
  agenda: [
    "calendar_list",
    "calendar_create",
    "calendar_update",
    "calendar_delete",
    "calendar_list_events",
    "calendar_freebusy",
    "calendar_create_event",
    "calendar_update_event",
    "calendar_delete_event",
    "calendar_respond_invitation",
    "calendar_create_meet_link",
  ],
  web_search: ["web_search"],
}

const ALL_TOOL_NAMES = [...new Set(Object.values(TOOL_GROUPS).flat())]

function parseEnabledTools(raw: string | undefined): Set<string> | null {
  const groups = String(raw ?? "")
    .split(",")
    .map((part) => part.trim().toLowerCase())
    .filter(Boolean)
  if (groups.length === 0) return null
  const enabled = new Set<string>()
  for (const group of groups) {
    for (const tool of TOOL_GROUPS[group] ?? []) {
      enabled.add(tool)
    }
  }
  return enabled.size > 0 ? enabled : null
}

function encodePath(path: string): string {
  return path
    .replace(/^\/+/, "")
    .split("/")
    .filter(Boolean)
    .map((seg) => encodeURIComponent(seg))
    .join("/")
}

function toolText(data: unknown) {
  return {
    content: [
      {
        type: "text" as const,
        text: typeof data === "string" ? data : JSON.stringify(data, null, 2),
      },
    ],
  }
}

/** Ultid list APIs use page + page_size, not limit + offset. */
function listPaginationParams(limit?: number, offset?: number): URLSearchParams {
  const qs = new URLSearchParams()
  const pageSize =
    limit != null && Number.isFinite(limit) && limit > 0
      ? Math.min(Math.trunc(limit), 500)
      : undefined
  if (pageSize != null) {
    qs.set("page_size", String(pageSize))
    const off = offset != null && Number.isFinite(offset) && offset > 0 ? Math.trunc(offset) : 0
    qs.set("page", String(Math.floor(off / pageSize) + 1))
  } else if (offset != null && Number.isFinite(offset) && offset > 0) {
    qs.set("page_size", "50")
    qs.set("page", String(Math.floor(Math.trunc(offset) / 50) + 1))
  }
  return qs
}

function withListPagination(
  base: URLSearchParams,
  limit?: number,
  offset?: number
): URLSearchParams {
  const qs = new URLSearchParams(base)
  for (const [key, value] of listPaginationParams(limit, offset)) {
    qs.set(key, value)
  }
  return qs
}

function applyHiddenMailboxParams(
  qs: URLSearchParams,
  include_spam?: boolean,
  include_trash?: boolean
) {
  if (include_spam) qs.set("include_spam", "true")
  if (include_trash) qs.set("include_trash", "true")
}

async function ultiFetch(
  token: string,
  path: string,
  init?: RequestInit
): Promise<unknown> {
  const headers = new Headers(init?.headers)
  headers.set("Accept", "application/json")
  if (token) headers.set("Authorization", `Bearer ${token}`)
  const res = await fetch(`${API_BASE}${path}`, { ...init, headers })
  const text = await res.text()
  if (!res.ok) {
    throw new Error(`ulti ${path} failed (${res.status}): ${text.slice(0, 500)}`)
  }
  if (!text) return { ok: true, status: res.status }
  try {
    return JSON.parse(text)
  } catch {
    return text
  }
}

function resolveToken(req: express.Request): string {
  return (
    String(req.headers["x-ulti-token"] ?? req.headers.authorization ?? "").replace(
      /^Bearer\s+/i,
      ""
    ) || ""
  )
}

function createServer(getToken: () => string, enabledTools: Set<string> | null) {
  const server = new McpServer({
    name: "ultimail-mcp",
    version: "0.1.0",
  })

  const allow = (toolName: string) => !enabledTools || enabledTools.has(toolName)

  if (allow("mail_search")) {
    server.tool(
      "mail_search",
      "Search mail messages (spam and trash excluded by default)",
      {
        query: z.string(),
        account_id: z.string().optional(),
        limit: z.number().optional(),
        offset: z.number().optional(),
        include_spam: z.boolean().optional(),
        include_trash: z.boolean().optional(),
      },
      async ({ query, account_id, limit, offset, include_spam, include_trash }) => {
        const qs = withListPagination(new URLSearchParams({ q: query }), limit, offset)
        if (account_id) qs.set("account_id", account_id)
        applyHiddenMailboxParams(qs, include_spam, include_trash)
        const data = await ultiFetch(getToken(), `/mail/search?${qs}`)
        return toolText(data)
      }
    )
  }

  if (allow("mail_read")) {
    server.tool(
      "mail_read",
      "Read a mail message by id",
      { message_id: z.string(), account_id: z.string().optional() },
      async ({ message_id, account_id }) => {
        const qs = account_id ? `?account_id=${encodeURIComponent(account_id)}` : ""
        const data = await ultiFetch(getToken(), `/mail/messages/${message_id}${qs}`)
        return toolText(data)
      }
    )
  }

  if (allow("mail_list")) {
    server.tool(
      "mail_list",
      "List mail messages in a folder (spam and trash excluded by default unless folder=spam/trash)",
      {
        folder: z.string().optional(),
        account_id: z.string().optional(),
        limit: z.number().optional(),
        offset: z.number().optional(),
        include_spam: z.boolean().optional(),
        include_trash: z.boolean().optional(),
      },
      async ({ folder, account_id, limit, offset, include_spam, include_trash }) => {
        let qs = new URLSearchParams()
        if (folder) qs.set("folder", folder)
        if (account_id) qs.set("account_id", account_id)
        qs = withListPagination(qs, limit, offset)
        applyHiddenMailboxParams(qs, include_spam, include_trash)
        const suffix = qs.size > 0 ? `?${qs}` : ""
        const data = await ultiFetch(getToken(), `/mail/messages${suffix}`)
        return toolText(data)
      }
    )
  }

  if (allow("mail_send")) {
    server.tool(
      "mail_send",
      "Send a mail message",
      {
        account_id: z.string(),
        to: z.array(z.string()),
        cc: z.array(z.string()).optional(),
        bcc: z.array(z.string()).optional(),
        subject: z.string(),
        body_text: z.string().optional(),
        body_html: z.string().optional(),
        in_reply_to: z.string().optional(),
        reply_to_message_id: z.string().optional(),
        schedule_at: z.string().optional(),
      },
      async (body) => {
        const data = await ultiFetch(getToken(), "/mail/send", {
          method: "POST",
          headers: {
            "Content-Type": "application/json",
            "Idempotency-Key": randomUUID(),
          },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("mail_update_labels")) {
    server.tool(
      "mail_update_labels",
      "Update labels on a mail message",
      { message_id: z.string(), labels: z.array(z.string()) },
      async ({ message_id, labels }) => {
        const data = await ultiFetch(getToken(), `/mail/messages/${message_id}/labels`, {
          method: "PUT",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ labels }),
        })
        return toolText(data)
      }
    )
  }

  if (allow("mail_update_flags")) {
    server.tool(
      "mail_update_flags",
      "Update flags on a mail message",
      { message_id: z.string(), flags: z.array(z.string()) },
      async ({ message_id, flags }) => {
        const data = await ultiFetch(getToken(), `/mail/messages/${message_id}/flags`, {
          method: "PUT",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ flags }),
        })
        return toolText(data)
      }
    )
  }

  if (allow("mail_delete")) {
    server.tool(
      "mail_delete",
      "Delete a mail message",
      { message_id: z.string() },
      async ({ message_id }) => {
        const data = await ultiFetch(getToken(), `/mail/messages/${message_id}`, {
          method: "DELETE",
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_list")) {
    server.tool(
      "drive_list",
      "List drive files in a folder",
      {
        path: z.string().optional(),
        limit: z.number().optional(),
        offset: z.number().optional(),
      },
      async ({ path, limit, offset }) => {
        const encoded = encodePath(path ?? "")
        const qs = listPaginationParams(limit, offset)
        const suffix = qs.size > 0 ? `?${qs}` : ""
        const data = await ultiFetch(getToken(), `/drive/files/${encoded}${suffix}`)
        return toolText(data)
      }
    )
  }

  if (allow("drive_file_info")) {
    server.tool(
      "drive_file_info",
      "Get drive file or folder metadata",
      { path: z.string() },
      async ({ path }) => {
        const encoded = encodePath(path)
        const data = await ultiFetch(getToken(), `/drive/files/info/${encoded}`)
        return toolText(data)
      }
    )
  }

  if (allow("drive_create_file")) {
    server.tool(
      "drive_create_file",
      "Create a new office file (document, spreadsheet, presentation, drawing)",
      {
        parent_path: z.string().optional(),
        name: z.string(),
        kind: z.enum(["document", "spreadsheet", "presentation", "drawing"]),
      },
      async ({ parent_path, name, kind }) => {
        const data = await ultiFetch(getToken(), "/drive/files/new", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ parent_path, name, kind }),
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_create_folder")) {
    server.tool(
      "drive_create_folder",
      "Create a drive folder",
      { path: z.string() },
      async ({ path }) => {
        const encoded = encodePath(path)
        const data = await ultiFetch(getToken(), `/drive/folders/${encoded}`, {
          method: "POST",
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_delete")) {
    server.tool(
      "drive_delete",
      "Delete a drive file or folder",
      { path: z.string() },
      async ({ path }) => {
        const encoded = encodePath(path)
        const data = await ultiFetch(getToken(), `/drive/files/${encoded}`, {
          method: "DELETE",
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_move")) {
    server.tool(
      "drive_move",
      "Move a drive file or folder",
      {
        source: z.string(),
        destination: z.string(),
        source_root: z.string().optional(),
        source_root_id: z.string().optional(),
        destination_root: z.string().optional(),
        destination_root_id: z.string().optional(),
      },
      async (body) => {
        const data = await ultiFetch(getToken(), "/drive/move", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_rename")) {
    server.tool(
      "drive_rename",
      "Rename a drive file or folder",
      {
        path: z.string(),
        new_name: z.string(),
        root: z.string().optional(),
        root_id: z.string().optional(),
      },
      async (body) => {
        const data = await ultiFetch(getToken(), "/drive/rename", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_share_create")) {
    server.tool(
      "drive_share_create",
      "Create a drive share link or internal share",
      {
        path: z.string(),
        share_type: z.number().optional(),
        permissions: z.number().optional(),
        role: z.enum(["owner", "editor", "viewer"]).optional(),
        mode: z.enum(["public", "internal", "contact"]).optional(),
        share_with: z.string().optional(),
        note: z.string().optional(),
        root: z.string().optional(),
        root_id: z.string().optional(),
      },
      async (body) => {
        const data = await ultiFetch(getToken(), "/drive/shares", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_share_update")) {
    server.tool(
      "drive_share_update",
      "Update an existing drive share",
      {
        share_id: z.string(),
        permissions: z.number().optional(),
        role: z.enum(["owner", "editor", "viewer"]).optional(),
        expire_date: z.string().optional(),
        password: z.string().optional(),
      },
      async ({ share_id, ...body }) => {
        const data = await ultiFetch(getToken(), `/drive/shares/${share_id}`, {
          method: "PUT",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("drive_share_delete")) {
    server.tool(
      "drive_share_delete",
      "Delete a drive share",
      { share_id: z.string() },
      async ({ share_id }) => {
        const data = await ultiFetch(getToken(), `/drive/shares/${share_id}`, {
          method: "DELETE",
        })
        return toolText(data)
      }
    )
  }

  if (allow("contacts_search")) {
    server.tool(
      "contacts_search",
      "Search contacts",
      {
        query: z.string(),
        book_id: z.string().optional(),
        limit: z.number().optional(),
        offset: z.number().optional(),
      },
      async ({ query, book_id, limit, offset }) => {
        const qs = withListPagination(new URLSearchParams({ q: query }), limit, offset)
        if (book_id) qs.set("book_id", book_id)
        const data = await ultiFetch(getToken(), `/contacts/search?${qs}`)
        return toolText(data)
      }
    )
  }

  if (allow("contacts_list_books")) {
    server.tool(
      "contacts_list_books",
      "List address books",
      {},
      async () => {
        const data = await ultiFetch(getToken(), "/contacts/books")
        return toolText(data)
      }
    )
  }

  if (allow("contacts_list")) {
    server.tool(
      "contacts_list",
      "List contacts in an address book",
      {
        book_id: z.string(),
        limit: z.number().optional(),
        offset: z.number().optional(),
      },
      async ({ book_id, limit, offset }) => {
        const qs = listPaginationParams(limit, offset)
        const suffix = qs.size > 0 ? `?${qs}` : ""
        const data = await ultiFetch(getToken(), `/contacts/books/${book_id}${suffix}`)
        return toolText(data)
      }
    )
  }

  if (allow("contacts_get")) {
    server.tool(
      "contacts_get",
      "Get a contact by path (book_id/contact_uid)",
      { path: z.string() },
      async ({ path }) => {
        const encoded = encodePath(path)
        const data = await ultiFetch(getToken(), `/contacts/${encoded}`)
        return toolText(data)
      }
    )
  }

  if (allow("contacts_create")) {
    server.tool(
      "contacts_create",
      "Create a contact in an address book",
      {
        book_id: z.string(),
        contact: z.record(z.string(), z.unknown()),
      },
      async ({ book_id, contact }) => {
        const data = await ultiFetch(getToken(), `/contacts/books/${book_id}`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(contact),
        })
        return toolText(data)
      }
    )
  }

  if (allow("contacts_update")) {
    server.tool(
      "contacts_update",
      "Update a contact (requires etag from get)",
      {
        path: z.string(),
        if_match: z.string(),
        contact: z.record(z.string(), z.unknown()),
      },
      async ({ path, if_match, contact }) => {
        const encoded = encodePath(path)
        const data = await ultiFetch(getToken(), `/contacts/${encoded}`, {
          method: "PUT",
          headers: {
            "Content-Type": "application/json",
            "If-Match": if_match,
          },
          body: JSON.stringify(contact),
        })
        return toolText(data)
      }
    )
  }

  if (allow("contacts_delete")) {
    server.tool(
      "contacts_delete",
      "Delete a contact by path",
      { path: z.string() },
      async ({ path }) => {
        const encoded = encodePath(path)
        const data = await ultiFetch(getToken(), `/contacts/${encoded}`, {
          method: "DELETE",
        })
        return toolText(data)
      }
    )
  }

  if (allow("suite_search")) {
    server.tool(
      "suite_search",
      "Unified search across mail, drive, contacts (mail spam/trash excluded by default)",
      {
        query: z.string(),
        types: z.string().optional(),
        limit: z.number().optional(),
        offset: z.number().optional(),
        include_spam: z.boolean().optional(),
        include_trash: z.boolean().optional(),
      },
      async ({ query, types, limit, offset, include_spam, include_trash }) => {
        const qs = withListPagination(new URLSearchParams({ q: query }), limit, offset)
        if (types) qs.set("types", types)
        applyHiddenMailboxParams(qs, include_spam, include_trash)
        const data = await ultiFetch(getToken(), `/search?${qs}`)
        return toolText(data)
      }
    )
  }

  if (allow("web_search")) {
    server.tool(
      "web_search",
      "Search the public web using the configured search provider (Brave, Bing, SearXNG, DuckDuckGo or custom JSON API).",
      {
        query: z.string(),
        count: z.number().optional(),
      },
      async ({ query, count }) => {
        const qs = new URLSearchParams({ q: query })
        if (count != null && Number.isFinite(count) && count > 0) {
          qs.set("count", String(Math.min(Math.trunc(count), 20)))
        }
        const data = await ultiFetch(getToken(), `/search/web?${qs}`)
        return toolText(data)
      }
    )
  }

  if (allow("docs_read")) {
    server.tool(
      "docs_read",
      "Read UltiDocs document JSON (.ultidoc sidecar path)",
      { path: z.string() },
      async ({ path }) => {
        const encoded = encodePath(path)
        const data = await ultiFetch(getToken(), `/drive/download/${encoded}`)
        return toolText(data)
      }
    )
  }

  if (allow("docs_save")) {
    server.tool(
      "docs_save",
      "Save UltiDocs TipTap content to sidecar path",
      {
        path: z.string(),
        document: z.record(z.string(), z.unknown()),
      },
      async ({ path, document }) => {
        const data = await ultiFetch(getToken(), "/richtext/save", {
          method: "PUT",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ path, document }),
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_list")) {
    server.tool(
      "calendar_list",
      "List calendars",
      {},
      async () => {
        const data = await ultiFetch(getToken(), "/calendar/")
        return toolText(data)
      }
    )
  }

  if (allow("calendar_create")) {
    server.tool(
      "calendar_create",
      "Create a calendar",
      {
        id: z.string().optional(),
        display_name: z.string(),
        color: z.string().optional(),
      },
      async (body) => {
        const data = await ultiFetch(getToken(), "/calendar/", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_update")) {
    server.tool(
      "calendar_update",
      "Update calendar display name or color",
      {
        cal_id: z.string(),
        display_name: z.string().optional(),
        color: z.string().optional(),
      },
      async ({ cal_id, ...body }) => {
        const data = await ultiFetch(getToken(), `/calendar/${cal_id}`, {
          method: "PATCH",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_delete")) {
    server.tool(
      "calendar_delete",
      "Delete a calendar",
      { cal_id: z.string() },
      async ({ cal_id }) => {
        const data = await ultiFetch(getToken(), `/calendar/${cal_id}`, {
          method: "DELETE",
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_list_events")) {
    server.tool(
      "calendar_list_events",
      "List events in a calendar",
      {
        cal_id: z.string(),
        limit: z.number().optional(),
        offset: z.number().optional(),
      },
      async ({ cal_id, limit, offset }) => {
        const qs = listPaginationParams(limit, offset)
        const suffix = qs.size > 0 ? `?${qs}` : ""
        const data = await ultiFetch(getToken(), `/calendar/${cal_id}/events${suffix}`)
        return toolText(data)
      }
    )
  }

  if (allow("calendar_freebusy")) {
    server.tool(
      "calendar_freebusy",
      "Query free/busy for attendees in a time range",
      {
        start: z.string(),
        end: z.string(),
        attendees: z.array(z.string()),
      },
      async (body) => {
        const data = await ultiFetch(getToken(), "/calendar/freebusy", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_create_event")) {
    server.tool(
      "calendar_create_event",
      "Create a calendar event",
      {
        cal_id: z.string(),
        event: z.record(z.string(), z.unknown()),
      },
      async ({ cal_id, event }) => {
        const data = await ultiFetch(getToken(), `/calendar/${cal_id}/events`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(event),
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_update_event")) {
    server.tool(
      "calendar_update_event",
      "Update a calendar event (requires If-Match etag)",
      {
        event_path: z.string(),
        if_match: z.string(),
        event: z.record(z.string(), z.unknown()),
      },
      async ({ event_path, if_match, event }) => {
        const encoded = encodePath(event_path)
        const data = await ultiFetch(getToken(), `/calendar/events/${encoded}`, {
          method: "PUT",
          headers: {
            "Content-Type": "application/json",
            "If-Match": if_match,
          },
          body: JSON.stringify(event),
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_delete_event")) {
    server.tool(
      "calendar_delete_event",
      "Delete a calendar event",
      { event_path: z.string() },
      async ({ event_path }) => {
        const encoded = encodePath(event_path)
        const data = await ultiFetch(getToken(), `/calendar/events/${encoded}`, {
          method: "DELETE",
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_respond_invitation")) {
    server.tool(
      "calendar_respond_invitation",
      "Respond to a calendar invitation",
      {
        event_path: z.string(),
        email: z.string().optional(),
        response: z.enum(["accepted", "declined", "tentative", "needs-action"]),
        if_match: z.string().optional(),
      },
      async ({ event_path, ...body }) => {
        const encoded = encodePath(event_path)
        const data = await ultiFetch(getToken(), `/calendar/events/response/${encoded}`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify(body),
        })
        return toolText(data)
      }
    )
  }

  if (allow("calendar_create_meet_link")) {
    server.tool(
      "calendar_create_meet_link",
      "Create a video meeting link for an event",
      {
        event_path: z.string(),
        if_match: z.string().optional(),
      },
      async ({ event_path, if_match }) => {
        const encoded = encodePath(event_path)
        const data = await ultiFetch(getToken(), `/calendar/events/meet-link/${encoded}`, {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ if_match }),
        })
        return toolText(data)
      }
    )
  }

  if (enabledTools && enabledTools.size === 0) {
    console.warn("ultimail-mcp: no tools enabled for groups", [...enabledTools])
  } else if (enabledTools) {
    const disabled = ALL_TOOL_NAMES.filter((name) => !enabledTools.has(name))
    if (disabled.length > 0) {
      console.log("ultimail-mcp disabled tools:", disabled.join(", "))
    }
  }

  return server
}

const app = express()
app.use(express.json({ limit: "4mb" }))

type TransportMap = Map<string, StreamableHTTPServerTransport | SSEServerTransport>
const streamableTransports: TransportMap = new Map()
const sseTransports: TransportMap = new Map()

app.get("/health", (_req, res) => {
  res.json({ ok: true, tools: ALL_TOOL_NAMES })
})

app.all("/mcp", async (req, res) => {
  const token = resolveToken(req)
  const enabledTools = parseEnabledTools(String(req.headers["x-ulti-enabled-tools"] ?? ""))
  try {
    const sessionId = String(req.headers["mcp-session-id"] ?? "")
    let transport: StreamableHTTPServerTransport | undefined

    if (sessionId && streamableTransports.has(sessionId)) {
      transport = streamableTransports.get(sessionId) as StreamableHTTPServerTransport
    } else if (!sessionId && req.method === "POST" && isInitializeRequest(req.body)) {
      transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => randomUUID(),
        onsessioninitialized: (sid) => {
          if (transport) streamableTransports.set(sid, transport)
        },
      })
      transport.onclose = () => {
        const sid = transport?.sessionId
        if (sid) streamableTransports.delete(sid)
      }
      const server = createServer(() => token, enabledTools)
      await server.connect(transport)
    } else {
      res.status(400).json({
        jsonrpc: "2.0",
        error: { code: -32000, message: "Bad Request: No valid session ID provided" },
        id: null,
      })
      return
    }

    await transport.handleRequest(req, res, req.body)
  } catch (error) {
    console.error("streamable mcp error:", error)
    if (!res.headersSent) {
      res.status(500).json({
        jsonrpc: "2.0",
        error: { code: -32603, message: "Internal server error" },
        id: null,
      })
    }
  }
})

app.get("/sse", async (req, res) => {
  const token = resolveToken(req)
  const enabledTools = parseEnabledTools(String(req.headers["x-ulti-enabled-tools"] ?? ""))
  const transport = new SSEServerTransport("/messages", res)
  sseTransports.set(transport.sessionId, transport)
  res.on("close", () => sseTransports.delete(transport.sessionId))
  const server = createServer(() => token, enabledTools)
  await server.connect(transport)
})

app.post("/messages", async (req, res) => {
  const sessionId = String(req.query.sessionId ?? "")
  const transport = sseTransports.get(sessionId)
  if (!transport || !(transport instanceof SSEServerTransport)) {
    res.status(404).json({ error: "session not found" })
    return
  }
  await transport.handlePostMessage(req, res, req.body)
})

app.listen(PORT, () => {
  console.log(`ultimail-mcp listening on :${PORT} (streamable /mcp, legacy /sse)`)
})
