-- Run a shell command and return its merged stdout+stderr output.
--
-- The command is executed via `/bin/sh -c <command>`, so shell
-- features (pipes, redirections, globs, `&&`, etc.) all work the
-- usual way. Stdin is closed; the command cannot read from the
-- terminal.
--
-- Context-budget discipline:
--   - Output is capped at 50 KB (matching `read`'s cap), keeping the
--     **last** 50 KB on overflow. Compiler errors, test failures, and
--     stack traces all put the important content at the end, so a
--     head-keeping shell tool routinely truncates exactly the bytes
--     the agent needed.
--   - Once overflow is detected, the *complete* transcript is
--     streamed to a spill file under `$PANTO_HOME/shell-output/`. The
--     agent can then `read` it with start_line/end_line to inspect
--     specific regions — including the head that fell out of the
--     tail buffer.
--
-- Timeout: defaults to 30s. When the timer fires the child gets
-- SIGTERM; if it still hasn't exited 250 ms later, SIGKILL.
--
-- Stdout and stderr are merged into a single stream (in arrival
-- order). Distinguishing them would require either prefixing every
-- chunk (visually noisy) or surfacing two separate strings (not how
-- a string-returning tool result is shaped). Most shell commands the
-- agent runs interleave the two anyway.

local uv = require("luv")

local MAX_BYTES = 50 * 1024   -- 50 KB, same as `read`
local DEFAULT_TIMEOUT_S = 300   -- 5 minutes
local SIGKILL_GRACE_MS = 250

-- ---------------------------------------------------------------------------
-- Coroutine-synchronous coordination.
--
-- panto runs every tool handler in its own coroutine and drives a
-- single `uv.run()` to completion; the only legal way for a handler to
-- block is to yield on pending libuv work whose callback(s) resume it.
-- The shell tool waits on THREE independent events before it can build
-- its result: stdout EOF, stderr EOF, and process exit. `await_n`
-- captures that "resume once, after N signals" pattern.
--
-- `arm(signal)` is called immediately with a `signal` function; it must
-- arrange for `signal` to be called exactly `n` times across its libuv
-- callbacks. After the n-th signal, the coroutine is resumed exactly
-- once. `signal` is safe to call from any libuv callback.
local function await_n(n, arm)
    local co = assert(coroutine.running(),
        "shell: await_n must run inside a tool handler coroutine")
    local remaining, resumed = n, false
    local function signal()
        if remaining <= 0 then return end
        remaining = remaining - 1
        if remaining > 0 or resumed then return end
        resumed = true
        local ok, err = coroutine.resume(co)
        if not ok then
            -- This runs inside a libuv callback, not the handler's
            -- own frame; re-raising here would unwind the event loop.
            -- Log instead; the scheduler surfaces the unfinished
            -- coroutine as a tool error on reap.
            io.stderr:write("shell tool: resume failed: " .. tostring(err) .. "\n")
        end
    end
    arm(signal)
    coroutine.yield()
end

-- ---------------------------------------------------------------------------
-- Spill-file path resolution. We mirror Zig's `panto_home.resolveHome`:
-- PANTO_HOME wins, then XDG_DATA_HOME/panto, then HOME/.local/share/panto.
-- The shell-output subdirectory is created lazily.
-- ---------------------------------------------------------------------------

local function panto_home()
    local p = os.getenv("PANTO_HOME")
    if p and p ~= "" then return p end
    local xdg = os.getenv("XDG_DATA_HOME")
    if xdg and xdg ~= "" then return xdg .. "/panto" end
    local home = os.getenv("HOME")
    if home and home ~= "" then return home .. "/.local/share/panto" end
    return nil
end

local function spill_dir()
    local base = panto_home()
    if not base then return nil end
    return base .. "/shell-output"
end

-- `mkdir -p`-equivalent via luv. We can't shell out to mkdir here —
-- that would be circular — so we walk the path components and call
-- `uv.fs_mkdir` on each, ignoring EEXIST. Synchronous; this is a
-- once-per-spill cost.
local function mkdir_p(path)
    -- Split on `/` and rebuild. Absolute paths start with "/"; preserve.
    local parts = {}
    for piece in string.gmatch(path, "[^/]+") do
        parts[#parts + 1] = piece
    end
    local cur = (path:sub(1, 1) == "/") and "" or "."
    for _, piece in ipairs(parts) do
        cur = cur .. "/" .. piece
        local ok, err, name = uv.fs_mkdir(cur, tonumber("755", 8))
        if not ok and name ~= "EEXIST" then
            return nil, err
        end
    end
    return true
end

local function fresh_spill_path()
    local dir = spill_dir()
    if not dir then return nil, "no PANTO_HOME / XDG_DATA_HOME / HOME in env" end
    local ok, err = mkdir_p(dir)
    if not ok then return nil, err end
    -- Filename: <unix-seconds>-<microseconds>-<pid>.txt. The pid is
    -- the parent panto process; collision avoidance against concurrent
    -- tool calls comes from the microsecond component (uv.hrtime() is
    -- nanoseconds since some epoch).
    local sec, usec = uv.gettimeofday()
    local pid = uv.os_getpid()
    return string.format("%s/%d-%06d-%d.txt", dir, sec, usec, pid)
end

-- ---------------------------------------------------------------------------
-- Tool body
-- ---------------------------------------------------------------------------

return {
    name = "std.shell",
    description = string.format("Execute a shell command via `/bin/sh -c`. Returns merged stdout+stderr after the command exits, prefixed with an exit-status header. Output is truncated to the last %dKB; on truncation the full transcript is saved to `$PANTO_HOME/shell-output/` and its path is included for follow-up `read` calls. Default timeout 5 minutes.", MAX_BYTES / 1024),
    schema = {
        type = "object",
        properties = {
            command = { type = "string", description = "Shell command. Passed to `sh -c` verbatim; pipes, redirections, etc. work." },
            cwd = { type = "string", description = "Working directory (default: cwd)." },
            timeout = { type = "integer", description = "Hard timeout in seconds (default: " .. tostring(DEFAULT_TIMEOUT_S) .. ").", minimum = 1 },
        },
        required = { "command" },
    },
    handler = function(input)
        local command = input.command
        if type(command) ~= "string" or command == "" then
            return "Error: `command` must be a non-empty string."
        end
        local cwd = input.cwd
        if cwd ~= nil and type(cwd) ~= "string" then
            return "Error: `cwd` must be a string if provided."
        end
        local timeout_s = input.timeout or DEFAULT_TIMEOUT_S
        local timeout_ms = timeout_s * 1000

        if not coroutine.running() then
            -- Should never happen — panto always runs handlers in a
            -- coroutine — but be defensive rather than silently hang.
            return "Error: shell tool requires a coroutine scheduler " ..
                "(internal bug; report to panto)."
        end

        -- Pipes for the child's stdout/stderr. Stdin is `nil` in the
        -- stdio table → child reads from /dev/null equivalent (luv
        -- closes that fd in the child).
        local stdout_pipe = uv.new_pipe(false)
        local stderr_pipe = uv.new_pipe(false)

        -- Output accumulator: a FIFO of arrival-order chunks whose
        -- total length is always ≤ MAX_BYTES. On overflow we evict
        -- from the front (oldest bytes first) so the in-memory buffer
        -- always holds the *tail* of the output. The full unabridged
        -- transcript is mirrored to a spill file, opened lazily on
        -- first overflow and pre-loaded with everything seen so far.
        local chunks = {}          -- FIFO of strings; total ≤ MAX_BYTES
        local first = 1            -- index of oldest live chunk in `chunks`
        local last = 0             -- index of newest live chunk
        local total_bytes = 0
        local total_received = 0   -- everything we ever saw, pre-eviction
        local overflowed = false   -- true once a byte has been evicted
        local spill_path = nil
        local spill_handle = nil   -- uv fd (integer) once opened
        local spill_error = nil    -- if opening/writing the spill file fails
        local total_written_to_spill = 0

        -- NOTE: the spill helpers below use SYNCHRONOUS `uv.fs_*`
        -- calls on purpose. They run inside `append_chunk`, which runs
        -- inside the stdout/stderr read callbacks — i.e. in event-loop
        -- context, NOT in the handler coroutine's own frame. You cannot
        -- `coroutine.yield` (and thus cannot `await`) from a libuv
        -- callback, so these must be the blocking variants. The brief
        -- per-write stall is acceptable; spilling only happens once
        -- output exceeds the in-memory cap.
        local function open_spill_and_seed()
            -- Opens the spill file and writes every chunk currently
            -- live in the FIFO. Called exactly once, the moment we
            -- realize we're going to overflow. After this returns,
            -- the spill file contains the entire history; every
            -- subsequent chunk just appends.
            if spill_handle or spill_error then return end
            local path, err = fresh_spill_path()
            if not path then
                spill_error = err
                return
            end
            -- O_WRONLY | O_CREAT | O_TRUNC; mode 0644.
            local fd, oerr = uv.fs_open(path, "w", tonumber("644", 8))
            if not fd then
                spill_error = oerr or "unknown error opening spill file"
                return
            end
            spill_path = path
            spill_handle = fd
            for i = first, last do
                local c = chunks[i]
                local n, werr = uv.fs_write(fd, c, -1)
                if not n then
                    spill_error = werr or "spill write failed"
                    return
                end
                total_written_to_spill = total_written_to_spill + n
            end
        end

        local function spill_write(data)
            if not spill_handle then return end
            local n, werr = uv.fs_write(spill_handle, data, -1)
            if not n then
                spill_error = spill_error or werr or "spill write failed"
            else
                total_written_to_spill = total_written_to_spill + n
            end
        end

        local function append_chunk(data)
            total_received = total_received + #data

            -- First time we'd cross MAX_BYTES: open the spill file
            -- and seed it with everything we've buffered so far. From
            -- here on we always dual-write (spill gets the unabridged
            -- stream; chunks holds only the tail).
            if not overflowed and total_bytes + #data > MAX_BYTES then
                overflowed = true
                open_spill_and_seed()
            end
            if overflowed then
                spill_write(data)
            end

            -- Append to the ring. If the new chunk alone is bigger
            -- than MAX_BYTES, drop everything we have and keep just
            -- its tail. Otherwise, evict whole front chunks until
            -- there's room, then either trim the front-most or
            -- append intact.
            if #data >= MAX_BYTES then
                -- Drop the whole FIFO; new chunk is the tail.
                for i = first, last do chunks[i] = nil end
                first = last + 1
                total_bytes = 0
                local trimmed = data:sub(#data - MAX_BYTES + 1)
                last = last + 1
                chunks[last] = trimmed
                total_bytes = #trimmed
                return
            end

            while total_bytes + #data > MAX_BYTES and first <= last do
                local front = chunks[first]
                if total_bytes - #front + #data <= MAX_BYTES then
                    -- Trim the front chunk: drop only as many bytes
                    -- as needed to make room.
                    local need_to_drop = (total_bytes + #data) - MAX_BYTES
                    chunks[first] = front:sub(need_to_drop + 1)
                    total_bytes = total_bytes - need_to_drop
                else
                    -- Drop the entire front chunk.
                    chunks[first] = nil
                    total_bytes = total_bytes - #front
                    first = first + 1
                end
            end
            last = last + 1
            chunks[last] = data
            total_bytes = total_bytes + #data
        end

        -- Per-call outcome state, filled by the libuv callbacks below.
        local exit_code = nil
        local exit_signal = nil
        local timed_out = false
        local process_handle = nil
        local timeout_timer = nil
        local kill_timer = nil

        local function buffer_to_string()
            local out = {}
            local oi = 0
            for i = first, last do
                oi = oi + 1
                out[oi] = chunks[i]
            end
            return table.concat(out)
        end

        -- Spawn first so a spawn failure can return *before* we commit
        -- to awaiting (a failed spawn fires no events to resume us).
        local stdio = { nil, stdout_pipe, stderr_pipe }
        local opts = { args = { "-c", command }, stdio = stdio }
        if cwd then opts.cwd = cwd end

        local on_exit_signal   -- set by await_n before any event fires
        local handle, pid_or_err, name = uv.spawn("/bin/sh", opts, function(code, sig)
            exit_code = code
            exit_signal = sig
            -- The on-exit callback owns closing the process handle.
            if process_handle and not process_handle:is_closing() then
                process_handle:close()
            end
            on_exit_signal()
        end)

        if not handle then
            -- Spawn failed before any event could fire. Clean up the
            -- pipes and return without awaiting.
            local spawn_error = tostring(pid_or_err) ..
                (name and (" (" .. name .. ")") or "")
            stdout_pipe:close(); stderr_pipe:close()
            return "Error: failed to spawn shell: " .. spawn_error
        end
        process_handle = handle

        -- `await_n(3, ...)` resumes this coroutine exactly once, after
        -- all three of {stdout EOF, stderr EOF, process exit} signal.
        await_n(3, function(signal)
            -- The exit callback was armed above (before this awaiter
            -- existed); wire its signal now.
            on_exit_signal = signal

            -- stdout/stderr each signal once, on EOF or read error.
            -- luv delivers stream EOF as a nil `data` with no error.
            local function on_read(err, data)
                if err then
                    append_chunk("[read error: " .. tostring(err) .. "]\n")
                    signal()
                elseif data == nil then
                    signal()
                else
                    append_chunk(data)
                end
            end
            stdout_pipe:read_start(on_read)
            stderr_pipe:read_start(on_read)

            -- Timeout timer. If it fires, send SIGTERM and arm a short
            -- follow-up timer for SIGKILL. These timers do NOT signal;
            -- they cause the child to exit, which fires the exit/EOF
            -- callbacks that do. We close them here so they stop
            -- holding the event loop open.
            timeout_timer = uv.new_timer()
            timeout_timer:start(timeout_ms, 0, function()
                timed_out = true
                if process_handle and not process_handle:is_closing() then
                    process_handle:kill("sigterm")
                end
                timeout_timer:stop(); timeout_timer:close(); timeout_timer = nil
                kill_timer = uv.new_timer()
                kill_timer:start(SIGKILL_GRACE_MS, 0, function()
                    if process_handle and not process_handle:is_closing() then
                        process_handle:kill("sigkill")
                    end
                    kill_timer:stop(); kill_timer:close(); kill_timer = nil
                end)
            end)
        end)

        if spawn_error then
            return "Error: failed to spawn shell: " .. spawn_error
        end

        -- Resumed: all three events fired. Cancel any still-pending
        -- timers (the command may have finished before the timeout) so
        -- they don't keep the loop alive, then close pipes/spill.
        if timeout_timer and not timeout_timer:is_closing() then
            timeout_timer:stop(); timeout_timer:close()
        end
        if kill_timer and not kill_timer:is_closing() then
            kill_timer:stop(); kill_timer:close()
        end
        if not stdout_pipe:is_closing() then stdout_pipe:close() end
        if not stderr_pipe:is_closing() then stderr_pipe:close() end
        if spill_handle then uv.fs_close(spill_handle) end

        -- Build the result.
        local header
        if timed_out then
            header = string.format(
                "[shell] timed out after %ds; killed (exit code=%s signal=%s)",
                timeout_s, tostring(exit_code or "?"), tostring(exit_signal or "?")
            )
        elseif exit_signal and exit_signal ~= 0 then
            header = string.format(
                "[shell] terminated by signal %s (exit code=%s)",
                tostring(exit_signal), tostring(exit_code or "?")
            )
        else
            header = string.format("[shell] exit code %d", exit_code or -1)
        end

        local body = buffer_to_string()
        local parts = { header, "", body }
        if overflowed then
            if spill_path and not spill_error then
                parts[#parts + 1] = string.format(
                    "\n[truncated: kept last %d bytes; full %d-byte transcript " ..
                        "saved to %s. Use `read` with `start_line`/`end_line` " ..
                        "to inspect specific regions.]",
                    MAX_BYTES, total_written_to_spill, spill_path
                )
            else
                parts[#parts + 1] = string.format(
                    "\n[truncated: kept last %d bytes; could NOT spill to disk: %s]",
                    MAX_BYTES, spill_error or "unknown error"
                )
            end
        end
        return table.concat(parts, "\n")
    end,
}