--!strict
--[[
	Handler registry and the single error boundary around every command.

	Handlers raise failures with `Dispatch.fail(code, message, hint)`. The hint is
	forwarded verbatim to the agent, so it should say what to do next ("call find
	first to get a valid path"), not just what went wrong.
]]

local Dispatch = {}

export type Error = {
	code: string,
	message: string,
	hint: string?,
}

export type Handler = (params: { [string]: any }) -> any

local handlers: { [string]: Handler } = {}

--[[
	When the server stops waiting for each running command, keyed by the
	command's thread.

	The server drops a timed-out command it has not delivered yet, but one
	already delivered can still be sitting behind the mutation lock, or behind
	a slow editor write. Starting it then does work nobody will hear about --
	and the agent, told "timed out", may be retrying it at that moment. So the
	deadline travels with the command and is checked where work would begin.

	Both clocks are this machine's, so there is no skew to allow for. Weak keys,
	so a thread that dies mid-command does not leave its entry behind.
]]
local deadlines: { [thread]: number } = setmetatable({}, { __mode = "k" }) :: any

function Dispatch.deadline(): number?
	return deadlines[coroutine.running()]
end

--[[
	Refuses to go on once the command's deadline has passed. Pass the deadline
	explicitly from code that runs on another thread -- an engine callback, or
	a spawned task -- which the table above cannot see.
]]
function Dispatch.checkDeadline(explicit: number?)
	local deadline = explicit or Dispatch.deadline()
	if deadline and DateTime.now().UnixTimestampMillis >= deadline then
		Dispatch.fail(
			"DEADLINE_EXPIRED",
			"The request timed out before Studio started this step, so it was not done.",
			"Earlier steps of the same request may have applied. Check the current state before retrying."
		)
	end
end

function Dispatch.register(op: string, handler: Handler)
	handlers[op] = handler
end

--[[
	Registers every entry of `map` under the prefix, e.g. `registerAll("script",
	{ read = fn })` exposes "script.read".
]]
function Dispatch.registerAll(prefix: string, map: { [string]: Handler })
	for name, handler in map do
		handlers[prefix .. "." .. name] = handler
	end
end

--[[
	Raises a structured failure. Always call this rather than `error("...")` so
	the agent receives a code it can branch on and a hint it can act on.
]]
function Dispatch.fail(code: string, message: string, hint: string?): never
	error({ code = code, message = message, hint = hint }, 0)
end

export type Result = {
	id: string,
	ok: boolean,
	data: any?,
	error: Error?,
}

--[[
	Runs one command and returns the frame to POST back. Never throws: an
	unhandled error inside a handler becomes a HANDLER_ERROR result, because a
	crashed dispatch would leave the server's promise hanging until it times out.
]]
function Dispatch.invoke(id: string, op: string, params: { [string]: any }?, deadlineMs: number?): Result
	local handler = handlers[op]
	if not handler then
		return {
			id = id,
			ok = false,
			error = {
				code = "UNKNOWN_OP",
				message = string.format('This Studio plugin has no handler for "%s".', op),
				hint = "The plugin is older than the MCP server. Update the Studio MCP "
					.. "plugin to the version matching your server.",
			},
		}
	end

	local thread = coroutine.running()
	deadlines[thread] = deadlineMs
	local ok, result = pcall(function()
		Dispatch.checkDeadline()
		return handler(params or {})
	end)
	deadlines[thread] = nil
	if ok then
		return { id = id, ok = true, data = result }
	end

	if typeof(result) == "table" and (result :: any).code then
		return { id = id, ok = false, error = result :: Error }
	end

	return {
		id = id,
		ok = false,
		error = {
			code = "HANDLER_ERROR",
			message = string.format('"%s" failed: %s', op, tostring(result)),
		},
	}
end

return Dispatch