-- stalwart_tool_check: Kong's access phase, asking the tool check before a tool call goes on.
--
-- The request is an MCP JSON-RPC message, or a plain call {"tool": ..., "arguments": {...}}.
-- Only `tools/call` is a call: every other MCP message (initialize, tools/list, a
-- notification) goes on unchecked, since it runs nothing.
--
-- A call that carries its agent's receipt -- the agent's session history, ending on the
-- check that allowed this call, in the MCP `_meta.stalwart` or the X-Stalwart-Receipt
-- header (base64 JSON) -- is admitted on that check: `POST /v1/tool-receipts` says whether
-- the history verifies and its last entry allows exactly this call, and the router spends
-- the receipt on it: presented again, it is refused (`receipt_spent`), except for that same
-- call sent again where re-sending is safe. A call without one is
-- checked in full (`POST /v1/tool-checks`), its history in X-Stalwart-History and a
-- person's approval in X-Stalwart-Confirmation; the answer's entry and token come back in
-- X-Stalwart-Entry and X-Stalwart-History-Token, for the client to add to its history.
-- Kong keeps no state.
--
-- Both routes are the customer's own enterprise qls-router's, in their own network: a
-- configuration naming no endpoint refuses every call (503 `router_not_configured`).

local http = require("resty.http")
local cjson = require("cjson.safe")

-- An empty JSON array must stay one: the history's entries are signed over their exact
-- JSON, and `[]` re-encoded as `{}` would make a receipt that cannot verify.
cjson.decode_array_with_array_mt(true)

local M = {}

-- What each answer of the check does to the request: the decision table the recipe's test reads.
M.OUTCOMES = {
  allow = { status = nil },                -- the call goes on to the tool server
  confirm = { status = 409, code = "confirmation_required" },
  refuse = { status = 403, code = "refused" },
}

function M.decide(answer)
  local outcome = M.OUTCOMES[answer.answer] or M.OUTCOMES.refuse
  return outcome
end

-- The call a request makes, and the receipt it carries; nil for a message that calls nothing.
function M.call_of(body)
  if body.jsonrpc ~= nil then
    if body.method ~= "tools/call" then
      return nil
    end
    local params = body.params or {}
    local meta = params._meta or {}
    return { tool = params.name, arguments = params.arguments or {}, receipt = meta.stalwart }
  end
  return { tool = body.tool, arguments = body.arguments or {}, receipt = nil }
end

function M.access(conf)
  local body = cjson.decode(kong.request.get_raw_body() or "") or {}
  local call = M.call_of(body)
  if call == nil then
    return
  end
  local receipt = call.receipt
  if receipt == nil then
    local header = kong.request.get_header("X-Stalwart-Receipt")
    if header ~= nil then
      receipt = cjson.decode(ngx.decode_base64(header) or "")
    end
  end

  if conf.endpoint == nil then
    -- No router to ask: nothing goes through.
    return kong.response.exit(503, { code = "router_not_configured",
                                     message = "set STALWART_ROUTER_URL to your qls-router" })
  end

  local endpoint, request
  if receipt ~= nil and conf.receipt_endpoint ~= nil then
    endpoint = conf.receipt_endpoint
    -- No agent is sent: the router reads it from the history's signed entry.
    request = { tool = call.tool, arguments = call.arguments, session = receipt.session or {} }
  else
    endpoint = conf.endpoint
    local history = cjson.decode(kong.request.get_header("X-Stalwart-History") or "") or {}
    request = {
      tool = call.tool,
      arguments = call.arguments,
      agent = kong.request.get_header(conf.agent_header),
      session = { id = history.session_id or "", history = history.entries or {}, history_token = history.token or "" },
      confirmation = kong.request.get_header("X-Stalwart-Confirmation"),
    }
  end

  local client = http.new()
  local response, failure = client:request_uri(endpoint, {
    method = "POST",
    body = cjson.encode(request),
    headers = { ["Content-Type"] = "application/json", ["Authorization"] = "Bearer " .. (conf.key or "") },
  })
  if not response then
    -- The check could not be asked: nothing goes through.
    return kong.response.exit(503, { code = "check_unreachable", message = failure })
  end
  -- An answer that is not the check's -- an error status, a body that does not read -- refuses.
  local answer = (response.status == 200 and cjson.decode(response.body)) or { answer = "refuse", reason = "unreadable" }
  local outcome = M.decide(answer)
  if answer.entry ~= nil then
    kong.response.set_header("X-Stalwart-Entry", cjson.encode(answer.entry))
    kong.response.set_header("X-Stalwart-History-Token", answer.history_token or "")
  end
  if outcome.status then
    return kong.response.exit(outcome.status, {
      code = outcome.code, reason = answer.reason, reasons = answer.reasons, confirmation = answer.confirmation,
    })
  end
end

return M