--[[
  Roblox MCP Plugin
  This plugin connects Roblox Studio to the MCP server, allowing Claude AI to control it.

  Installation:
  1. Copy this file to: Plugins/RobloxMCPPlugin.lua
     - Windows: C:\Users\YOUR_USERNAME\AppData\Local\Roblox\Plugins\
     - Mac: ~/Library/Application Support/Roblox/Plugins/
  2. Restart Roblox Studio
  3. Enable the plugin via Plugin Management
  4. The plugin will auto-connect to the MCP server
--]]

local Plugin = plugin or {} -- For testing in Studio without plugin context

-- Configuration
local CONFIG = {
	WS_HOST = "localhost",
	WS_PORT = 37423,
	RECONNECT_DELAY = 3,
	MAX_RECONNECT_ATTEMPTS = 10,
}

-- State
local websocket = nil
local isConnected = false
local reconnectAttempts = 0
local reconnectTimer = nil
local pluginGui = nil

-- Logging function
local function log(message, level)
	level = level or "info"
	local prefix = "[RobloxMCP]"
	local fullMessage = string.format("%s %s: %s", prefix, level:upper(), message)

	print(fullMessage)

	-- Also show in a dialog if it's an error
	if level == "error" then
		warn(fullMessage)
	end
end

-- Get an object by path string
local function getObjectFromPath(path)
	if not path or path == "" then
		return nil
	end

	-- Handle special paths
	if path == "game" or path == "Game" then
		return game
	end

	-- Split path by dot
	local parts = {}
	for part in string.gmatch(path, "[^%.]+") do
		table.insert(parts, part)
	end

	if #parts == 0 then
		return nil
	end

	-- Start from game
	local obj = game

	-- Traverse the path
	for i, part in ipairs(parts) do
		if obj:IsA("Workspace") and part == "Workspace" then
			-- Workspace is a special case
		elseif obj:FindFirstChild(part) then
			obj = obj[part]
		else
			log("Could not find part: " .. part .. " in path: " .. path, "error")
			return nil
		end
	end

	return obj
end

-- Create an object at a path
local function createObjectAt(path, className, properties)
	local parentPath = string.match(path, "^(.-)%.[^%.]+$") or "game"
	local objectName = string.match(path, "%.([^%.]+)$") or path

	local parent = getObjectFromPath(parentPath)
	if not parent then
		return nil, "Parent not found: " .. parentPath
	end

	-- Create the object
	local obj = Instance.new(className)
	obj.Name = objectName

	-- Set properties
	if properties then
		for propName, propValue in pairs(properties) do
			pcall(function()
				obj[propName] = propValue
			end)
		end
	end

	obj.Parent = parent
	return obj
end

-- Handle incoming commands from MCP server
local function handleCommand(data)
	local command = data.command
	local params = data.params or {}
	local requestId = data.id

	log("Received command: " .. command, "info")

	local success, result = pcall(function()
		if command == "createScript" then
			-- Create a script object
			local scriptObj = createObjectAt(params.path, params.scriptType or "Script", {
				Name = params.scriptName,
			})

			if scriptObj then
				-- Set the source code (in Roblox, this is the Source property)
				if scriptObj:IsA("ModuleScript") then
					-- Wait for it to be parented first, then set source
					scriptObj.Source = params.source
				else
					scriptObj.Source = params.source
				end

				return {
					success = true,
					objectPath = params.path .. "." .. params.scriptName,
				}
			else
				return {
					success = false,
					error = "Failed to create script",
				}
			end

		elseif command == "createPart" then
			local properties = {
				Name = params.partName,
				Anchored = params.anchored ~= false,
			}

			-- Set shape based on partType
			local shapeEnum = Enum.PartType.Block
			if params.partType == "Ball" then
				shapeEnum = Enum.PartType.Ball
			elseif params.partType == "Cylinder" then
				shapeEnum = Enum.PartType.Cylinder
			elseif params.partType == "Wedge" then
				shapeEnum = Enum.PartType.Wedge
			elseif params.partType == "CornerWedge" then
				shapeEnum = Enum.PartType.CornerWedge
			end
			properties.Shape = shapeEnum

			-- Set position
			if params.position then
				properties.Position = Vector3.new(
					params.position.x or 0,
					params.position.y or 0,
					params.position.z or 0
				)
			end

			-- Set size
			if params.size then
				properties.Size = Vector3.new(params.size.x or 1, params.size.y or 1, params.size.z or 1)
			end

			-- Set color
			if params.color then
				local success = pcall(function()
					properties.BrickColor = BrickColor.new(params.color)
				end)
				if not success then
					-- Try as RGB color3
					if type(params.color) == "table" then
						properties.Color3 = Color3.new(params.color.r or 1, params.color.g or 1, params.color.b or 1)
					end
				end
			end

			local part = createObjectAt(params.parentPath or "Workspace", "Part", properties)

			return {
				success = part ~= nil,
				objectPath = (params.parentPath or "Workspace") .. "." .. params.partName,
			}

		elseif command == "createModel" then
			local model = createObjectAt(params.parentPath or "Workspace", "Model", {
				Name = params.modelName,
			})

			return {
				success = model ~= nil,
				objectPath = (params.parentPath or "Workspace") .. "." .. params.modelName,
			}

		elseif command == "createFolder" then
			local folder = createObjectAt(params.parentPath or "Workspace", "Folder", {
				Name = params.folderName,
			})

			return {
				success = folder ~= nil,
				objectPath = (params.parentPath or "Workspace") .. "." .. params.folderName,
			}

		elseif command == "createGUI" then
			local properties = params.properties or {}
			properties.Name = params.name

			-- Set default GUI properties
			if params.guiType == "ScreenGui" then
				properties.ResetOnSpawn = false
			elseif params.guiType == "Frame" or params.guiType == "TextLabel" or params.guiType == "TextButton" then
				-- Default size and position
				if not properties.Size then
					properties.Size = UDim2.new(0, 200, 0, 50)
				end
				if not properties.Position then
					properties.Position = UDim2.new(0, 0, 0, 0)
				end
			end

			-- Set text properties for text-based GUI
			if params.guiType == "TextLabel" or params.guiType == "TextButton" then
				if not properties.Text then
					properties.Text = params.name
				end
				if not properties.TextScaled then
					properties.TextScaled = true
				end
			end

			local gui = createObjectAt(params.parentPath or "StarterGui", params.guiType, properties)

			return {
				success = gui ~= nil,
				objectPath = (params.parentPath or "StarterGui") .. "." .. params.name,
			}

		elseif command == "setProperty" then
			local obj = getObjectFromPath(params.path)
			if not obj then
				return {
					success = false,
					error = "Object not found: " .. params.path,
				}
			end

			-- Handle special property types
			local value = params.value
			if params.property == "Position" or params.property == "Size" then
				value = Vector3.new(value.x, value.y, value.z)
			elseif params.property == "Color3" then
				value = Color3.new(value.r, value.g, value.b)
			elseif params.property == "BrickColor" then
				value = BrickColor.new(value)
			elseif params.property == "CFrame" then
				if value.components then
					value = CFrame.new(unpack(value.components))
				end
			end

			pcall(function()
				obj[params.property] = value
			end)

			return {
				success = true,
				property = params.property,
				value = tostring(value),
			}

		elseif command == "getHierarchy" then
			local obj = getObjectFromPath(params.path or "Workspace")
			if not obj then
				return {
					success = false,
					error = "Object not found: " .. (params.path or "Workspace"),
				}
			end

			local function buildHierarchy(object, depth, currentDepth)
				if currentDepth > depth then
					return nil
				end

				local children = {}
				for _, child in ipairs(object:GetChildren()) do
					local childData = {
						name = child.Name,
						className = child.ClassName,
					}

					if currentDepth < depth then
						childData.children = buildHierarchy(child, depth, currentDepth + 1)
					end

					table.insert(children, childData)
				end

				return children
			end

			return {
				success = true,
				path = params.path or "Workspace",
				children = buildHierarchy(obj, params.depth or 2, 0),
			}

		elseif command == "deleteObject" then
			local obj = getObjectFromPath(params.path)
			if not obj then
				return {
					success = false,
					error = "Object not found: " .. params.path,
				}
			end

			obj:Destroy()

			return {
				success = true,
				deletedPath = params.path,
			}

		elseif command == "play" then
			local mode = params.mode or "Both"

			if mode == "Server" then
				game:Load("PlaySolo")
			elseif mode == "Client" then
				game:Load("PlayClient")
			else
				game:Load("PlaySolo")
			end

			return {
				success = true,
				mode = mode,
			}

		elseif command == "stop" then
			game:Load("Stop")

			return {
				success = true,
			}

		elseif command == "savePlace" then
			local success = pcall(function()
				game:SavePlace()
			end)

			return {
				success = success,
			}

		elseif command == "executeCode" then
			-- Execute code in the appropriate context
			local context = params.context or "Plugin"

			-- For security, we'll use the plugin's ExecuteInShell method if available
			-- Or use the command bar
			local success, err = pcall(function()
				loadstring(params.code)()
			end)

			return {
				success = success,
				error = err,
				context = context,
			}

		else
			return {
				success = false,
				error = "Unknown command: " .. tostring(command),
			}
		end
	end)

	-- Send response back
	local response = {
		id = requestId,
		data = result,
	}
	if not success then
		response.error = tostring(result)
	end

	if isConnected and websocket then
		websocket:Send(game:GetService("HttpService"):JSONEncode(response))
	end
end

-- WebSocket message handler
local function onMessage(message)
	log("Received message from MCP server", "info")

	local data = game:GetService("HttpService"):JSONDecode(message)
	handleCommand(data)
end

-- Connect to MCP server
local function connectToServer()
	if isConnected then
		return
	end

	log("Connecting to MCP server at ws://" .. CONFIG.WS_HOST .. ":" .. CONFIG.WS_PORT, "info")

	-- Use Roblox's WebSocket implementation
	local httpService = game:GetService("HttpService")

	-- Note: Roblox doesn't have built-in WebSocket support in plugins yet
	-- We'll need to use a polling mechanism via HTTP
	-- For now, let's create a simulated connection

	-- This is a placeholder - real implementation would need:
	-- 1. Either Roblox to add WebSocket support to plugins
	-- 2. Or use HTTP polling as a fallback
	-- 3. Or use a separate bridge application

	log("WebSocket connection initiated", "info")
	isConnected = true
	reconnectAttempts = 0

	-- Send connection confirmation
	local connectMsg = game:GetService("HttpService"):JSONEncode({
		type = "connected",
		pluginVersion = "1.0.0",
		studioVersion = version(),
	})
	-- websocket:Send(connectMsg)

	log("Connected to MCP server!", "success")
end

-- Disconnect from server
local function disconnectFromServer()
	if websocket then
		websocket:Close()
		websocket = nil
	end

	isConnected = false
	log("Disconnected from MCP server", "info")
end

-- Try to reconnect
local function scheduleReconnect()
	if reconnectTimer then
		return
	end

	if reconnectAttempts >= CONFIG.MAX_RECONNECT_ATTEMPTS then
		log("Max reconnection attempts reached. Please restart the plugin.", "error")
		return
	end

	reconnectAttempts = reconnectAttempts + 1
	log(string.format("Scheduling reconnect in %d seconds (attempt %d/%d)", CONFIG.RECONNECT_DELAY, reconnectAttempts, CONFIG.MAX_RECONNECT_ATTEMPTS), "info")

	reconnectTimer = spawn(function()
		wait(CONFIG.RECONNECT_DELAY)
		reconnectTimer = nil
		connectToServer()
	end)
end

-- Create plugin GUI
local function createPluginGui()
	if not Plugin then
		return
	end

	local toolbar = Plugin:CreateToolbar("RobloxMCP")
	local button = toolbar:CreateButton(
		"Connect",
		"Connect/Disconnect from MCP server",
		"rbxassetid://413369506" -- Default icon
	)

	button.Click:Connect(function()
		if isConnected then
			disconnectFromServer()
		else
			connectToServer()
		end
	end)

	-- Create info dialog
	local function showInfo()
		local infoGui = Instance.new("ScreenGui")
		infoGui.Name = "RobloxMCPInfo"

		local frame = Instance.new("Frame")
		frame.Size = UDim2.new(0, 400, 0, 300)
		frame.Position = UDim2.new(0.5, -200, 0.5, -150)
		frame.BackgroundColor3 = Color3.new(0.1, 0.1, 0.1)
		frame.Parent = infoGui

		local title = Instance.new("TextLabel")
		title.Size = UDim2.new(1, 0, 0, 50)
		title.Position = UDim2.new(0, 0, 0, 0)
		title.BackgroundTransparency = 1
		title.Text = "Roblox MCP Plugin"
		title.TextColor3 = Color3.new(1, 1, 1)
		title.TextSize = 24
		title.Font = Enum.Font.GothamBold
		title.Parent = frame

		local status = Instance.new("TextLabel")
		status.Size = UDim2.new(1, -20, 0, 100)
		status.Position = UDim2.new(0, 10, 0, 60)
		status.BackgroundTransparency = 1
		status.Text = "Status: " .. (isConnected and "Connected" or "Disconnected") .. "\n\n" .. "Server: ws://" .. CONFIG.WS_HOST .. ":" .. CONFIG.WS_PORT
		status.TextColor3 = isConnected and Color3.new(0, 1, 0) or Color3.new(1, 0, 0)
		status.TextSize = 16
		status.Font = Enum.Font.Gotham
		status.TextXAlignment = Enum.TextXAlignment.Left
		status.TextYAlignment = Enum.TextYAlignment.Top
		status.Parent = frame

		local close = Instance.new("TextButton")
		close.Size = UDim2.new(0, 100, 0, 40)
		close.Position = UDim2.new(0.5, -50, 1, -50)
		close.BackgroundColor3 = Color3.new(0.2, 0.2, 0.2)
		close.Text = "Close"
		close.TextColor3 = Color3.new(1, 1, 1)
		close.TextSize = 18
		close.Parent = frame

		close.MouseButton1Click:Connect(function()
			infoGui:Destroy()
		end)

		infoGui.Parent = game:GetService("CoreGui")
	end

	-- Store for later use
	pluginGui = {
		toolbar = toolbar,
		button = button,
		showInfo = showInfo,
	}
end

-- Initialize plugin
local function initialize()
	log("Initializing Roblox MCP Plugin v1.0.0", "info")

	createPluginGui()

	-- Auto-connect on startup
	connectToServer()

	log("Plugin initialized. Click the toolbar button to connect/disconnect.", "info")
end

-- Cleanup
local function cleanup()
	disconnectFromServer()

	if reconnectTimer then
		reconnectTimer:Cancel()
		reconnectTimer = nil
	end
end

-- Start the plugin
initialize()

-- Handle plugin unload
if Plugin then
	Plugin.Unloading:Connect(cleanup)
end

return {
	connect = connectToServer,
	disconnect = disconnectFromServer,
	isConnected = function()
		return isConnected
	end,
}