-- subagents/luatool.lua
--
-- The `subagents.lua` model-facing tool: run a transient, model-authored Lua
-- workflow without writing a definition to disk. The source must evaluate to
-- `subagents.workflow(function(ctx, input) ... end)`; the tool then executes it
-- with the tool's `prompt` as the workflow input and formats the terminal
-- results for the calling model. Optional inline agent profiles are overlaid
-- for this execution only; they are never persisted or added to discovery.
--
-- The source is loaded in text mode only (`load(source, chunkname, "t", env)`)
-- against a restricted `_ENV`. That environment holds a safe slice of the
-- standard library plus `subagents.workflow`; it has no `os`, `io`, `debug`,
-- `package`, `require`, `load`, `dofile`, `coroutine`, `setmetatable`, or
-- `getmetatable`, and `print` is a no-op so generated code cannot scribble on
-- the TUI. `string`, `table`, `math`, and `utf8` are shallow copies, so a guest
-- that reassigns `table.insert` only breaks itself, and `string.dump` is
-- removed from the copy.
--
-- `pcall`/`xpcall` are deliberately absent: they would let a guest catch the
-- instruction-budget error and spin again in fresh 10M-instruction chunks. A
-- guest has no need to recover from its own errors — the handler reports them —
-- and with no `pcall`, `coroutine`, or metatable access left, nothing in the
-- environment can trap the hook error before it reaches the host.
--
-- Known, accepted gaps in that sandbox:
--
-- * The real string metatable is still reachable through any string literal
--   (`("").dump`), so `string.dump` is obtainable. Without `load` there is no
--   way to turn bytecode back into a running function, so this is noise rather
--   than an escape.
-- * `ctx:agent` returns handles the guest can read and scribble on. Nothing
--   reachable from one starts work: the running job — which owns the child's
--   agent and could start turns outside the cap — lives in a private side table
--   in workflow.lua, as do the job cap, the counter, and the profile set, so a
--   guest holding `ctx` and its handles cannot raise its own cap or reach the
--   host.
--
-- Runaway generated Lua is bounded two ways: `max_jobs = 32` caps how many
-- children one transient workflow may start, and a debug count hook is armed
-- for the duration of the guest callback and disarmed as soon as it returns.
-- The hook lives here rather than in workflow.lua because the guest has no
-- `debug` library but the host does; workflow.execute only exposes the
-- on_resume/on_yield seam the hook needs. The guest runs on the tool handler's
-- own coroutine (an await parks and resumes exactly that coroutine when a child
-- settles), so the budget covers the whole run rather than one slice; awaiting
-- a child executes no instructions, so only real spinning trips it.

local workflow = require("subagents.workflow")
local run = require("subagents.run")
local spawn = require("subagents.spawn")

local MAX_JOBS = 32
local INSTRUCTION_BUDGET = 10000000
local CHUNK_NAME = "subagents.lua"

local M = {}

M.max_jobs = MAX_JOBS

local function shallow_copy(source, skip)
    local copy = {}
    for key, value in pairs(source) do
        if key ~= skip then
            copy[key] = value
        end
    end
    return copy
end

-- Build a fresh restricted environment per call: the guest may mutate anything
-- it can reach, so nothing here is shared between invocations.
local function build_env()
    local env = {
        assert = assert,
        error = error,
        ipairs = ipairs,
        next = next,
        pairs = pairs,
        select = select,
        tonumber = tonumber,
        tostring = tostring,
        type = type,
        string = shallow_copy(string, "dump"),
        table = shallow_copy(table),
        math = shallow_copy(math),
        utf8 = shallow_copy(utf8),
        print = function() end,
        subagents = { workflow = workflow.workflow },
    }
    env._G = env
    return env
end

M.build_env = build_env

-- Overlay workflow-local profiles without mutating the activation-time set.
-- Inline names intentionally win, matching normal layered profile precedence
-- while keeping their lifetime bounded to this one execute() call.
local function workflow_profiles(inline, discovered)
    if inline == nil then
        return discovered, nil
    end
    if type(inline) ~= "table" then
        return nil, "agents must be an array"
    end

    local count = 0
    for key in pairs(inline) do
        if type(key) ~= "number" then
            return nil, "agents must be an array"
        end
        count = count + 1
    end
    if count ~= #inline then
        return nil, "agents must be a dense array"
    end

    local base = spawn.profiles(discovered)
    local by_name = {}
    for name, profile in pairs(base.by_name or {}) do
        by_name[name] = profile
    end
    local seen = {}
    for index, raw in ipairs(inline) do
        if type(raw) ~= "table" then
            return nil, string.format("agents[%d] must be an object", index)
        end
        local name = raw.name
        if type(name) ~= "string" or name:match("^%s*$") then
            return nil, string.format("agents[%d].name must be a non-empty string", index)
        end
        if seen[name] then
            return nil, string.format("duplicate inline agent '%s'", name)
        end
        local system_prompt = raw.system_prompt
        if type(system_prompt) ~= "string" or system_prompt:match("^%s*$") then
            return nil, string.format("agents[%d].system_prompt must be a non-empty string", index)
        end
        if raw.description ~= nil and type(raw.description) ~= "string" then
            return nil, string.format("agents[%d].description must be a string when given", index)
        end
        seen[name] = true
        by_name[name] = {
            name = name,
            description = raw.description or "",
            body = system_prompt,
            layer = "workflow",
        }
    end

    local list = {}
    for _, profile in pairs(by_name) do
        list[#list + 1] = profile
    end
    table.sort(list, function(a, b) return a.name < b.name end)
    return { list = list, by_name = by_name, warnings = base.warnings or {} }, nil
end

local function budget_hook()
    error("instruction budget exceeded", 0)
end

-- subagents.run's block formatter expects string output; a structured worker's
-- output is a decoded table, so it is re-encoded first.
local function format_one(result)
    if type(result.output) == "table" then
        local flattened = {}
        for key, value in pairs(result) do
            flattened[key] = value
        end
        flattened.output = workflow.output_text(result)
        return run.format_result(flattened)
    end
    return run.format_result(result)
end

-- Format whatever the workflow callback returned. Result-shaped tables (the
-- common case: one settled result, or an array of them) render as the same
-- plain "key: value" block subagents.run uses; anything else is encoded
-- compactly so the model still sees it.
local function format_return(value)
    if value == nil then
        return "The workflow returned no value."
    end
    if type(value) ~= "table" then
        return tostring(value)
    end
    if value.status ~= nil then
        return format_one(value)
    end
    local blocks, count = {}, 0
    for index, entry in ipairs(value) do
        if type(entry) ~= "table" or entry.status == nil then
            blocks = nil
            break
        end
        blocks[index] = format_one(entry)
        count = count + 1
    end
    if blocks and count > 0 then
        return table.concat(blocks, "\n\n")
    end
    return workflow.json_encode(value)
end

-- Tool handler for `subagents.lua`. `profiles` is the discovered profile set
-- from activation; when omitted the workflow API discovers it lazily.
function M.handle(input, profiles)
    if type(input) ~= "table" then
        return "Error: expected an input object"
    end
    if type(input.prompt) ~= "string" or input.prompt == "" then
        return "Error: prompt is required and must be a non-empty string"
    end
    if type(input.source) ~= "string" or input.source == "" then
        return "Error: source is required and must be a non-empty string"
    end

    local profiles_for_run, profiles_err = workflow_profiles(input.agents, profiles)
    if not profiles_for_run then
        return "Error: " .. profiles_err
    end

    local chunk, load_err = load(input.source, CHUNK_NAME, "t", build_env())
    if not chunk then
        return "Error: source did not compile: " .. tostring(load_err)
    end

    local built_ok, built = pcall(chunk)
    if not built_ok then
        return "Error: source failed to run: " .. tostring(built)
    end
    if not workflow.is_workflow(built) then
        return "Error: source must return subagents.workflow(function(ctx, input) ... end)"
    end

    local armed = nil
    local ran_ok, result = pcall(workflow.execute, built, input.prompt, {
        max_jobs = MAX_JOBS,
        profiles = profiles_for_run,
        on_resume = function(co)
            armed = co
            debug.sethook(co, budget_hook, "", INSTRUCTION_BUDGET)
        end,
        on_yield = function(co)
            debug.sethook(co)
            armed = nil
        end,
    })
    if armed then
        debug.sethook(armed)
    end

    if not ran_ok then
        return "Error: " .. tostring(result)
    end
    return format_return(result)
end

return M