#!/usr/bin/env node
import express from "express";
import fs from "node:fs";
import path from "node:path";
import { fileURLToPath } from "node:url";
import dotenv from "dotenv";
import OpenAI from "openai";
import { MCPClient } from "mcp-use";

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

const app = express();
const port = Number(process.env.PORT) || 3000;
const publicDir = path.join(__dirname, "public");
const workspaceRoot = path.resolve(__dirname, "..");
const mcpConfigPath = path.join(workspaceRoot, ".vscode", "mcp.json");

dotenv.config({ path: path.join(workspaceRoot, ".env") });
console.log(`OPENAI_API_KEY loaded: ${Boolean(process.env.OPENAI_API_KEY)}`);

const openai = new OpenAI({ apiKey: process.env.OPENAI_API_KEY });
const MCP_REQUEST_TIMEOUT_MS =
  Number(process.env.MCP_REQUEST_TIMEOUT_MS) || 120000;
const MCP_RETRY_COUNT = Number(process.env.MCP_RETRY_COUNT) || 2;
const MCP_RETRY_DELAY_MS = Number(process.env.MCP_RETRY_DELAY_MS) || 2000;

function loadMcpConfig() {
  let raw = {};
  try {
    raw = JSON.parse(fs.readFileSync(mcpConfigPath, "utf8"));
  } catch (error) {
    console.error("Failed to read MCP config:", mcpConfigPath, error);
    return { mcpServers: {}, error: "Unable to read .vscode/mcp.json" };
  }

  const servers = raw.servers || {};
  const mcpServers = {};

  for (const [name, server] of Object.entries(servers)) {
    if (!server || typeof server !== "object") {
      continue;
    }
    let cwd = server.cwd || workspaceRoot;
    if (typeof cwd === "string") {
      cwd = cwd.replaceAll("${workspaceFolder}", workspaceRoot);
      if (!path.isAbsolute(cwd)) {
        cwd = path.resolve(workspaceRoot, cwd);
      }
    }
    const resolvedArgs = (server.args || []).map((arg) => {
      if (typeof arg !== "string" || !cwd) {
        return arg;
      }
      if (arg.startsWith("-")) {
        return arg;
      }
      const resolvedPath = path.resolve(cwd, arg);
      return fs.existsSync(resolvedPath) ? resolvedPath : arg;
    });
    mcpServers[name] = {
      command: server.command,
      args: resolvedArgs,
      cwd
    };
  }

  console.log("Loaded MCP servers:", Object.keys(mcpServers));
  return { mcpServers, error: null };
}

const mcpConfig = loadMcpConfig();
const mcpClient = MCPClient.fromDict({ mcpServers: mcpConfig.mcpServers });

async function getSession(serverName) {
  console.log(`[mcp] getSession start: ${serverName}`);
  let session = mcpClient.getSession(serverName);
  if (!session) {
    const start = Date.now();
    session = await mcpClient.createSession(serverName, true);
    console.log(
      `[mcp] session created: ${serverName} in ${Date.now() - start}ms`
    );
  }
  console.log(`[mcp] getSession ready: ${serverName}`);
  return session;
}

function isTimeoutError(error) {
  const message = String(error?.message || error || "");
  return (
    message.includes("timed out") ||
    message.includes("timeout") ||
    message.includes("-32001")
  );
}

function delay(ms) {
  return new Promise((resolve) => setTimeout(resolve, ms));
}

async function withRetries(fn, label) {
  let lastError;
  for (let attempt = 0; attempt <= MCP_RETRY_COUNT; attempt += 1) {
    try {
      if (attempt > 0) {
        console.warn(`[mcp] retrying ${label} (attempt ${attempt + 1})`);
      }
      return await fn();
    } catch (error) {
      lastError = error;
      if (!isTimeoutError(error) || attempt === MCP_RETRY_COUNT) {
        throw error;
      }
      await delay(MCP_RETRY_DELAY_MS);
    }
  }
  throw lastError;
}

function toSafeToolName(serverName, toolName) {
  const combined = `${serverName}__${toolName}`;
  return combined.replace(/[^a-zA-Z0-9_-]/g, "_");
}

function normalizeSchema(inputSchema) {
  if (!inputSchema || typeof inputSchema !== "object") {
    return { type: "object", properties: {}, additionalProperties: true };
  }
  if (!inputSchema.type) {
    return { ...inputSchema, type: "object" };
  }
  return inputSchema;
}

async function buildToolIndex() {
  const index = new Map();
  const toolSpecs = [];
  const failures = [];

  for (const serverName of Object.keys(mcpConfig.mcpServers)) {
    console.log(`[mcp] listing tools for: ${serverName}`);
    try {
      const start = Date.now();
      const session = await getSession(serverName);
      const tools = await withRetries(
        () => session.listTools({ timeoutMs: MCP_REQUEST_TIMEOUT_MS }),
        `listTools:${serverName}`
      );
      console.log(
        `[mcp] tools loaded for ${serverName}: ${tools.length} in ${Date.now() - start}ms`
      );

      for (const tool of tools) {
        const safeName = toSafeToolName(serverName, tool.name);
        index.set(safeName, { serverName, toolName: tool.name });
        toolSpecs.push({
          type: "function",
          function: {
            name: safeName,
            description: `[${serverName}] ${tool.description || tool.name}`,
            parameters: normalizeSchema(tool.inputSchema)
          }
        });
      }
    } catch (error) {
      failures.push({
        server: serverName,
        error: String(error?.message || error)
      });
      console.error(`[mcp] tools load failed for ${serverName}:`, error);
    }
  }

  return { toolSpecs, index, failures };
}

app.use(express.json({ limit: "200kb" }));
app.use(express.static(publicDir));

app.get("/", (_req, res) => {
  res.sendFile(path.join(publicDir, "index.html"));
});

app.get("/api/mcp/servers", (_req, res) => {
  if (mcpConfig.error) {
    res.status(500).json({ error: mcpConfig.error });
    return;
  }
  res.json({ servers: Object.keys(mcpConfig.mcpServers) });
});

app.get("/api/mcp/tools", async (req, res) => {
  const server = String(req.query.server || "").trim();
  if (!server) {
    res.status(400).json({ error: "Server is required." });
    return;
  }

  try {
    const start = Date.now();
    const session = await getSession(server);
    const tools = await withRetries(
      () => session.listTools({ timeoutMs: MCP_REQUEST_TIMEOUT_MS }),
      `listTools:${server}`
    );
    console.log(
      `[mcp] /api/mcp/tools ${server}: ${tools.length} in ${Date.now() - start}ms`
    );
    res.json({ tools });
  } catch (error) {
    console.error(`[mcp] /api/mcp/tools error ${server}:`, error);
    res.status(500).json({ error: String(error?.message || error) });
  }
});

app.post("/api/mcp/call", async (req, res) => {
  const server = String(req.body?.server || "").trim();
  const tool = String(req.body?.tool || "").trim();
  const args = req.body?.args || {};

  if (!server || !tool) {
    res.status(400).json({ error: "Server and tool are required." });
    return;
  }

  try {
    const start = Date.now();
    const session = await getSession(server);
    const result = await withRetries(
      () => session.callTool(tool, args, { timeoutMs: MCP_REQUEST_TIMEOUT_MS }),
      `callTool:${server}.${tool}`
    );
    console.log(
      `[mcp] /api/mcp/call ${server}.${tool} in ${Date.now() - start}ms`
    );
    res.json({ result });
  } catch (error) {
    console.error(`[mcp] /api/mcp/call error ${server}.${tool}:`, error);
    res.status(500).json({ error: String(error?.message || error) });
  }
});

app.post("/api/agent", async (req, res) => {
  const prompt = String(req.body?.prompt || "").trim();
  if (!prompt) {
    res.status(400).json({ error: "Prompt is required." });
    return;
  }

  if (!process.env.OPENAI_API_KEY) {
    res.status(500).json({ error: "OPENAI_API_KEY is not set." });
    return;
  }

  try {
    const model = process.env.OPENAI_MODEL || "gpt-4o";
    console.log(`[agent] prompt: ${prompt}`);
    const buildStart = Date.now();
    const { toolSpecs, index, failures } = await buildToolIndex();
    console.log(
      `[agent] tool index built: ${toolSpecs.length} tools in ${Date.now() - buildStart}ms`
    );
    if (failures.length > 0) {
      console.warn(
        `[agent] tool index failures: ${failures.map((item) => item.server).join(", ")}`
      );
    }

    const messages = [
      {
        role: "system",
        content:
          "You are an Agentforce-style assistant. Use MCP tools to answer. " +
          "If tools are needed, call them. Return a concise final response."
      },
      { role: "user", content: prompt }
    ];

    const toolTrace = [];
    for (let i = 0; i < 5; i += 1) {
      const response = await openai.chat.completions.create({
        model,
        messages,
        tools: toolSpecs,
        tool_choice: "auto"
      });

      const message = response.choices[0]?.message;
      if (!message) {
        res.status(500).json({ error: "No response from model." });
        return;
      }

      if (!message.tool_calls || message.tool_calls.length === 0) {
        res.json({ text: message.content || "", toolTrace });
        return;
      }

      messages.push(message);

      for (const call of message.tool_calls) {
        const mapping = index.get(call.function.name);
        if (!mapping) {
          res
            .status(400)
            .json({ error: `Unknown tool: ${call.function.name}` });
          return;
        }

        let args = {};
        try {
          args = call.function.arguments
            ? JSON.parse(call.function.arguments)
            : {};
        } catch (error) {
          res.status(400).json({ error: "Tool arguments must be valid JSON." });
          return;
        }

        const session = await getSession(mapping.serverName);
        const toolStart = Date.now();
        const result = await withRetries(
          () =>
            session.callTool(mapping.toolName, args, {
              timeoutMs: MCP_REQUEST_TIMEOUT_MS
            }),
          `callTool:${mapping.serverName}.${mapping.toolName}`
        );
        console.log(
          `[agent] tool ${mapping.serverName}.${mapping.toolName} in ${Date.now() - toolStart}ms`
        );
        toolTrace.push({
          server: mapping.serverName,
          tool: mapping.toolName,
          args,
          result
        });

        messages.push({
          role: "tool",
          tool_call_id: call.id,
          content: JSON.stringify(result)
        });
      }
    }

    res.status(500).json({ error: "Tool loop limit reached." });
  } catch (error) {
    res.status(500).json({ error: String(error?.message || error) });
  }
});

app.listen(port, () => {
  console.log(`Agentforce UI running on http://localhost:${port}`);
});