// server/mcpServer.ts
import { Elysia } from "elysia";
import { getMcpTools } from "../lib/mcp_tool_convert";

/**
 * Refactored Elysia-based MCP server
 * - Fixes inconsistent "text/json" handling by normalizing response extraction
 * - Robust executeTool: supports path/query/header/cookie/body params (if provided in x-props)
 * - Proper baseUrl/path normalization and URLSearchParams building (repeated keys for arrays)
 * - Consistent MCP content conversion: always returns either { type: 'json', data } or { type: 'text', text }
 * - Safer error handling, batch support, Promise.allSettled to avoid full failure on single-item error
 * - Lightweight in-memory tools cache with explicit init endpoint (keeps original behavior)
 */

/* -------------------------
   Environment & Globals
   ------------------------- */
if (!process.env.BUN_PUBLIC_BASE_URL) {
  throw new Error("BUN_PUBLIC_BASE_URL environment variable is not set");
}

const OPENAPI_URL = `${process.env.BUN_PUBLIC_BASE_URL.replace(/\/+$/, "")}/docs/json`;
const FILTER_TAG = "mcp";

let tools: any[] = [];

/* -------------------------
   MCP Types
   ------------------------- */
type JSONRPCRequest = {
  jsonrpc: "2.0";
  id: string | number;
  method: string;
  params?: any;
  credentials?: any;
};

type JSONRPCResponse = {
  jsonrpc: "2.0";
  id: string | number | null;
  result?: any;
  error?: { code: number; message: string; data?: any };
};

/* -------------------------
   Helpers
   ------------------------- */

/** Ensure baseUrl doesn't end with slash; ensure path begins with slash */
function joinBasePath(base: string, path: string) {
  const normalizedBase = base.replace(/\/+$/, "");
  const normalizedPath = path ? (path.startsWith("/") ? path : `/${path}`) : "";
  return `${normalizedBase}${normalizedPath}`;
}

/** Serialize query object to repeated-key QS when arrays provided */
function buildQueryString(q: Record<string, any>): string {
  const parts: string[] = [];
  for (const [k, v] of Object.entries(q)) {
    if (v === undefined || v === null) continue;
    if (Array.isArray(v)) {
      for (const item of v) {
        parts.push(`${encodeURIComponent(k)}=${encodeURIComponent(String(item))}`);
      }
    } else if (typeof v === "object") {
      parts.push(`${encodeURIComponent(k)}=${encodeURIComponent(JSON.stringify(v))}`);
    } else {
      parts.push(`${encodeURIComponent(k)}=${encodeURIComponent(String(v))}`);
    }
  }
  return parts.length ? `?${parts.join("&")}` : "";
}

/** Safely extract "useful" payload from a fetch result:
 * Prefer resp.data if present, otherwise resp itself.
 * If resp is a string, keep as string.
 */
function extractRaw(result: { data: any } | any) {
  // If result shaped as { data: ... } prefer inner .data
  if (result && typeof result === "object" && "data" in result) {
    return result.data;
  }
  return result;
}

/** Convert various payloads into MCP content shape */
function convertToMcpContent(payload: any) {
  if (typeof payload === "string") {
    return { type: "text", text: payload };
  }

  if (payload == null) {
    return { type: "text", text: String(payload) };
  }

  // If payload looks like an image/audio wrapper produced by converter
  if (payload?.__mcp_type === "image" && payload.base64) {
    return { type: "image", data: payload.base64, mimeType: payload.mimeType || "image/png" };
  }
  if (payload?.__mcp_type === "audio" && payload.base64) {
    return { type: "audio", data: payload.base64, mimeType: payload.mimeType || "audio/mpeg" };
  }

  // If already an object/array → return JSON
  if (typeof payload === "object") {
    return { type: "json", data: payload };
  }

  // Fallback — stringify
  try {
    return { type: "text", text: JSON.stringify(payload) };
  } catch {
    return { type: "text", text: String(payload) };
  }
}

/* -------------------------
   executeTool (robust)
   ------------------------- */
/**
 * Execute a tool converted from OpenAPI -> expected x-props shape:
 * x-props may contain:
 *  - method, path
 *  - parameters: [{ name, in, required? }]
 *
 * If x.parameters present, we inspect args and place them accordingly.
 */
export async function executeTool(
  tool: any,
  args: Record<string, any> = {},
  baseUrl: string,
  xPayload: Record<string, any> = {}
) {
  const x = tool["x-props"] || {};
  const method = (x.method || "GET").toUpperCase();

  // Start with provided path (may contain {param})
  let path = x.path ?? `/${tool.name}`;

  // Headers, cookies, query, body collection
  const headers: Record<string, any> = {
    "Content-Type": "application/json",
    ...(x.defaultHeaders || {}),
  };
  const query: Record<string, any> = {};
  const cookies: string[] = [];
  let bodyPayload: any = undefined;

  // If parameters described, map args accordingly
  if (Array.isArray(x.parameters)) {
    for (const p of x.parameters) {
      try {
        const name: string = p.name;
        const value = args?.[name];

        // skip undefined unless required — we let API validate required semantics
        if (value === undefined) continue;

        switch (p.in) {
          case "path":
            if (path.includes(`{${name}}`)) {
              path = path.replace(new RegExp(`{${name}}`, "g"), encodeURIComponent(String(value)));
            } else {
              // fallback to query
              query[name] = value;
            }
            break;

          case "query":
            query[name] = value;
            break;

          case "header":
            headers[name] = value;
            break;

          case "cookie":
            cookies.push(`${name}=${value}`);
            break;

          case "body":
          case "requestBody":
            bodyPayload = value;
            break;

          default:
            // unknown location -> place into body
            bodyPayload = bodyPayload ?? {};
            bodyPayload[name] = value;
            break;
        }
      } catch (err) {
        // best-effort: skip problematic param
        console.warn(`[MCP] Skipping parameter ${String(p?.name)} due to error:`, err);
      }
    }
  } else {
    // no param descriptions: assume all args are body
    bodyPayload = Object.keys(args || {}).length ? args : undefined;
  }

  if (cookies.length) {
    headers["Cookie"] = cookies.join("; ");
  }

  // Build full URL
  const urlBase = baseUrl || process.env.BUN_PUBLIC_BASE_URL!;
  let url = joinBasePath(urlBase, path);
  const qs = buildQueryString(query);
  if (qs) url += qs;

  // Build RequestInit
  const opts: RequestInit & { headers?: Record<string, any> } = { method, headers };

  // Body handling for applicable methods
  const bodyMethods = new Set(["POST", "PUT", "PATCH", "DELETE"]);
  const contentTypeLower = (headers["Content-Type"] || "").toLowerCase();

  if (bodyMethods.has(method) && bodyPayload !== undefined) {
    // Support tiny formdata marker shape: { __formdata: true, entries: [ [k,v], ... ] }
    if (bodyPayload && bodyPayload.__formdata === true && Array.isArray(bodyPayload.entries)) {
      const form = new FormData();
      for (const [k, v] of bodyPayload.entries) {
        form.append(k, v as any);
      }
      // Let fetch set boundary
      delete opts.headers!["Content-Type"];
      opts.body = form as any;
    } else if (contentTypeLower.includes("application/x-www-form-urlencoded")) {
      opts.body = new URLSearchParams(bodyPayload as Record<string, string>).toString();
    } else if (contentTypeLower.includes("multipart/form-data")) {
      // If caller explicitly requested multipart but didn't pass FormData — convert object to form
      const form = new FormData();
      if (typeof bodyPayload === "object") {
        for (const [k, v] of Object.entries(bodyPayload)) {
          form.append(k, (v as any) as any);
        }
      } else {
        form.append("payload", String(bodyPayload));
      }
      delete opts.headers!["Content-Type"];
      opts.body = form as any;
    } else {
      // Default JSON
      opts.body = JSON.stringify(bodyPayload);
    }
  }

  // Execute fetch
  console.log(`[MCP] → ${method} ${url}`);
  for(const [key, value] of Object.entries(xPayload)) {
    opts.headers![key] = value;
  }
  const res = await fetch(url, opts);

  const resContentType = (res.headers.get("content-type") || "").toLowerCase();

  let data: any;
  try {
    if (resContentType.includes("application/json")) {
      data = await res.json();
    } else {
      data = await res.text();
    }
  } catch (err) {
    // fallback to text
    try {
      data = await res.text();
    } catch {
      data = null;
    }
  }

  return {
    success: res.ok,
    status: res.status,
    method,
    url,
    path,
    headers: res.headers,
    data,
  };
}

/* -------------------------
   JSON-RPC Handler
   ------------------------- */
async function handleMCPRequestAsync(request: JSONRPCRequest, xPayload: Record<string, any>): Promise<JSONRPCResponse> {
  const { id, method, params } = request;

  const makeError = (code: number, message: string, data?: any): JSONRPCResponse => ({
    jsonrpc: "2.0",
    id: id ?? null,
    error: { code, message, data },
  });

  switch (method) {
    case "initialize":
      return {
        jsonrpc: "2.0",
        id,
        result: {
          protocolVersion: "2024-11-05",
          capabilities: { tools: {} },
          serverInfo: { name: "elysia-mcp-server", version: "1.0.0" },
        },
      };

    case "tools/list":
      return {
        jsonrpc: "2.0",
        id,
        result: {
          tools: tools.map((t) => {
            const inputSchema =
              typeof t.inputSchema === "object" && t.inputSchema?.type === "object"
                ? t.inputSchema
                : { type: "object", properties: {}, required: [] };
            return {
              name: t.name,
              description: t.description || "No description provided",
              inputSchema,
              "x-props": t["x-props"],
            };
          }),
        },
      };

    case "tools/call": {
      const toolName = params?.name;
      const tool = tools.find((t) => t.name === toolName);
      if (!tool) return makeError(-32601, `Tool '${toolName}' not found`);

      try {
        const baseUrl = (params?.credentials?.baseUrl as string) || process.env.BUN_PUBLIC_BASE_URL || "http://localhost:3000";
        const args = params?.arguments || {};

        const result = await executeTool(tool, args, baseUrl, xPayload);

        // Extract the meaningful payload (prefer nested .data if present)
        const raw = extractRaw(result.data);

        // Normalize content shape consistently:
        const contentItem = convertToMcpContent(raw ?? result.data ?? result);

        return {
          jsonrpc: "2.0",
          id,
          result: {
            content: [contentItem],
          },
        };
      } catch (err: any) {
        // avoid leaking secrets — small debug
        const dbg = { message: err?.message };
        return makeError(-32603, err?.message ?? "Internal error", dbg);
      }
    }

    case "ping":
      return { jsonrpc: "2.0", id, result: {} };

    default:
      return makeError(-32601, `Method '${method}' not found`);
  }
}

/* -------------------------
   Elysia App & Routes
   ------------------------- */
export const MCPRoute = new Elysia({ tags: ["MCP Server"] })
  .post("/mcp", async ({ request, set, headers }) => {
    set.headers["Content-Type"] = "application/json";
    set.headers["Access-Control-Allow-Origin"] = "*";

    // Lazy load the tools (keeps previous behavior)
    if (!tools.length) {
      try {
        tools = await getMcpTools(OPENAPI_URL, FILTER_TAG);
      } catch (err) {
        console.error("[MCP] Failed to load tools during lazy init:", err);
      }
    }

    const xPayload = {
      ['x-user']: headers['x-user'] || "",
      ['x-phone']: headers['x-phone'] || ""
    }

    try {
      const body = await request.json();

      // If batch array -> allSettled for resilience
      if (Array.isArray(body)) {
        const promises = body.map((req: JSONRPCRequest) => handleMCPRequestAsync(req, xPayload));
        const settled = await Promise.allSettled(promises);
        const responses = settled.map((s) =>
          s.status === "fulfilled"
            ? s.value
            : ({
                jsonrpc: "2.0",
                id: null,
                error: {
                  code: -32000,
                  message: "Unhandled handler error",
                  data: String((s as PromiseRejectedResult).reason),
                },
              } as JSONRPCResponse)
        );
        return responses;
      }

      const single = await handleMCPRequestAsync(body as JSONRPCRequest, xPayload);
      return single;
    } catch (err: any) {
      set.status = 400;
      return {
        jsonrpc: "2.0",
        id: null,
        error: { code: -32700, message: "Parse error", data: err?.message ?? String(err) },
      } as JSONRPCResponse;
    }
  })

  /* Debug / management endpoints */
  .get("/mcp/tools", async ({ set }) => {
    set.headers["Access-Control-Allow-Origin"] = "*";
    if (!tools.length) {
      try {
        tools = await getMcpTools(OPENAPI_URL, FILTER_TAG);
      } catch (err) {
        console.error("[MCP] Failed to load tools for /mcp/tools:", err);
      }
    }
    return {
      tools: tools.map((t) => ({
        name: t.name,
        description: t.description,
        inputSchema: t.inputSchema,
        "x-props": t["x-props"],
      })),
    };
  })

  .get("/mcp/status", ({ set }) => {
    set.headers["Access-Control-Allow-Origin"] = "*";
    return { status: "active", timestamp: Date.now() };
  })

  .get("/health", ({ set }) => {
    set.headers["Access-Control-Allow-Origin"] = "*";
    return { status: "ok", timestamp: Date.now(), tools: tools.length };
  })

  // Force re-init (useful for admin / CI)
  .get("/mcp/init", async ({ set }) => {
    set.headers["Access-Control-Allow-Origin"] = "*";
    try {
      tools = await getMcpTools(OPENAPI_URL, FILTER_TAG);
      return { success: true, message: "MCP initialized", tools: tools.length };
    } catch (err) {
      set.status = 500;
      return { success: false, message: "Failed to initialize tools", error: String(err) };
    }
  })

  /* CORS preflight */
  .options("/mcp", ({ set }) => {
    set.headers["Access-Control-Allow-Origin"] = "*";
    set.headers["Access-Control-Allow-Methods"] = "GET,POST,OPTIONS";
    set.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization,X-API-Key";
    set.status = 204;
    return "";
  })
  .options("/mcp/tools", ({ set }) => {
    set.headers["Access-Control-Allow-Origin"] = "*";
    set.headers["Access-Control-Allow-Methods"] = "GET,OPTIONS";
    set.headers["Access-Control-Allow-Headers"] = "Content-Type,Authorization,X-API-Key";
    set.status = 204;
    return "";
  });

/* -------------------------
   End
   ------------------------- */