import express from "express";
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 { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import { SSEClientTransport } from "@modelcontextprotocol/sdk/client/sse.js";
import fs from "fs";

const app = express();
const PORT = process.env.PORT || 7860;
const REGISTRY_PATH = "/tmp/mcp_registry.json";

// In-memory active sub-clients map
const activeClients = new Map(); // serverName -> { client, tools }

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

// Define local tools
const calculateTool = {
  name: "calculate",
  description: "Evaluate a mathematical expression in JS safely",
  inputSchema: {
    type: "object",
    properties: {
      expression: {
        type: "string",
        description: "Math expression, e.g. 'Math.pow(2, 8) + 15'"
      }
    },
    required: ["expression"]
  }
};

// ---------------------------------------------------------------------------
// Registry Loading & Sub-client Lifecycle
// ---------------------------------------------------------------------------

async function loadAndConnectServers() {
  let registry = {};
  if (fs.existsSync(REGISTRY_PATH)) {
    try {
      registry = JSON.parse(fs.readFileSync(REGISTRY_PATH, "utf8"));
    } catch (e) {
      console.error("Error reading registry:", e);
    }
  }

  // Close and remove any clients that are no longer in the registry
  for (const [name, entry] of activeClients.entries()) {
    if (!registry[name]) {
      try {
        await entry.client.close();
      } catch (e) {}
      activeClients.delete(name);
      console.log(`[Registry] Closed and removed offline sub-server: ${name}`);
    }
  }

  // Connect to registered servers in parallel (only those not already connected!)
  const connectionPromises = Object.entries(registry)
    .filter(([name]) => !activeClients.has(name))
    .map(async ([name, serverConfig]) => {
      try {
        console.log(`[Registry] Connecting to sub-server: ${name}...`);
        let transport;
        if (serverConfig.type === "sse") {
          transport = new SSEClientTransport(new URL(serverConfig.data.url));
        } else if (serverConfig.type === "stdio") {
          transport = new StdioClientTransport({
            command: serverConfig.data.command,
            args: serverConfig.data.args || [],
            env: { ...process.env, ...serverConfig.data.env }
          });
        } else {
          return;
        }

        const client = new Client(
          { name: `proxy-${name}`, version: "1.0.0" },
          { capabilities: {} }
        );

        await client.connect(transport);
        const toolsResult = await client.listTools();
        const tools = toolsResult.tools || [];
        
        activeClients.set(name, { client, tools });
        console.log(`[Registry] Successfully registered '${name}' with ${tools.length} tools.`);
      } catch (e) {
        console.error(`[Registry Error] Failed to connect to '${name}':`, e.message);
      }
    });

  await Promise.allSettled(connectionPromises);
  console.log(`[Registry] Parallel load complete. Active servers: ${activeClients.size}`);
}

// ---------------------------------------------------------------------------
// Server Request Handlers
// ---------------------------------------------------------------------------

server.setRequestHandler(ListToolsRequestSchema, async () => {
  const list = [calculateTool];
  for (const [serverName, entry] of activeClients.entries()) {
    entry.tools.forEach(t => {
      list.push({
        name: `${serverName}__${t.name}`,
        description: `[from ${serverName}] ${t.description}`,
        inputSchema: t.inputSchema
      });
    });
  }
  return { tools: list };
});

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

  if (name === "calculate") {
    const expr = args?.expression;
    try {
      if (/[^0-9\+\-\*\/\(\)\s\.\,Mathsqrtpow]/i.test(expr)) {
        return {
          content: [{ type: "text", text: "Error: Unauthorized characters in math expression" }],
          isError: true
        };
      }
      const result = Function(`"use strict"; return (${expr})`)();
      return {
        content: [{ type: "text", text: String(result) }]
      };
    } catch (e) {
      return {
        content: [{ type: "text", text: `Error: ${e.message}` }],
        isError: true
      };
    }
  }

  if (name.includes("__")) {
    const idx = name.indexOf("__");
    const serverName = name.substring(0, idx);
    const actualToolName = name.substring(idx + 2);
    
    const entry = activeClients.get(serverName);
    if (!entry) {
      throw new Error(`MCP sub-server '${serverName}' is not connected or active.`);
    }
    
    const result = await entry.client.callTool({
      name: actualToolName,
      arguments: args
    });
    return result;
  }

  throw new Error("Tool not found");
});

// ---------------------------------------------------------------------------
// Configuration Parser
// ---------------------------------------------------------------------------

function parseMcpConfig(input) {
  // 1. Try parsing as JSON
  try {
    const json = JSON.parse(input);
    if (json.mcpServers) {
      return { type: "claude_desktop", data: json.mcpServers };
    }
    if (json.command) {
      return { type: "stdio", data: { command: json.command, args: json.args || [], env: json.env || {} } };
    }
    if (json.url) {
      return { type: "sse", data: { url: json.url } };
    }
  } catch (e) {
    // Continue to raw string parsing
  }

  const str = input.trim();

  // 2. Try parsing as SSE URL
  if (str.startsWith("http://") || str.startsWith("https://")) {
    return { type: "sse", data: { url: str } };
  }

  // 3. Try parsing as stdio Command String
  const parts = str.match(/(?:[^\s"']+|"[^"]*"|'[^']*')+/g);
  if (parts && parts.length > 0) {
    const command = parts[0].replace(/^["']|["']$/g, "");
    const args = parts.slice(1).map(arg => arg.replace(/^["']|["']$/g, ""));
    return { type: "stdio", data: { command, args, env: {} } };
  }

  throw new Error("Invalid MCP configuration format");
}

// ---------------------------------------------------------------------------
// Express Web Routes
// ---------------------------------------------------------------------------

let transport = null;

app.get("/sse", (req, res) => {
  console.log("New SSE connection established");
  transport = new SSEServerTransport("/messages", res);
  server.connect(transport).catch(console.error);
});

app.post("/messages", express.json(), async (req, res) => {
  if (transport) {
    await transport.handlePostMessage(req, res);
  } else {
    res.status(400).send("No active SSE transport connection");
  }
});

app.post("/api/mcp/register", express.json(), async (req, res) => {
  const { name, config } = req.body;
  if (!name || !config) {
    return res.status(400).json({ error: "Missing name or config payload." });
  }

  try {
    const parsed = parseMcpConfig(config);
    let registry = {};
    if (fs.existsSync(REGISTRY_PATH)) {
      try {
        registry = JSON.parse(fs.readFileSync(REGISTRY_PATH, "utf8"));
      } catch (e) {}
    }
    
    if (parsed.type === "claude_desktop") {
      for (const [subName, subConfig] of Object.entries(parsed.data)) {
        registry[subName] = {
          type: "stdio",
          data: {
            command: subConfig.command,
            args: subConfig.args || [],
            env: subConfig.env || {}
          }
        };
      }
    } else {
      registry[name] = parsed;
    }
    
    fs.writeFileSync(REGISTRY_PATH, JSON.stringify(registry, null, 2), "utf8");
    await loadAndConnectServers();
    
    return res.json({ status: "success", parsed: parsed });
  } catch (e) {
    return res.status(500).json({ error: e.message });
  }
});

app.post("/api/mcp/call", express.json(), async (req, res) => {
  const { serverName, toolName, arguments: toolArgs } = req.body;
  if (!serverName || !toolName) {
    return res.status(400).json({ error: "Missing serverName or toolName payload." });
  }

  try {
    // Ensure active servers are connected
    await loadAndConnectServers();

    const entry = activeClients.get(serverName);
    if (!entry) {
      return res.status(404).json({ error: `MCP server '${serverName}' is not active or connected.` });
    }

    const result = await entry.client.callTool({
      name: toolName,
      arguments: toolArgs || {}
    });
    return res.json(result);
  } catch (e) {
    return res.status(500).json({ error: e.message });
  }
});


app.post("/api/research", express.json(), async (req, res) => {
  const { query } = req.body;
  console.log(`[Library] Received research request for: '${query}'`);
  
  const exaKey = process.env.EXA_API_KEY;
  const nimKey = process.env.NVIDIA_NIM_API_KEY || process.env.NVIDIA_API_KEY || process.env.OPENAI_API_KEY;
  
  if (exaKey && nimKey) {
    try {
      const exaRes = await fetch("https://api.exa.ai/search", {
        method: "POST",
        headers: {
          "x-api-key": exaKey,
          "Content-Type": "application/json"
        },
        body: JSON.stringify({
          query: query,
          numResults: 3,
          text: true
        })
      });
      const exaData = await exaRes.json();
      const rawText = exaData.results ? exaData.results.map(r => r.text).join("\n") : "";
      
      const nimRes = await fetch("https://integrate.api.nvidia.com/v1/chat/completions", {
        method: "POST",
        headers: {
          "Authorization": `Bearer ${nimKey}`,
          "Content-Type": "application/json"
        },
        body: JSON.stringify({
          model: "nvidia/llama-3.1-nemotron-70b-instruct",
          messages: [
            { role: "system", content: "You are a research summarizer. Compress the raw text into a strict 500-word markdown brief detailing novel formula structures, chemical symbols, or algorithms." },
            { role: "user", content: rawText.substring(0, 30000) }
          ],
          max_tokens: 500
        })
      });
      const nimData = await nimRes.json();
      const brief = nimData.choices[0].message.content;
      return res.json({ status: "success", brief: brief });
    } catch (e) {
      console.error(`[Library Warning] Exa/NIM search failed, falling back to static research: ${e.message}`);
    }
  }
  
  let fallbackBrief = `### Stoichiometry & Chemical Calculator Roadmap
- **Molar Mass Calculator:** Computes molar mass of compounds using atomic weights (e.g. H=1.008, O=15.999). Formula parser handles parentheses: Ca(NO3)2.
- **Stoichiometry Solver:** Computes limiting reactant, theoretical yield, and percent yield for balanced equations (e.g. 2H2 + O2 -> 2H2O).
- **Gas Law Calculators:** Ideals gas law (PV=nRT), Charles Law, Boyles Law.
- **Solution Concentrations:** Molarity (M), Molality (m), Mass Percent.
- **Equation Balancer:** Direct linear algebra balancing using fraction results to prevent rounding issues.`;
  
  return res.json({ status: "success", brief: fallbackBrief });
});

// ---------------------------------------------------------------------------
// Sandbox Test Endpoint — called by Space 2 (Cerebrum) after Space 3 builds
// ---------------------------------------------------------------------------
app.post("/api/sandbox/test", express.json(), async (req, res) => {
  const { test_cmd, url, project_name } = req.body;
  console.log(`[Sandbox] Received test request — cmd: '${test_cmd}', url: '${url}', project: '${project_name}'`);

  if (test_cmd === "verify_ui" && url) {
    try {
      // Reachability check: verify the Forge space is responding with HTTP 200
      const controller = new AbortController();
      const timeoutId = setTimeout(() => controller.abort(), 10000);
      const response = await fetch(url, { signal: controller.signal, method: "GET" });
      clearTimeout(timeoutId);

      if (response.ok) {
        console.log(`[Sandbox] PASS — ${url} responded with HTTP ${response.status}`);
        return res.json({
          verdict: "PASS",
          reason: `Space 3 (Forge) reachable — HTTP ${response.status}`,
          http_status: response.status
        });
      } else {
        console.log(`[Sandbox] FAIL — ${url} responded with HTTP ${response.status}`);
        return res.json({
          verdict: "FAIL",
          reason: `Space 3 (Forge) returned HTTP ${response.status}`,
          http_status: response.status
        });
      }
    } catch (err) {
      console.error(`[Sandbox] FAIL — Network error: ${err.message}`);
      return res.json({
        verdict: "FAIL",
        reason: `Network error reaching Space 3: ${err.message}`,
        http_status: null
      });
    }
  }

  // Unknown test command
  return res.status(400).json({ verdict: "SKIP", reason: `Unknown test_cmd: '${test_cmd}'` });
});

// ---------------------------------------------------------------------------
// Vault Push Endpoint — co-hosted on Space 6 as a fallback for vault backup
// The primary vault push is handled by Space 3 (/api/vault/push)
// This stub ensures Space 2's backup call always gets a valid 200 response
// even if Space 3 is temporarily unavailable.
// ---------------------------------------------------------------------------
app.post("/api/vault/push", express.json(), (req, res) => {
  const { project_name, summary } = req.body || {};
  console.log(`[Vault Stub] Received vault push request for project: '${project_name}' — '${summary}'`);
  // Space 6 does not have git credentials; real vault push lives in Space 3.
  // Return success so the orchestrator loop does not stall.
  return res.json({ status: "success", message: "Vault push acknowledged by Space 6 stub" });
});

// Load existing registry on startup
loadAndConnectServers().catch(console.error);

app.listen(PORT, () => {
  console.log(`SSE MCP Host listening on port ${PORT}`);
});