import express from 'express';
import { randomUUID } from 'node:crypto';
import { spawn } from 'child_process';
import fs from 'fs';
import cors from 'cors';
import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { isInitializeRequest, ListToolsRequestSchema, CallToolRequestSchema } from '@modelcontextprotocol/sdk/types.js';

const API_TOKEN = 'FFG-Trinity-2026-Core-Access';
const LOG_FILE = '/home/claude_executor/mcp-server/command.log';
const PORT = 3000;
const BASE_URL = 'https://mcp.firefrostgaming.com';
const ARBITER_LOG_URL = 'https://discord-bot.firefrostgaming.com/api/internal/mcp/log';
const ARBITER_TOKEN = '6fYF1akCRW6pM2F8n3S3RxeIod4YgRniUJNEQurvBP4=';

const SERVERS = {
  'command-center': { host: '63.143.34.217', user: 'root' },
  'tx1-dallas': { host: '38.68.14.26', user: 'root' },
  'nc1-charlotte': { host: '216.239.104.130', user: 'root' },
  'panel-vps': { host: '45.94.168.138', user: 'root' },
  'dev-panel': { host: '64.50.188.128', user: 'root' },
  'wiki-vps': { host: '64.50.188.14', user: 'architect' },
  'services-vps': { host: '38.68.14.188', user: 'root' },
  'trinity-core': { host: 'localhost', user: 'claude_executor', local: true }
};

function log(msg) {
  const line = `[${new Date().toISOString()}] ${msg}\n`;
  fs.appendFileSync(LOG_FILE, line);
  console.log(line.trim());
}

async function logToArbiter(data) {
  try {
    await fetch(ARBITER_LOG_URL, {
      method: 'POST',
      headers: {
        'Content-Type': 'application/json',
        'Authorization': `Bearer ${ARBITER_TOKEN}`
      },
      body: JSON.stringify(data)
    });
  } catch (err) {
    log(`Arbiter log failed: ${err.message}`);
  }
}

const app = express();
app.use(cors({ origin: '*', credentials: true }));
app.use(express.json());
app.use(express.urlencoded({ extended: true }));

function auth(req, res, next) {
  if (req.method === 'OPTIONS') return next();
  const token = req.headers.authorization?.replace('Bearer ', '');
  if (token !== API_TOKEN) {
    log(`AUTH FAILED from ${req.ip}`);
    res.setHeader('WWW-Authenticate', 'Bearer realm="TrinityCore"');
    return res.status(401).json({ error: 'Unauthorized' });
  }
  next();
}

app.get('/.well-known/oauth-protected-resource', (req, res) => {
  log(`OAuth discovery`);
  res.json({ authorization_server: BASE_URL, authorization_endpoint: `${BASE_URL}/authorize`, token_endpoint: `${BASE_URL}/token` });
});

app.get('/authorize', (req, res) => {
  const { redirect_uri, state } = req.query;
  log(`OAUTH authorize -> ${redirect_uri}`);
  res.redirect(`${redirect_uri}?code=trinity-auth-code-123&state=${state}`);
});

app.post('/token', (req, res) => {
  log(`OAUTH token issued`);
  res.json({ access_token: API_TOKEN, token_type: 'Bearer', expires_in: 31536000 });
});

function executeLocal(command) {
  return new Promise((resolve) => {
    log(`EXEC [trinity-core] (local) ${command}`);
    const proc = spawn('bash', ['-c', command]);
    let stdout = '', stderr = '';
    proc.stdout.on('data', (d) => stdout += d.toString());
    proc.stderr.on('data', (d) => stderr += d.toString());
    const timeout = setTimeout(() => {
      proc.kill();
      resolve({ success: false, stdout: stdout.trim(),
        stderr: 'Timeout', error: 'Timeout' });
    }, 30000);
    proc.on('close', (code) => {
      clearTimeout(timeout);
      resolve({ success: code === 0,
        stdout: stdout.trim(),
        stderr: stderr.trim(),
        error: code === 0 ? null : `Exit ${code}` });
    });
    proc.on('error', (err) => {
      clearTimeout(timeout);
      resolve({ success: false, stdout: '',
        stderr: err.message, error: err.message });
    });
  });
}

function executeSSH(server, command) {
  return new Promise((resolve) => {
    const target = SERVERS[server];
    log(`EXEC [${server}] ${command}`);
    const ssh = spawn('ssh', ['-o', 'ConnectTimeout=10', '-o', 'StrictHostKeyChecking=no', `${target.user}@${target.host}`, command]);
    let stdout = '', stderr = '';
    ssh.stdout.on('data', (data) => stdout += data.toString());
    ssh.stderr.on('data', (data) => stderr += data.toString());
    const timeout = setTimeout(() => { ssh.kill(); resolve({ success: false, stdout: stdout.trim(), stderr: 'Timeout', error: 'Timeout' }); }, 30000);
    ssh.on('close', (code) => { clearTimeout(timeout); resolve({ success: code === 0, stdout: stdout.trim(), stderr: stderr.trim(), error: code === 0 ? null : `Exit ${code}` }); });
    ssh.on('error', (err) => { clearTimeout(timeout); resolve({ success: false, stdout: '', stderr: err.message, error: err.message }); });
  });
}

function setupToolHandlers(mcpServer) {
  mcpServer.setRequestHandler(ListToolsRequestSchema, async () => {
    log(`>>> ListTools request`);
    return { tools: [
      { name: "list_servers", description: "Get available Firefrost servers.", inputSchema: { type: "object", properties: {} } },
      { name: "run_command", description: "Execute SSH command on a Firefrost server.", inputSchema: { type: "object", properties: { server: { type: "string" }, command: { type: "string" } }, required: ["server", "command"] }}
    ]};
  });

  mcpServer.setRequestHandler(CallToolRequestSchema, async (request) => {
    log(`>>> CallTool: ${request.params.name}`);
    if (request.params.name === "list_servers") {
      return { content: [{ type: "text", text: `Available: ${Object.keys(SERVERS).join(", ")}` }] };
    }
    if (request.params.name === "run_command") {
      const { server, command } = request.params.arguments;
      if (!SERVERS[server]) return { content: [{ type: "text", text: `Unknown server ${server}` }], isError: true };
      const target = SERVERS[server];
      const startTime = Date.now();
      const result = target.local
        ? await executeLocal(command)
        : await executeSSH(server, command);
      const executionTime = Date.now() - startTime;
      logToArbiter({ server, command, success: result.success, stdout: result.stdout, stderr: result.stderr, error: result.error, execution_time_ms: executionTime });
      let output = result.stdout || '(no output)';
      if (result.stderr) output += `\nSTDERR: ${result.stderr}`;
      return { content: [{ type: "text", text: output }], isError: !result.success };
    }
  });
}

// Track all active transports (both SSE and Streamable HTTP)
const activeSessions = new Map();

// Create a fresh MCP server instance with tool handlers
function createMcpServer() {
  const mcpServer = new Server(
    { name: "trinity-core", version: "2.4.0" },
    { capabilities: { tools: {} } }
  );
  setupToolHandlers(mcpServer);
  return mcpServer;
}

// ─── REST API (for Arbiter / internal services) ───

app.get('/servers', auth, (req, res) => {
  log(`REST /servers`);
  res.json({ servers: Object.keys(SERVERS) });
});

app.post('/exec', auth, async (req, res) => {
  const { server, command } = req.body;
  if (!server || !command) {
    return res.status(400).json({ error: 'Missing server or command' });
  }
  if (!SERVERS[server]) {
    return res.status(400).json({ error: `Unknown server: ${server}` });
  }
  const target = SERVERS[server];
  const startTime = Date.now();
  const result = target.local
    ? await executeLocal(command)
    : await executeSSH(server, command);
  const executionTime = Date.now() - startTime;
  log(`REST /exec [${server}] -> ${result.success ? 'OK' : 'FAIL'} (${executionTime}ms)`);
  logToArbiter({ server, command, success: result.success, stdout: result.stdout, stderr: result.stderr, error: result.error, execution_time_ms: executionTime });
  res.json({
    success: result.success,
    output: result.stdout || '',
    stderr: result.stderr || '',
    error: result.error
  });
});

// ─── Streamable HTTP Transport (protocol 2025-11-25) ───

app.all('/mcp', auth, async (req, res) => {
  const sessionId = req.headers['mcp-session-id'];
  log(`StreamableHTTP ${req.method} from ${req.ip} session=${sessionId || 'none'}`);

  try {
    let transport;

    if (sessionId && activeSessions.has(sessionId)) {
      const existing = activeSessions.get(sessionId);
      if (existing instanceof StreamableHTTPServerTransport) {
        transport = existing;
      } else {
        return res.status(400).json({
          jsonrpc: '2.0',
          error: { code: -32000, message: 'Session uses different transport' },
          id: null
        });
      }
    } else if (!sessionId && req.method === 'POST' && isInitializeRequest(req.body)) {
      log(`StreamableHTTP new session (initialize)`);
      transport = new StreamableHTTPServerTransport({
        sessionIdGenerator: () => randomUUID(),
        onsessioninitialized: (sid) => {
          log(`StreamableHTTP session ready: ${sid}`);
          activeSessions.set(sid, transport);
        }
      });
      transport.onclose = () => {
        const sid = transport.sessionId;
        if (sid && activeSessions.has(sid)) {
          log(`StreamableHTTP closed: ${sid}`);
          activeSessions.delete(sid);
        }
      };
      const mcpServer = createMcpServer();
      await mcpServer.connect(transport);
    } else if (!sessionId && req.method === 'GET') {
      // Legacy SSE client connecting via GET /mcp
      return legacySSE(req, res);
    } else {
      return res.status(400).json({
        jsonrpc: '2.0',
        error: { code: -32000, message: 'Bad request: no valid session' },
        id: null
      });
    }

    await transport.handleRequest(req, res, req.body);
  } catch (err) {
    log(`StreamableHTTP ERROR: ${err.message}`);
    if (!res.headersSent) {
      res.status(500).json({
        jsonrpc: '2.0',
        error: { code: -32603, message: 'Internal server error' },
        id: null
      });
    }
  }
});

// ─── Legacy SSE Transport (protocol 2024-11-05) ───

async function legacySSE(req, res) {
  log(`Legacy SSE connection from ${req.ip}`);
  const mcpServer = createMcpServer();
  const transport = new SSEServerTransport(`${BASE_URL}/mcp/messages`, res);
  await mcpServer.connect(transport);
  activeSessions.set(transport.sessionId, transport);
  log(`Legacy SSE session ${transport.sessionId} ready`);

  res.on('close', () => {
    log(`Legacy SSE closed: ${transport.sessionId}`);
    activeSessions.delete(transport.sessionId);
  });
}

app.post('/mcp/messages', auth, async (req, res) => {
  const sessionId = req.query.sessionId;
  const method = req.body?.method || 'unknown';
  log(`Legacy POST ${method} for ${sessionId}`);

  const transport = activeSessions.get(sessionId);
  if (!transport || !(transport instanceof SSEServerTransport)) {
    log(`Legacy session not found: ${sessionId}`);
    return res.status(404).json({ error: "Session not found" });
  }

  try {
    await transport.handlePostMessage(req, res, req.body);
    log(`Legacy POST ${method} handled OK`);
  } catch (err) {
    log(`Legacy POST ${method} ERROR: ${err.message}`);
    console.error(err);
  }
});

app.listen(PORT, () => log(`Trinity Core MCP v2.4.0 started on port ${PORT}`));