-- lmcp.lua — Lightweight MCP server in pure Lua
-- Zero external dependencies (uses built-in socket or luasocket)
-- SPDX-License-Identifier: MIT

local json = require('json')

local lmcp = {}
lmcp.__index = lmcp

-- Protocol constants
local MCP_VERSION = "2025-03-26"
local JSONRPC = "2.0"

function lmcp.new(name, opts)
    opts = opts or {}
    local self = setmetatable({}, lmcp)
    self.name = name or "lmcp"
    self.version = opts.version or "0.1.0"
    self.host = opts.host or "0.0.0.0"
    self.port = opts.port or 8080
    self.tools = {}
    self._session_id = nil
    return self
end

-- Register a tool
function lmcp:tool(name, description, params_schema, handler)
    self.tools[name] = {
        name = name,
        description = description,
        inputSchema = params_schema or { type = "object", properties = {} },
        handler = handler,
    }
    return self
end

-- JSON-RPC response helpers
local function jsonrpc_result(id, result)
    return json.encode({ jsonrpc = JSONRPC, id = id, result = result })
end

local function jsonrpc_error(id, code, message)
    return json.encode({
        jsonrpc = JSONRPC,
        id = id,
        error = { code = code, message = message },
    })
end

-- Handle a single JSON-RPC request
function lmcp:handle_request(req)
    local method = req.method
    local id = req.id  -- nil for notifications

    if method == "initialize" then
        self._session_id = self._session_id or tostring(os.time())
        return jsonrpc_result(id, {
            protocolVersion = MCP_VERSION,
            capabilities = {
                tools = { listChanged = false },
            },
            serverInfo = {
                name = self.name,
                version = self.version,
            },
        })

    elseif method == "notifications/initialized" then
        return nil  -- notification, no response

    elseif method == "ping" then
        return jsonrpc_result(id, {})

    elseif method == "tools/list" then
        local tool_list = {}
        for _, t in pairs(self.tools) do
            tool_list[#tool_list + 1] = {
                name = t.name,
                description = t.description,
                inputSchema = t.inputSchema,
            }
        end
        return jsonrpc_result(id, { tools = tool_list })

    elseif method == "tools/call" then
        local params = req.params or {}
        local tool_name = params.name
        local arguments = params.arguments or {}
        local tool = self.tools[tool_name]
        if not tool then
            return jsonrpc_error(id, -32601, "Tool not found: " .. tostring(tool_name))
        end
        local ok, result = pcall(tool.handler, arguments)
        if ok then
            local content
            if type(result) == "string" then
                content = {{ type = "text", text = result }}
            elseif type(result) == "table" and result.type then
                content = { result }
            elseif type(result) == "table" then
                content = {{ type = "text", text = json.encode(result) }}
            else
                content = {{ type = "text", text = tostring(result) }}
            end
            return jsonrpc_result(id, { content = content, isError = false })
        else
            return jsonrpc_result(id, {
                content = {{ type = "text", text = "Error: " .. tostring(result) }},
                isError = true,
            })
        end

    else
        return jsonrpc_error(id, -32601, "Method not found: " .. tostring(method))
    end
end

-- ---- HTTP Server (raw sockets) ----

local function parse_http_request(client)
    -- Read request line
    local line, err = client:receive('*l')
    if not line then return nil, err end

    local method, path, version = line:match('^(%S+)%s+(%S+)%s+(%S+)')
    if not method then return nil, 'bad request line' end

    -- Read headers
    local headers = {}
    while true do
        line, err = client:receive('*l')
        if not line or line == '' then break end
        local k, v = line:match('^(%S+):%s*(.*)')
        if k then headers[k:lower()] = v end
    end

    -- Read body
    local body = ''
    local content_length = tonumber(headers['content-length'] or 0)
    if content_length > 0 then
        body, err = client:receive(content_length)
        if not body then return nil, err end
    end

    return {
        method = method,
        path = path,
        version = version,
        headers = headers,
        body = body,
    }
end

local function send_response(client, status, headers, body)
    local parts = { string.format('HTTP/1.1 %s', status) }
    headers['Content-Length'] = tostring(#body)
    headers['Connection'] = 'close'
    for k, v in pairs(headers) do
        parts[#parts + 1] = k .. ': ' .. v
    end
    parts[#parts + 1] = ''
    parts[#parts + 1] = body
    client:send(table.concat(parts, '\r\n'))
end

local function send_sse_event(client, data)
    client:send('event: message\r\ndata: ' .. data .. '\r\n\r\n')
end

function lmcp:serve_request(client)
    client:settimeout(5)
    local req, err = parse_http_request(client)
    if not req then
        client:close()
        return
    end

    local path = req.path
    local accept = req.headers['accept'] or ''

    -- GET /mcp — SSE endpoint (for session establishment)
    if req.method == 'GET' and path:match('^/mcp') then
        -- SSE stream — send headers and keep alive briefly
        local sse_headers = {
            'HTTP/1.1 200 OK',
            'Content-Type: text/event-stream',
            'Cache-Control: no-cache',
            'Connection: keep-alive',
            'Access-Control-Allow-Origin: *',
        }
        client:send(table.concat(sse_headers, '\r\n') .. '\r\n\r\n')

        -- Send endpoint event pointing to POST /mcp
        local endpoint_data = json.encode({
            endpoint = '/mcp',
            sessionId = self._session_id or tostring(os.time()),
        })
        client:send('event: endpoint\r\ndata: ' .. endpoint_data .. '\r\n\r\n')

        -- Keep connection open briefly for any SSE messages
        client:settimeout(0.1)
        -- In a full implementation we'd keep this open for server-initiated messages
        -- For now, the POST endpoint handles request-response
        client:close()
        return
    end

    -- POST /mcp — JSON-RPC endpoint
    if req.method == 'POST' and path:match('^/mcp') then
        if req.body == '' then
            send_response(client, '400 Bad Request',
                { ['Content-Type'] = 'application/json' },
                jsonrpc_error(nil, -32700, 'Empty body'))
            client:close()
            return
        end

        local ok, rpc_req = pcall(json.decode, req.body)
        if not ok then
            send_response(client, '400 Bad Request',
                { ['Content-Type'] = 'application/json' },
                jsonrpc_error(nil, -32700, 'Parse error'))
            client:close()
            return
        end

        -- Handle request
        local response = self:handle_request(rpc_req)

        if response then
            -- Check if client accepts SSE
            if accept:find('text/event%-stream') then
                local sse_headers = {
                    'HTTP/1.1 200 OK',
                    'Content-Type: text/event-stream',
                    'Cache-Control: no-cache',
                    'Access-Control-Allow-Origin: *',
                    'Connection: close',
                }
                client:send(table.concat(sse_headers, '\r\n') .. '\r\n\r\n')
                send_sse_event(client, response)
            else
                send_response(client, '200 OK',
                    { ['Content-Type'] = 'application/json',
                      ['Access-Control-Allow-Origin'] = '*' },
                    response)
            end
        else
            -- Notification — no response body
            send_response(client, '202 Accepted',
                { ['Content-Type'] = 'application/json',
                  ['Access-Control-Allow-Origin'] = '*' },
                '')
        end
        client:close()
        return
    end

    -- OPTIONS (CORS preflight)
    if req.method == 'OPTIONS' then
        send_response(client, '204 No Content', {
            ['Access-Control-Allow-Origin'] = '*',
            ['Access-Control-Allow-Methods'] = 'GET, POST, OPTIONS',
            ['Access-Control-Allow-Headers'] = 'Content-Type, Accept',
        }, '')
        client:close()
        return
    end

    -- Fallback
    send_response(client, '404 Not Found',
        { ['Content-Type'] = 'text/plain' },
        'Not Found')
    client:close()
end

function lmcp:run()
    -- Try luasocket first, fall back to Lua 5.4+ built-in (if available)
    local socket
    local ok, sock = pcall(require, 'socket')
    if ok then
        socket = sock
    else
        error('luasocket required: install with "luarocks install luasocket" or your package manager')
    end

    local server = assert(socket.bind(self.host, self.port))
    server:settimeout(1)

    local addr, port = server:getsockname()
    io.stderr:write(string.format("lmcp: %s v%s listening on %s:%d/mcp\n",
        self.name, self.version, addr, port))

    local running = true
    -- Handle SIGINT gracefully (Lua doesn't have signal handlers,
    -- but the timeout-based accept loop means Ctrl+C works)
    while running do
        local client = server:accept()
        if client then
            local ok, err = pcall(self.serve_request, self, client)
            if not ok then
                io.stderr:write("lmcp: request error: " .. tostring(err) .. "\n")
                pcall(client.close, client)
            end
        end
    end
end

return lmcp