-- idalib_mcp.lua -- launch, health-check, or stop the IDA "idalib-mcp" supervisor.
-- Registers the tool "idalib_mcp". Logs live under %TEMP%\idalib-mcp.
return function(registry)

local LOG_DIR = (os.getenv("TEMP") or ".") .. "\\idalib-mcp"
local DEFAULT_PORT = 8745

local function valid_host(host)
    return type(host) == "string" and host:match("^[%w%.%-]+$") ~= nil
end

local function valid_port(port)
    local n = tonumber(port)
    return n ~= nil and n >= 1 and n <= 65535 and n == math.floor(n)
end

-- Run a command and return its stdout (stderr merged), or nil.
local function sh(command)
    local pipe = io.popen(command .. " 2>nul")
    if not pipe then return nil end
    local output = pipe:read("*a")
    pipe:close()
    return output
end

-- POST one MCP "initialize" to the endpoint. Returns healthy, http_code, body.
local function probe(host, port)
    local body, out = LOG_DIR .. "\\probe-" .. port .. ".json", LOG_DIR .. "\\probe-" .. port .. ".out"
    local f = io.open(body, "w")
    if not f then return nil, nil, "cannot write probe body" end
    f:write('{"jsonrpc":"2.0","id":1,"method":"initialize","params":{'
        .. '"protocolVersion":"2025-06-18","capabilities":{},'
        .. '"clientInfo":{"name":"girl-idalib-tool","version":"1.0"}}}')
    f:close()
    local cmd = 'curl.exe -s -m 3 -o "' .. out .. '" -w "%{http_code}" --noproxy 127.0.0.1 -X POST '
        .. '"http://' .. host .. ':' .. port .. '/mcp" '
        .. '-H "Content-Type: application/json" '
        .. '-H "Accept: application/json, text/event-stream" '
        .. '--data-binary @"' .. body .. '"'
    local pipe = io.popen(cmd)
    local code = pipe and pipe:read("*a") or ""
    if pipe then pipe:close() end
    local http = tonumber((code:gsub("%s", "")))
    local g = io.open(out, "r")
    local resp = g and g:read("*a") or ""
    if g then g:close() end
    if http == nil then return false, nil, resp end
    local healthy = http == 200 and resp:find("ida-pro-mcp", 1, true) ~= nil
    return healthy, http, resp
end

-- Find python/idalib-mcp processes whose command line serves this port.
local function find_processes(port)
    local text = sh('wmic process where "name=\'python.exe\' or name=\'idalib-mcp.exe\'" '
        .. 'get ProcessId,Name,CommandLine /format:list') or ""
    local wanted = tostring(port)
    local procs, fields = {}, nil
    local function flush()
        if not fields then return end
        local pid, name, cmd = fields.ProcessId, fields.Name, fields.CommandLine
        fields = nil
        if pid and name and cmd
            and (name == "python.exe" or name == "idalib-mcp.exe")
            and cmd:find("idalib-mcp", 1, true) then
            local flag_port = cmd:match("%-%-port%s*(%d+)")
            local serves = (flag_port == nil and wanted == tostring(DEFAULT_PORT))
                or flag_port == wanted
            if serves then
                procs[#procs + 1] = { pid = tonumber(pid) or pid, name = name, cmd = cmd }
            end
        end
    end
    for line in (text .. "\n"):gmatch("([^\r\n]+)") do
        if line:match("^%s*$") then
            flush()
        else
            local key, value = line:match("^%s*([%w]+)%s*=(.*)$")
            if key then
                fields = fields or {}
                fields[key] = value
            end
        end
    end
    flush()
    table.sort(procs, function(a, b) return a.pid < b.pid end)
    return procs
end

local function log_path_for(port)
    return LOG_DIR .. "\\server-" .. os.date("%Y%m%d-%H%M%S") .. "-" .. port .. ".log"
end

-- Detached hidden launch; returns only after ~0.1s, server boots in background.
local function spawn(host, port, binary, logpath)
    local inner = "idalib-mcp --host " .. host .. " --port " .. port
    if binary and binary ~= "" then inner = inner .. ' "' .. binary .. '"' end
    local cmd = 'cmd /c start "" /b cmd /c "(' .. inner .. ') >> ' .. logpath .. ' 2>&1"'
    local ok, why, code = os.execute(cmd)
    return ok, why, code
end

local function wait_ready(host, port, seconds)
    local attempts = math.max(2, math.ceil((tonumber(seconds) or 30) * 1000 / 250))
    local healthy, http, body
    for _ = 1, attempts do
        healthy, http, body = probe(host, port)
        if healthy then return true, http end
        sleep(250)
    end
    return false, http, body
end

local function tail_newest_log(lines)
    local listing = sh('cmd /c dir /b /o-n "' .. LOG_DIR .. '\\server-*.log"') or ""
    local newest = (listing:gmatch("([^\r\n]+)")())
    if not newest or newest == "" then return nil, "no server logs yet" end
    local f = io.open(LOG_DIR .. "\\" .. newest, "r")
    if not f then return nil, "cannot open log" end
    local all = f:read("*a")
    f:close()
    local out, count = {}, 0
    for line in all:gmatch("([^\r\n]*)\r?\n") do
        count = count + 1
        out[#out + 1] = line
    end
    local n = math.max(1, math.min(tonumber(lines) or 40, 500))
    local first = math.max(1, count - n + 1)
    return table.concat(out, "\n", first, count), newest
end

registry:register {
    name = "idalib_mcp",
    description = "Ensure the IDA idalib-mcp MCP supervisor is running and reachable "
        .. "(default http://127.0.0.1:8745/mcp). action=status (default) probes MCP initialize. "
        .. "start spawns it detached (hidden, log under %TEMP%\\idalib-mcp) and waits until ready. "
        .. "stop kills the instance serving the port. restart = stop + start. "
        .. "log tails the newest server log. Optional: port, host, binary (input path to open), "
        .. "wait_seconds (start/restart), lines (log tail).",
    input_schema = {
        type = "object",
        properties = {
            action = { type = "string", enum = { "status", "start", "stop", "restart", "log" },
                description = "What to do; defaults to status." },
            port = { type = "integer", description = "TCP port; default 8745." },
            host = { type = "string", description = "Probe/bind host; default 127.0.0.1." },
            binary = { type = "string", description = "Optional binary for idalib-mcp to open on startup." },
            wait_seconds = { type = "integer", description = "Max wait for readiness on start/restart; default 30." },
            lines = { type = "integer", description = "Tail length for action=log; default 40." },
        },
    },
    permission = "shell",
    resource = function(args) return "idalib-mcp:" .. tostring(args.action or "status") end,
    run = function(args)
        args = args or {}
        local action = tostring(args.action or "status")
        local port = args.port or DEFAULT_PORT
        local host = args.host or "127.0.0.1"
        local url = "http://" .. host .. ":" .. port .. "/mcp"
        if not valid_host(host) then error("invalid host: " .. tostring(host), 0) end
        if not valid_port(port) then error("invalid port: " .. tostring(port), 0) end
        if action == "log" then
            local text, name = tail_newest_log(args.lines)
            if not text then return { ok = false, action = "log", message = name } end
            return { ok = true, action = "log", log = name, tail = text }
        end

        local healthy, http, body = probe(host, port)

        if action == "status" then
            return { ok = true, action = "status", url = url, ready = healthy == true,
                http = http, pids = find_processes(port) }
        end

        if action == "start" or action == "restart" then
            if action == "restart" and (healthy or (http ~= nil and http > 0)) then
                local procs = find_processes(port)
                for _, p in ipairs(procs) do
                    local k = io.popen('taskkill /PID ' .. p.pid .. ' /F /T 2>nul')
                    if k then k:read("*a"); k:close() end
                end
                sleep(500)
                healthy, http, body = probe(host, port)
            end
            if healthy then
                return { ok = true, action = "start", url = url, ready = true, already_running = true,
                    pids = find_processes(port) }
            end
            if http ~= nil and http > 0 then
                return { ok = false, action = "start", url = url, ready = false,
                    error = "port " .. port .. " is occupied by a different HTTP service",
                    http = http, hint = body:sub(1, 200) }
            end
            if args.binary and args.binary ~= "" and not io.open(args.binary, "rb") then
                return { ok = false, action = "start", url = url, ready = false,
                    error = "binary not found: " .. tostring(args.binary) }
            end
            os.execute('if not exist "' .. LOG_DIR .. '" mkdir "' .. LOG_DIR .. '" >nul 2>&1')
            local logpath = log_path_for(port)
            spawn(host, port, args.binary, logpath)
            local ready, last_http, last_body = wait_ready(host, port, args.wait_seconds)
            if not ready then
                local text, name = tail_newest_log(25)
                return { ok = false, action = "start", url = url, ready = false,
                    error = "server did not become ready within "
                        .. tostring(args.wait_seconds or 30) .. "s",
                    http = last_http, log = name, log_tail = text }
            end
            local procs = find_processes(port)
            local meta = io.open((logpath:gsub("%.log$", ".meta")), "w")
            if meta then
                meta:write("url=" .. url .. "\nstarted=" .. os.date("%Y-%m-%d %H:%M:%S") .. "\n"
                    .. "log=" .. logpath .. "\n"
                    .. "pids=" .. table.concat(
                        (function() local t = {} for _, p in ipairs(procs) do t[#t + 1] = tostring(p.pid) end return t end)(),
                        ",") .. "\n")
                meta:close()
            end
            return { ok = true, action = "start", url = url, ready = true,
                log = logpath, pids = procs }
        end

        if action == "stop" then
            local procs = find_processes(port)
            if #procs == 0 and (http == nil or http == 0) then
                return { ok = true, action = "stop", url = url, stopped = {}, already_down = true }
            end
            local stopped, details = {}, {}
            for _, p in ipairs(procs) do
                local k = io.popen('taskkill /PID ' .. p.pid .. ' /F /T 2>&1')
                local out = k and k:read("*a") or ""
                if k then k:close() end
                details[#details + 1] = { pid = p.pid, name = p.name, taskkill = (out:gsub("%s+", " ")) }
                stopped[#stopped + 1] = p.pid
            end
            sleep(400)
            local down = select(1, probe(host, port))
            return { ok = true, action = "stop", url = url, stopped = stopped,
                details = details, ready_after = down == true }
        end

        error("unknown action: " .. action, 0)
    end,
}

end