--[[
	RoSocket - WebSockets for Roblox game servers over long-polling HTTP.
	Pyramus Studios - https://discord.pyramus.dev

	Put this ModuleScript in ServerScriptService (server-only; requires "Allow HTTP Requests").

	local RoSocket = require(game.ServerScriptService.RoSocket)
	RoSocket.configure({
		BaseUrl = "https://your-rosocket-host",
		Token = game:GetService("HttpService"):GetSecret("rosocket_token"), -- or a string
	})

	local socket = RoSocket.connect("wss://echo.websocket.org")
	socket:OnMessage(function(data, isBinary) print(data) end)
	socket:Send("hello")

	Each socket owns two hidden (unparented) BindableEvents:
	  socket.Inbound  - fired with ("open") | ("message", data, isBinary) | ("error", message) | ("close", code, reason)
	  socket.Outbound - :Fire(data, isBinary?) to send; pass it to other scripts that only need to send
]]

local HttpService = game:GetService("HttpService")
local RunService = game:GetService("RunService")

assert(RunService:IsServer(), "RoSocket can only be used from server scripts")

local RoSocket = {}

local config = {
	BaseUrl = nil,
	Token = nil,
	PollTimeout = 20,
	MaxBatch = 100,
}

local activeSockets = {}

function RoSocket.configure(options)
	assert(type(options) == "table", "RoSocket.configure expects a table")
	if options.BaseUrl ~= nil then
		config.BaseUrl = (string.gsub(options.BaseUrl, "/+$", ""))
	end
	if options.Token ~= nil then
		config.Token = options.Token
	end
	if options.PollTimeout ~= nil then
		config.PollTimeout = math.clamp(options.PollTimeout, 1, 25)
	end
end

local function authorization()
	if typeof(config.Token) == "Secret" then
		return config.Token:AddPrefix("Bearer ")
	end
	return "Bearer " .. tostring(config.Token)
end

-- Returns (success, statusCode, decodedBody). statusCode is 0 for transport errors.
local function request(method, path, body)
	assert(config.BaseUrl and config.Token, "Call RoSocket.configure({ BaseUrl = ..., Token = ... }) first")
	local options = {
		Url = config.BaseUrl .. "/api/v1" .. path,
		Method = method,
		Headers = { Authorization = authorization() },
	}
	if body ~= nil then
		options.Headers["Content-Type"] = "application/json"
		options.Body = HttpService:JSONEncode(body)
	end

	local ok, response = pcall(HttpService.RequestAsync, HttpService, options)
	if not ok then
		return false, 0, { error = tostring(response) }
	end

	local data = nil
	if response.Body and #response.Body > 0 then
		local decoded, result = pcall(HttpService.JSONDecode, HttpService, response.Body)
		if decoded then
			data = result
		end
	end
	return response.Success, response.StatusCode, data
end

local function errorText(status, data)
	return string.format("HTTP %d: %s", status, (type(data) == "table" and data.error) or "request failed")
end

local Socket = {}
Socket.__index = Socket

function Socket:_dispatch(event)
	if event.type == "open" then
		self.State = "open"
		self.Inbound:Fire("open")
	elseif event.type == "message" then
		self.Inbound:Fire("message", event.data, event.binary == true)
	elseif event.type == "error" then
		self.Inbound:Fire("error", event.message)
	elseif event.type == "close" then
		self:_finish(event.code, event.reason)
	end
end

function Socket:_finish(code, reason)
	if self.State == "closed" then
		return
	end
	self.State = "closed"
	activeSockets[self] = nil
	self.Inbound:Fire("close", code or 1006, reason or "")
	-- Give deferred handlers time to run before tearing the events down.
	task.delay(5, function()
		self.Inbound:Destroy()
		self.Outbound:Destroy()
	end)
end

function Socket:_pollLoop()
	local backoff = 1
	while self.State ~= "closed" do
		local path = string.format("/sockets/%s/poll?ack=%d&timeout=%d", self.Id, self._ack, config.PollTimeout)
		local ok, status, data = request("GET", path)
		if self.State == "closed" then
			break
		end

		if ok and type(data) == "table" then
			backoff = 1
			for _, event in ipairs(data.events or {}) do
				if event.seq > self._ack then
					self._ack = event.seq
					self:_dispatch(event)
				end
			end
			if data.closed then
				self:_finish(1006, "Socket closed")
			end
		elseif status == 401 or status == 403 or status == 404 then
			self.Inbound:Fire("error", errorText(status, data))
			self:_finish(1006, "Socket no longer available")
		else
			task.wait(backoff)
			backoff = math.min(backoff * 2, 30)
		end
	end
end

function Socket:_enqueue(data, isBinary)
	assert(type(data) == "string", "RoSocket: messages must be strings (JSONEncode tables first)")
	if self.State == "closed" then
		return
	end
	table.insert(self._outQueue, { data = data, binary = isBinary == true })
	if not self._sending then
		task.spawn(self._pump, self)
	end
end

function Socket:_pump()
	self._sending = true
	while #self._outQueue > 0 and self.State ~= "closed" do
		local batch = {}
		while #batch < config.MaxBatch and #self._outQueue > 0 do
			table.insert(batch, table.remove(self._outQueue, 1))
		end

		for attempt = 1, 3 do
			local ok, status, data = request("POST", "/sockets/" .. self.Id .. "/send", { messages = batch })
			if ok then
				break
			end
			-- 4xx won't succeed on retry; transport/5xx errors get two retries.
			if (status >= 400 and status < 500) or attempt == 3 then
				self.Inbound:Fire("error", "Send failed: " .. errorText(status, data))
				break
			end
			task.wait(attempt)
		end
	end
	self._sending = false
	if self._closing then
		self:_sendClose()
	end
end

function Socket:_sendClose()
	if self._closeSent then
		return
	end
	self._closeSent = true
	request("DELETE", "/sockets/" .. self.Id)
	self:_finish(1000, "Closed by game")
end

function Socket:Send(data, isBinary)
	self:_enqueue(data, isBinary)
end

function Socket:Close()
	if self.State == "closed" or self._closing then
		return
	end
	self._closing = true
	if not self._sending then
		task.spawn(self._sendClose, self)
	end
end

local function onKind(kind)
	return function(self, callback)
		return self.Inbound.Event:Connect(function(eventKind, ...)
			if eventKind == kind then
				callback(...)
			end
		end)
	end
end

Socket.OnOpen = onKind("open")
Socket.OnMessage = onKind("message")
Socket.OnError = onKind("error")
Socket.OnClose = onKind("close")

function RoSocket.connect(url, protocols)
	assert(type(url) == "string", "RoSocket.connect expects a ws:// or wss:// URL")
	local ok, status, data = request("POST", "/sockets", { url = url, protocols = protocols })
	if not ok or type(data) ~= "table" or not data.id then
		error("RoSocket: failed to open socket - " .. errorText(status, data), 2)
	end

	local self = setmetatable({
		Id = data.id,
		Url = data.url,
		State = "connecting",
		_ack = 0,
		_outQueue = {},
		_sending = false,
		_closing = false,
		_closeSent = false,
	}, Socket)

	-- Never parented, so they're invisible to the rest of the DataModel.
	self.Inbound = Instance.new("BindableEvent")
	self.Inbound.Name = "RoSocketInbound"
	self.Outbound = Instance.new("BindableEvent")
	self.Outbound.Name = "RoSocketOutbound"
	self.Outbound.Event:Connect(function(payload, isBinary)
		self:_enqueue(payload, isBinary)
	end)

	activeSockets[self] = true
	-- Deferred so callers can connect OnOpen/OnMessage before the first event arrives.
	task.defer(self._pollLoop, self)
	return self
end

game:BindToClose(function()
	local pending = 0
	for socket in pairs(activeSockets) do
		pending += 1
		task.spawn(function()
			request("DELETE", "/sockets/" .. socket.Id)
			socket:_finish(1001, "Server shutting down")
			pending -= 1
		end)
	end
	local deadline = os.clock() + 5
	while pending > 0 and os.clock() < deadline do
		task.wait()
	end
end)

return RoSocket