-- mcp.lua — MCP (Model Context Protocol) JSON-RPC 2.0 client.
-- Phase 2 v1: HTTP POST per RPC against lmcp servers; no long-lived SSE
-- channel (lmcp doesn't push — capabilities.tools.listChanged = false).
-- See docs/PHASE2.md §3 (module changes) and §4 (transport).

local curl = require("ffi.curl")
local json = require("dkjson")

local M = {}
local Session = {}
Session.__index = Session

local MCP_PROTOCOL_VERSION = "2025-03-26"

-- ---------------------------------------------------------------- M.connect
-- Open a session. No network traffic yet — call session:initialize()
-- to actually round-trip initialize + tools/list.
-- opts:
--   alias       short name for this server (defaults to URL hostname)
--   auth_token  literal Bearer token
--   auth_env    env-var name to read the token from (used if auth_token nil)
function M.connect(url, opts)
    opts = opts or {}
    local auth = opts.auth_token
    if (not auth or auth == "") and opts.auth_env then
        local env = os.getenv(opts.auth_env)
        if env and env ~= "" then auth = env end
    end
    return setmetatable({
        url             = url,
        alias           = opts.alias or url:match("https?://([^:/]+)") or url,
        auth            = auth,
        next_id         = 1,
        tools           = nil,   -- populated by initialize()
        server_info     = nil,
        server_caps     = nil,
        version_warning = nil,   -- non-nil string if server returned different protocolVersion
    }, Session)
end

-- ---------------------------------------------------------------- headers
function Session:_headers()
    local h = { "Content-Type: application/json", "Accept: application/json" }
    if self.auth and self.auth ~= "" then
        h[#h + 1] = "Authorization: Bearer " .. self.auth
    end
    return h
end

-- ---------------------------------------------------------------- _rpc
-- One round-trip. Returns:
--   result_table, "ok"                          — JSON-RPC success
--   nil,          "rpc_error",       error_obj  — JSON-RPC envelope error
--   nil,          "transport_error", msg        — HTTP >=400 / libcurl / parse
-- If has_id == false this is a notification: lmcp returns HTTP 202 empty
-- body and we synthesize (true, "ok") on transport success.
function Session:_rpc(method, params, has_id)
    local req = { jsonrpc = "2.0", method = method, params = params or {} }
    if has_id ~= false then
        req.id = self.next_id
        self.next_id = self.next_id + 1
    end
    local body, status = curl.post(self.url, json.encode(req), self:_headers())
    if not body then
        return nil, "transport_error", tostring(status)  -- 2nd slot is errmsg
    end
    if status >= 400 then
        return nil, "transport_error",
               ("HTTP %d: %s"):format(status, body:sub(1, 200))
    end
    if has_id == false then
        return true, "ok"
    end
    local doc, _, derr = json.decode(body)
    if not doc then
        return nil, "transport_error", "malformed JSON: " .. tostring(derr)
    end
    if doc.error then
        return nil, "rpc_error", doc.error
    end
    return doc.result or {}, "ok"
end

-- ---------------------------------------------------------------- initialize
-- Round-trips initialize + sends notifications/initialized + caches tools/list.
-- Returns:
--   true,  "ok"                          — session ready
--   false, kind, err                     — first failing RPC (caller logs)
function Session:initialize()
    local r, kind, err = self:_rpc("initialize", {
        protocolVersion = MCP_PROTOCOL_VERSION,
        capabilities    = {},
        clientInfo      = { name = "aish", version = "phase2" },
    })
    if not r then return false, kind, err end
    self.server_info = r.serverInfo
    self.server_caps = r.capabilities
    local sv = r.protocolVersion
    if sv and sv ~= MCP_PROTOCOL_VERSION then
        self.version_warning =
            ("protocol version mismatch (sent %s, got %s); proceeding")
            :format(MCP_PROTOCOL_VERSION, tostring(sv))
    end

    -- notifications/initialized — fire-and-forget; failure non-fatal.
    self:_rpc("notifications/initialized", nil, false)

    -- Eagerly fetch tools (cache for session lifetime per
    -- capabilities.tools.listChanged = false).
    local tr, tkind, terr = self:_rpc("tools/list", {})
    if not tr then return false, tkind, terr end
    self.tools = tr.tools or {}
    return true, "ok"
end

-- ---------------------------------------------------------------- list_tools
-- Cached. Returns the tool list captured at initialize() time;
-- empty table if not initialized.
function Session:list_tools()
    return self.tools or {}
end

-- ---------------------------------------------------------------- call_tool
-- Returns:
--   result_table, "ok"                          — tool succeeded (content[])
--   result_table, "handler_error"               — tool ran but result.isError = true
--                                                  (caller passes content through
--                                                  to the model regardless;
--                                                  PHASE2-baseline.md §3 also
--                                                  notes isError may be false on
--                                                  actual failure — content is
--                                                  authoritative)
--   nil,          "rpc_error",       error_obj  — JSON-RPC envelope error
--   nil,          "transport_error", msg        — HTTP/libcurl/parse failure
function Session:call_tool(name, args)
    local r, kind, err = self:_rpc("tools/call",
                                   { name = name, arguments = args or {} })
    if not r then return nil, kind, err end
    if r.isError then return r, "handler_error" end
    return r, "ok"
end

-- ---------------------------------------------------------------- close
-- Drops cached state. lmcp has no session teardown — every RPC was
-- already Connection: close.
function Session:close()
    self.tools           = nil
    self.server_info     = nil
    self.server_caps     = nil
    self.version_warning = nil
end

return M