#!/usr/bin/env node

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import {
  CallToolRequestSchema,
  ListToolsRequestSchema,
} from '@modelcontextprotocol/sdk/types.js';
import express from 'express';
import cors from 'cors';
import { WebSocketServer } from 'ws';

// Express server for Roblox Studio plugin communication
const app = express();
app.use(cors());
app.use(express.json());

const HTTP_PORT = 37423;
const WS_PORT = 37424;

// Store connected Roblox Studio instances
let studioClients = new Set();
let pendingRequests = new Map();
let requestIdCounter = 0;

// HTTP polling support (alternative to WebSocket)
let pendingCommands = [];
let commandResults = new Map();
let commandIdCounter = 0;

// Create MCP server
const server = new Server(
  {
    name: 'roblox-studio-server',
    version: '1.0.0',
  },
  {
    capabilities: {
      tools: {},
    },
  }
);

// Helper function to send command to Roblox Studio and wait for response
async function sendToStudio(command, params = {}) {
  return new Promise((resolve, reject) => {
    const requestId = ++requestIdCounter;
    const commandId = ++commandIdCounter;

    // Set up response handler (supports both WebSocket and HTTP polling)
    const timeout = setTimeout(() => {
      pendingRequests.delete(requestId);
      commandResults.delete(commandId);
      reject(new Error('Request timeout - Roblox Studio did not respond'));
    }, 30000);

    // Store in both maps for WebSocket and HTTP compatibility
    pendingRequests.set(requestId, { resolve, reject, timeout });
    commandResults.set(commandId, { resolve, reject, timeout });

    // Add command to HTTP polling queue
    pendingCommands.push({ id: commandId, requestId, command, params, timestamp: Date.now() });

    // Try to send via WebSocket if available
    if (studioClients.size > 0) {
      const message = JSON.stringify({ id: commandId, requestId, command, params });
      studioClients.forEach(ws => {
        if (ws.readyState === ws.OPEN) {
          ws.send(message);
        }
      });
      console.error(`[MCP] Sent command ${command} (WS ID: ${commandId})`);
    } else {
      console.error(`[MCP] Queued command ${command} (HTTP ID: ${commandId}) - waiting for Roblox Studio to poll`);
    }
  });
}

// List available tools
server.setRequestHandler(ListToolsRequestSchema, async () => {
  return {
    tools: [
      {
        name: 'roblox_create_script',
        description: 'Create a new Lua script in Roblox Studio',
        inputSchema: {
          type: 'object',
          properties: {
            path: {
              type: 'string',
              description: 'Hierarchy path (e.g., "Workspace.Part.Script")',
            },
            scriptName: {
              type: 'string',
              description: 'Name of the script',
            },
            scriptType: {
              type: 'string',
              enum: ['Script', 'LocalScript', 'ModuleScript'],
              description: 'Type of script to create',
              default: 'Script',
            },
            source: {
              type: 'string',
              description: 'Lua source code',
            },
          },
          required: ['path', 'scriptName', 'source'],
        },
      },
      {
        name: 'roblox_create_part',
        description: 'Create a 3D part in the workspace',
        inputSchema: {
          type: 'object',
          properties: {
            parentPath: {
              type: 'string',
              description: 'Parent path (e.g., "Workspace" or "Workspace.Model")',
              default: 'Workspace',
            },
            partName: {
              type: 'string',
              description: 'Name of the part',
            },
            partType: {
              type: 'string',
              enum: ['Ball', 'Block', 'Cylinder', 'Wedge', 'CornerWedge'],
              description: 'Shape of the part',
              default: 'Block',
            },
            position: {
              type: 'object',
              properties: {
                x: { type: 'number' },
                y: { type: 'number' },
                z: { type: 'number' },
              },
              description: 'Position in 3D space',
            },
            size: {
              type: 'object',
              properties: {
                x: { type: 'number' },
                y: { type: 'number' },
                z: { type: 'number' },
              },
              description: 'Size of the part',
            },
            anchored: {
              type: 'boolean',
              description: 'Whether the part is anchored',
              default: true,
            },
            color: {
              type: 'string',
              description: 'BrickColor name (e.g., "Bright red")',
            },
          },
          required: ['partName'],
        },
      },
      {
        name: 'roblox_create_model',
        description: 'Create a new model container',
        inputSchema: {
          type: 'object',
          properties: {
            parentPath: {
              type: 'string',
              description: 'Parent path (e.g., "Workspace")',
              default: 'Workspace',
            },
            modelName: {
              type: 'string',
              description: 'Name of the model',
            },
          },
          required: ['modelName'],
        },
      },
      {
        name: 'roblox_set_property',
        description: 'Set a property on an existing object',
        inputSchema: {
          type: 'object',
          properties: {
            path: {
              type: 'string',
              description: 'Full path to the object',
            },
            property: {
              type: 'string',
              description: 'Property name to set',
            },
            value: {
              description: 'Property value (can be string, number, boolean, or object)',
            },
          },
          required: ['path', 'property', 'value'],
        },
      },
      {
        name: 'roblox_get_hierarchy',
        description: 'Get the hierarchy of objects in a path',
        inputSchema: {
          type: 'object',
          properties: {
            path: {
              type: 'string',
              description: 'Path to query (e.g., "Workspace")',
              default: 'Workspace',
            },
            depth: {
              type: 'number',
              description: 'How many levels deep to explore',
              default: 2,
            },
          },
          required: ['path'],
        },
      },
      {
        name: 'roblox_delete_object',
        description: 'Delete an object by path',
        inputSchema: {
          type: 'object',
          properties: {
            path: {
              type: 'string',
              description: 'Full path to the object to delete',
            },
          },
          required: ['path'],
        },
      },
      {
        name: 'roblox_play',
        description: 'Start playtest in Roblox Studio',
        inputSchema: {
          type: 'object',
          properties: {
            mode: {
              type: 'string',
              enum: ['Client', 'Server', 'Both'],
              description: 'Play mode',
              default: 'Both',
            },
          },
        },
      },
      {
        name: 'roblox_stop',
        description: 'Stop the current playtest',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'roblox_save_place',
        description: 'Save the current place',
        inputSchema: {
          type: 'object',
          properties: {},
        },
      },
      {
        name: 'roblox_execute_code',
        description: 'Execute arbitrary Lua code in the command bar',
        inputSchema: {
          type: 'object',
          properties: {
            code: {
              type: 'string',
              description: 'Lua code to execute',
            },
            context: {
              type: 'string',
              enum: ['Server', 'Client', 'Plugin'],
              description: 'Execution context',
              default: 'Plugin',
            },
          },
          required: ['code'],
        },
      },
      {
        name: 'roblox_create_folder',
        description: 'Create a folder object for organization',
        inputSchema: {
          type: 'object',
          properties: {
            parentPath: {
              type: 'string',
              description: 'Parent path (e.g., "Workspace")',
              default: 'Workspace',
            },
            folderName: {
              type: 'string',
              description: 'Name of the folder',
            },
          },
          required: ['folderName'],
        },
      },
      {
        name: 'roblox_create_gui',
        description: 'Create a basic GUI element (ScreenGui, Frame, TextButton, etc.)',
        inputSchema: {
          type: 'object',
          properties: {
            parentPath: {
              type: 'string',
              description: 'Parent path (e.g., "PlayerGui" or "StarterGui")',
              default: 'StarterGui',
            },
            guiType: {
              type: 'string',
              enum: ['ScreenGui', 'Frame', 'TextButton', 'TextLabel', 'TextBox', 'ImageLabel', 'ImageButton', 'ScrollingFrame'],
              description: 'Type of GUI element',
            },
            name: {
              type: 'string',
              description: 'Name of the GUI element',
            },
            properties: {
              type: 'object',
              description: 'Properties to set on the GUI element (size, position, text, color, etc.)',
            },
          },
          required: ['guiType', 'name'],
        },
      },
      {
        name: 'roblox_import_glb',
        description: 'Import a GLB 3D model into Roblox Studio',
        inputSchema: {
          type: 'object',
          properties: {
            glbData: {
              type: 'string',
              description: 'Base64-encoded GLB model data',
            },
            parentPath: {
              type: 'string',
              description: 'Parent path (e.g., "Workspace")',
              default: 'Workspace',
            },
            modelName: {
              type: 'string',
              description: 'Name for the imported model',
            },
          },
          required: ['glbData', 'modelName'],
        },
      },
    ],
  };
});

// Handle tool execution
server.setRequestHandler(CallToolRequestSchema, async (request) => {
  const { name, arguments: args } = request.params;

  console.error(`[MCP] Tool called: ${name}`, JSON.stringify(args));

  try {
    let result;

    switch (name) {
      case 'roblox_create_script':
        result = await sendToStudio('createScript', {
          path: args.path,
          scriptName: args.scriptName,
          scriptType: args.scriptType || 'Script',
          source: args.source,
        });
        break;

      case 'roblox_create_part':
        result = await sendToStudio('createPart', {
          parentPath: args.parentPath || 'Workspace',
          partName: args.partName,
          partType: args.partType || 'Block',
          position: args.position,
          size: args.size,
          anchored: args.anchored !== undefined ? args.anchored : true,
          color: args.color,
        });
        break;

      case 'roblox_create_model':
        result = await sendToStudio('createModel', {
          parentPath: args.parentPath || 'Workspace',
          modelName: args.modelName,
        });
        break;

      case 'roblox_set_property':
        result = await sendToStudio('setProperty', {
          path: args.path,
          property: args.property,
          value: args.value,
        });
        break;

      case 'roblox_get_hierarchy':
        result = await sendToStudio('getHierarchy', {
          path: args.path || 'Workspace',
          depth: args.depth || 2,
        });
        break;

      case 'roblox_delete_object':
        result = await sendToStudio('deleteObject', {
          path: args.path,
        });
        break;

      case 'roblox_play':
        result = await sendToStudio('play', {
          mode: args.mode || 'Both',
        });
        break;

      case 'roblox_stop':
        result = await sendToStudio('stop', {});
        break;

      case 'roblox_save_place':
        result = await sendToStudio('savePlace', {});
        break;

      case 'roblox_execute_code':
        result = await sendToStudio('executeCode', {
          code: args.code,
          context: args.context || 'Plugin',
        });
        break;

      case 'roblox_create_folder':
        result = await sendToStudio('createFolder', {
          parentPath: args.parentPath || 'Workspace',
          folderName: args.folderName,
        });
        break;

      case 'roblox_create_gui':
        result = await sendToStudio('createGUI', {
          parentPath: args.parentPath || 'StarterGui',
          guiType: args.guiType,
          name: args.name,
          properties: args.properties || {},
        });
        break;

      case 'roblox_import_glb':
        result = await sendToStudio('importGLB', {
          glbData: args.glbData,
          parentPath: args.parentPath || 'Workspace',
          modelName: args.modelName,
        });
        break;

      default:
        throw new Error(`Unknown tool: ${name}`);
    }

    return {
      content: [
        {
          type: 'text',
          text: JSON.stringify(result, null, 2),
        },
      ],
    };
  } catch (error) {
    return {
      content: [
        {
          type: 'text',
          text: JSON.stringify({
            success: false,
            error: error.message,
          }, null, 2),
        },
      ],
      isError: true,
    };
  }
});

// WebSocket server for real-time communication with Roblox Studio
const wss = new WebSocketServer({ port: WS_PORT });

wss.on('connection', (ws) => {
  console.error(`[WS] Roblox Studio connected!`);
  studioClients.add(ws);

  ws.on('message', (data) => {
    try {
      const message = JSON.parse(data);
      console.error(`[WS] Received:`, message);

      // Handle responses to pending requests
      if (message.id && pendingRequests.has(message.id)) {
        const { resolve, reject, timeout } = pendingRequests.get(message.id);
        clearTimeout(timeout);
        pendingRequests.delete(message.id);

        if (message.error) {
          reject(new Error(message.error));
        } else {
          resolve(message.data || message);
        }
      }
    } catch (e) {
      console.error(`[WS] Error parsing message:`, e);
    }
  });

  ws.on('close', () => {
    console.error(`[WS] Roblox Studio disconnected`);
    studioClients.delete(ws);
  });

  ws.on('error', (e) => {
    console.error(`[WS] Error:`, e);
    studioClients.delete(ws);
  });
});

// Health check endpoint
app.get('/health', (req, res) => {
  res.json({
    status: 'ok',
    studioConnected: studioClients.size > 0 || pendingCommands.length > 0,
    connections: studioClients.size,
    pendingCommands: pendingCommands.length,
  });
});

// HTTP polling endpoint for Roblox Studio
app.get('/poll', (req, res) => {
  const lastId = parseInt(req.query.last || '0');

  // Filter commands newer than lastId
  const newCommands = pendingCommands.filter(cmd => cmd.id > lastId);

  // Clean up old commands (older than 5 minutes)
  const now = Date.now();
  pendingCommands = pendingCommands.filter(cmd => now - cmd.timestamp < 300000);

  res.json({
    commands: newCommands,
    lastId: commandIdCounter,
  });
});

// Result endpoint for Roblox Studio to send command results
app.post('/result', (req, res) => {
  const { id, result } = req.body;

  // Store result for pending request
  if (commandResults.has(id)) {
    const { resolve, reject, timeout } = commandResults.get(id);
    clearTimeout(timeout);
    commandResults.delete(id);

    if (result && result.success === false) {
      reject(new Error(result.error || 'Command failed'));
    } else {
      resolve(result);
    }
  } else {
    // Also check pendingRequests for WebSocket compatibility
    if (pendingRequests.has(id)) {
      const { resolve, reject, timeout } = pendingRequests.get(id);
      clearTimeout(timeout);
      pendingRequests.delete(id);

      if (result && result.success === false) {
        reject(new Error(result.error || 'Command failed'));
      } else {
        resolve(result);
      }
    }
  }

  res.json({ success: true });
});

// Start Express server
app.listen(HTTP_PORT, () => {
  console.error(`HTTP server listening on port ${HTTP_PORT}`);
  console.error(`WebSocket server listening on port ${WS_PORT}`);
  console.error(`Waiting for Roblox Studio connection...`);
});

// Start MCP server
async function main() {
  const transport = new StdioServerTransport();
  await server.connect(transport);
  console.error('Roblox Studio MCP server running on stdio');
}

main().catch((error) => {
  console.error('Fatal error:', error);
  process.exit(1);
});