-- subagents/workflow.lua
--
-- The callback-based Lua workflow API. `M.workflow(fn)` wraps a
-- `function(ctx, input)` in a tagged table; `M.execute(wf, input, opts)` runs
-- it and returns whatever the callback returned. Human-authored panto
-- extensions require this module directly; the restricted `subagents.lua` tool
-- (subagents/luatool.lua) and the TOML DAG lowering
-- (subagents/toml_workflows.lua) are both built on the same primitives.
--
-- ctx surface:
--   ctx:agent{ agent=, prompt=, model?, reasoning?, output? } -> handle
--   handle:await()                -> one settled result
--   ctx:await(handles, "all")     -> array of settled results in input order
--   ctx:await(handles, "first")   -> first settled result, remaining handles
--
-- Profile resolution and spawn-spec construction are NOT duplicated here: both
-- come from subagents/spawn.lua (`build_spec(input, profiles)` / `spawn(spec)`)
-- so the tool > profile > primary precedence lives in exactly one place. The
-- only thing this file adds to the spec is the synthetic structured-output
-- tool, which it normalizes to { name, description, schema } (defaulting the
-- name to "emit_result") so the host seam always sees the same shape.
--
-- Edge cases and deliberate policies:
--
-- * Child failures are values, never errors. A rejected spawn produces a
--   pre-settled handle with status "failed" so a workflow can branch on it;
--   execute() only raises for programmer/guest errors (bad arguments, an
--   exceeded job budget, an error thrown by the callback itself).
-- * The callback runs on the caller's own coroutine — the tool handler's —
--   because `subagents.jobs.await` parks the running coroutine and resumes that
--   exact coroutine from the uv callback that saw the child settle. Wrapping the
--   callback in a nested coroutine would park the wrong thread and wedge the
--   handler. opts.on_resume/opts.on_yield bracket the callback with that
--   coroutine so a caller can arm a guard on it (the instruction budget the
--   sandbox needs); trusted callers just omit them.
-- * "first" mode may settle several jobs at once. Extra results are cached on
--   their handles rather than dropped, and a handle that already holds a
--   result is served from that cache without re-entering the host, so no
--   settled result is ever lost between awaits.
-- * Handles the callback never awaited are awaited ("all") after it returns,
--   purely so no child is orphaned; those results are discarded.
-- * Structured output is decoded from result.structured_json and validated
--   against output.schema by the one validator below: the JSON Schema subset
--   a child's output tool actually uses, ignoring keywords it does not know.
--   There is deliberately no second, rock-dependent path -- `jsonschema` needs
--   lrexlib-pcre and a system PCRE that stock macOS lacks, so it is not a
--   declared dependency, and a validator picked by whether a rock happens to
--   be installed would make the same output pass here and fail there. A
--   validation failure turns the result into status "failed"; it is never
--   reported as a successful structured result.
-- * The host seam is reached through `require("panto").ext` at call time, not
--   aliased at load time, matching subagents/spawn.lua so a test can install a
--   fake `panto` module before the first call rather than before the require.
--   subagents.jobs is required the same way, so load order between the two
--   never matters.

local spawn = require("subagents.spawn")

local M = {}

local workflow_mt = { __name = "subagents.workflow" }

-- ---------------------------------------------------------------------------
-- Host seam access
-- ---------------------------------------------------------------------------

local function host()
    return require("panto").ext
end

-- The job machinery, resolved at call time for the same reason as the host seam.
local function jobs()
    return require("subagents.jobs")
end

local function host_json()
    local ok, ext = pcall(host)
    if ok and type(ext) == "table" and type(ext.json) == "table" then
        return ext.json
    end
    return nil
end

-- Decode a JSON document. Panto installs its own codec as `panto.ext.json`;
-- the dkjson fallback only matters for a bare `lua` process running the specs.
local function json_decode(text)
    local json = host_json()
    if json and json.decode then
        return json.decode(text)
    end
    local ok, dkjson = pcall(require, "dkjson")
    if ok and type(dkjson) == "table" and dkjson.decode then
        local value, _, err = dkjson.decode(text)
        if err then
            error(err, 0)
        end
        return value
    end
    error("no JSON decoder available (panto.ext.json missing, dkjson not installed)", 0)
end

local function json_encode(value)
    local json = host_json()
    if json and json.encode then
        local ok, encoded = pcall(json.encode, value)
        if ok then
            return encoded
        end
    end
    local ok, dkjson = pcall(require, "dkjson")
    if ok and type(dkjson) == "table" and dkjson.encode then
        local encoded_ok, encoded = pcall(dkjson.encode, value)
        if encoded_ok then
            return encoded
        end
    end
    return tostring(value)
end

M.json_encode = json_encode

-- ---------------------------------------------------------------------------
-- Schema validation
-- ---------------------------------------------------------------------------

-- The validator: the JSON Schema subset that structured child output actually
-- uses. Anything it does not understand is ignored rather than rejected, so an
-- unrecognized keyword never fails a legitimate result.
local function is_array_like(value)
    local count = 0
    for key in pairs(value) do
        if type(key) ~= "number" then
            return false
        end
        count = count + 1
    end
    return count == #value
end

local function type_matches(value, expected)
    if expected == "object" then
        return type(value) == "table"
    elseif expected == "array" then
        return type(value) == "table" and is_array_like(value)
    elseif expected == "string" then
        return type(value) == "string"
    elseif expected == "number" then
        return type(value) == "number"
    elseif expected == "integer" then
        return type(value) == "number" and value == math.floor(value)
    elseif expected == "boolean" then
        return type(value) == "boolean"
    elseif expected == "null" then
        return value == nil or type(value) == "userdata"
    end
    return true
end

local function check_schema(value, schema, path)
    if type(schema) ~= "table" then
        return true
    end

    local expected = schema.type
    if type(expected) == "string" then
        if not type_matches(value, expected) then
            return false, string.format("%s: expected %s, got %s", path, expected, type(value))
        end
    elseif type(expected) == "table" then
        local any = false
        for _, candidate in ipairs(expected) do
            if type_matches(value, candidate) then
                any = true
                break
            end
        end
        if not any then
            return false, string.format("%s: no listed type matched %s", path, type(value))
        end
    end

    if type(schema.enum) == "table" then
        local found = false
        for _, allowed in ipairs(schema.enum) do
            if allowed == value then
                found = true
                break
            end
        end
        if not found then
            return false, string.format("%s: value is not one of the enumerated options", path)
        end
    end

    if type(value) == "string" then
        if type(schema.minLength) == "number" and #value < schema.minLength then
            return false, string.format("%s: shorter than minLength %d", path, schema.minLength)
        end
        if type(schema.maxLength) == "number" and #value > schema.maxLength then
            return false, string.format("%s: longer than maxLength %d", path, schema.maxLength)
        end
    end

    if type(value) == "number" then
        if type(schema.minimum) == "number" and value < schema.minimum then
            return false, string.format("%s: below minimum %s", path, tostring(schema.minimum))
        end
        if type(schema.maximum) == "number" and value > schema.maximum then
            return false, string.format("%s: above maximum %s", path, tostring(schema.maximum))
        end
    end

    if type(value) ~= "table" then
        return true
    end

    if type(schema.required) == "table" then
        for _, key in ipairs(schema.required) do
            if value[key] == nil then
                return false, string.format("%s: missing required property '%s'", path, tostring(key))
            end
        end
    end

    if type(schema.properties) == "table" then
        for key, sub in pairs(schema.properties) do
            if value[key] ~= nil then
                local ok, err = check_schema(value[key], sub, path .. "." .. tostring(key))
                if not ok then
                    return false, err
                end
            end
        end
        if schema.additionalProperties == false then
            for key in pairs(value) do
                if schema.properties[key] == nil then
                    return false, string.format("%s: unexpected property '%s'", path, tostring(key))
                end
            end
        end
    end

    if type(schema.items) == "table" then
        if type(schema.minItems) == "number" and #value < schema.minItems then
            return false, string.format("%s: fewer than minItems %d", path, schema.minItems)
        end
        if type(schema.maxItems) == "number" and #value > schema.maxItems then
            return false, string.format("%s: more than maxItems %d", path, schema.maxItems)
        end
        for index, item in ipairs(value) do
            local ok, err = check_schema(item, schema.items, string.format("%s[%d]", path, index))
            if not ok then
                return false, err
            end
        end
    end

    return true
end

-- ---------------------------------------------------------------------------
-- Result shaping
-- ---------------------------------------------------------------------------

local function copy_result(result)
    local shaped = {}
    if type(result) == "table" then
        for key, value in pairs(result) do
            shaped[key] = value
        end
    end
    if shaped.status == nil then
        shaped.status = "failed"
        shaped.error = shaped.error or "the host returned no result for this job"
        shaped.resumable = false
    end
    return shaped
end

local function fail(shaped, message)
    shaped.status = "failed"
    shaped.error = message
    shaped.output = nil
    return shaped
end

-- Turn a host result into the value a workflow callback sees. Only handles
-- carrying an output schema decode structured JSON; everything else passes
-- through untouched.
local function shape_result(result, handle)
    local shaped = copy_result(result)
    local schema = handle and handle.output_schema
    if schema == nil or shaped.status ~= "completed" then
        return shaped
    end

    -- An empty tool input is a real provider case (a chat-style provider
    -- finalizes an argument-less call with ""), and it is not a validation
    -- failure: nothing was produced to validate.
    local raw = shaped.structured_json
    if raw == nil or raw == "" then
        return fail(shaped, "structured output missing: the child produced no structured result")
    end

    local decoded_ok, decoded = pcall(json_decode, raw)
    if not decoded_ok then
        return fail(shaped, "structured output failed validation: " .. tostring(decoded))
    end

    local valid, message = check_schema(decoded, schema, "output")
    if not valid then
        return fail(shaped, "structured output failed validation: " .. tostring(message or "schema mismatch"))
    end

    shaped.output = decoded
    return shaped
end

-- The host may hand back a single result table or an array of them; both are
-- normalized to an array here. A result always carries `status`, which is what
-- distinguishes the two shapes.
local function as_result_array(value)
    if type(value) ~= "table" then
        return {}
    end
    if value.status ~= nil then
        return { value }
    end
    return value
end

-- ---------------------------------------------------------------------------
-- Handles and context
-- ---------------------------------------------------------------------------

local handle_mt = {}
handle_mt.__index = handle_mt
handle_mt.__name = "subagents.handle"

function handle_mt:await()
    return self.ctx:await({ self }, "all")[1]
end

local ctx_mt = {}
ctx_mt.__index = ctx_mt
ctx_mt.__name = "subagents.ctx"

-- Per-run state lives here, not on ctx: the sandboxed guest holds the ctx table
-- and would otherwise be able to raise its own job cap (`ctx.max_jobs = nil`),
-- reset the counter, or read the profile set. ctx itself is an empty table
-- exposing only `agent` and `await`. Weak keys so a finished run is collectable.
local state = setmetatable({}, { __mode = "k" })

-- The started job stays off the handle for the same reason: a subagents.jobs
-- handle owns the child's agent and job userdata, so a guest holding a workflow
-- handle would otherwise reach `agent:run_async` directly and start children
-- outside the job cap. The guest sees only `result` and `await`.
local job_of = setmetatable({}, { __mode = "k" })

function ctx_mt:agent(input)
    if type(input) ~= "table" then
        error("ctx:agent expects a table of { agent =, prompt =, ... }", 2)
    end
    local s = state[self]
    if not s then
        error("ctx:agent must be called on a workflow context (use ctx:agent{...})", 2)
    end
    if s.max_jobs and s.job_count >= s.max_jobs then
        error(string.format("workflow job limit exceeded (max %d)", s.max_jobs), 2)
    end

    -- spawn.build_spec owns profile lookup, model/reasoning precedence, the
    -- child-role system messages, and the synthetic output tool. A nil profile
    -- set means "use the cached discovery", which is what it already does.
    local spec, spec_err = spawn.build_spec(input, s.profiles)
    if not spec then
        error(tostring(spec_err), 2)
    end

    local output_schema = nil
    if type(spec.output) == "table" then
        output_schema = spec.output.schema
    end

    local handle = setmetatable({
        ctx = self,
        spec = spec,
        output_schema = output_schema,
        result = nil,
    }, handle_mt)

    local job, job_err = spawn.spawn(spec)
    if not job then
        -- A rejected spawn is a child failure, not a workflow error.
        handle.result = {
            id = nil,
            status = "failed",
            error = tostring(job_err or "the host refused to start the child"),
            resumable = false,
        }
    else
        job_of[handle] = job
    end

    s.job_count = s.job_count + 1
    s.outstanding[#s.outstanding + 1] = handle
    return handle
end

-- Pick the first handle (in input order) that already holds a settled result,
-- returning it with the remaining handles.
local function take_settled(handles)
    for index, handle in ipairs(handles) do
        if handle.result ~= nil then
            local remaining = {}
            for other_index, other in ipairs(handles) do
                if other_index ~= index then
                    remaining[#remaining + 1] = other
                end
            end
            return handle.result, remaining
        end
    end
    return nil, nil
end

local function pending_handles(handles)
    local pending, started = {}, {}
    for _, handle in ipairs(handles) do
        if handle.result == nil and job_of[handle] ~= nil then
            pending[#pending + 1] = handle
            started[#started + 1] = job_of[handle]
        end
    end
    return pending, started
end

function ctx_mt:await(handles, mode)
    mode = mode or "all"
    if type(handles) ~= "table" then
        error("ctx:await expects an array of handles", 2)
    end
    if getmetatable(handles) == handle_mt then
        handles = { handles }
    end
    if mode ~= "all" and mode ~= "first" then
        error("ctx:await mode must be \"all\" or \"first\"", 2)
    end

    if mode == "all" then
        local pending, started = pending_handles(handles)
        if #pending > 0 then
            local results = as_result_array(jobs().await(started, "all"))
            for index, handle in ipairs(pending) do
                handle.result = shape_result(results[index], handle)
            end
        end
        local out = {}
        for index, handle in ipairs(handles) do
            out[index] = handle.result or copy_result(nil)
        end
        return out
    end

    local ready, remaining = take_settled(handles)
    if ready ~= nil then
        return ready, remaining
    end

    local pending, started = pending_handles(handles)
    if #pending == 0 then
        return nil, {}
    end

    local results, still_pending = jobs().await(started, "first")
    results = as_result_array(results)

    -- Everything not listed as still-running has settled; pair those handles
    -- with the returned results in order. The listed jobs are the very handles
    -- that went in, so they match by identity.
    local settled = pending
    if type(still_pending) == "table" and #still_pending > 0 then
        local still_running = {}
        for _, job in ipairs(still_pending) do
            still_running[job] = true
        end
        settled = {}
        for _, handle in ipairs(pending) do
            if not still_running[job_of[handle]] then
                settled[#settled + 1] = handle
            end
        end
    end
    for index, result in ipairs(results) do
        local handle = settled[index]
        if handle then
            handle.result = shape_result(result, handle)
        end
    end

    ready, remaining = take_settled(handles)
    if ready == nil then
        -- The await returned without settling anything; treat the batch as
        -- failed rather than spinning forever on the same handles.
        local first = pending[1]
        first.result = copy_result(nil)
        return take_settled(handles)
    end
    return ready, remaining
end

-- ---------------------------------------------------------------------------
-- Workflow objects and execution
-- ---------------------------------------------------------------------------

function M.workflow(fn)
    if type(fn) ~= "function" then
        error("subagents.workflow expects a function(ctx, input)", 2)
    end
    return setmetatable({ run = fn }, workflow_mt)
end

function M.is_workflow(value)
    return type(value) == "table" and getmetatable(value) == workflow_mt
end

-- Settle any handle the callback left running so a returning workflow never
-- orphans a child. Results are intentionally discarded.
local function drain(ctx)
    local pending = {}
    for _, handle in ipairs(state[ctx].outstanding) do
        if handle.result == nil then
            pending[#pending + 1] = handle
        end
    end
    if #pending == 0 then
        return
    end
    pcall(function()
        ctx:await(pending, "all")
    end)
end

-- Run `wf` against `input`. opts:
--   max_jobs   -- cap on ctx:agent calls (nil = unbounded; the sandbox passes 32)
--   profiles   -- discovered profile set for spawn.build_spec (nil = discover)
--   on_resume  -- called with the running coroutine before the callback starts
--   on_yield   -- called with the same coroutine once it has finished
--
-- The callback runs on the CALLER's coroutine, never a nested one:
-- `subagents.jobs.await` parks whichever coroutine is running when it suspends
-- and resumes exactly that coroutine when the job settles. A nested coroutine
-- would be the thread parked and resumed, leaving the tool handler that yielded
-- around it suspended forever.
function M.execute(wf, input, opts)
    if not M.is_workflow(wf) then
        error("subagents.workflow.execute expects a workflow object", 2)
    end
    opts = opts or {}

    local ctx = setmetatable({}, ctx_mt)
    state[ctx] = {
        profiles = opts.profiles,
        max_jobs = opts.max_jobs,
        job_count = 0,
        outstanding = {},
    }

    local co = coroutine.running()
    if opts.on_resume then
        opts.on_resume(co)
    end
    -- pcall is yieldable in 5.4, so the callback may still await across it.
    local packed = table.pack(pcall(wf.run, ctx, input))
    if opts.on_yield then
        opts.on_yield(co)
    end

    drain(ctx)
    if not packed[1] then
        error(packed[2], 0)
    end
    return table.unpack(packed, 2, packed.n)
end

-- A child's output text: a structured result decodes to a table, which is
-- re-encoded compactly so anything model-visible is still a string.
function M.output_text(result)
    local output = result.output
    if type(output) == "table" then
        return json_encode(output)
    end
    if output == nil or output == "" then
        return tostring(result.error or "")
    end
    return tostring(output)
end

return M