-- A scriptable stand-in for the `panto` module the extension talks to.
--
-- Every module reaches the host through `require("panto")` at call time, never
-- as a load-time alias, so a spec can install this fake after the modules are
-- already loaded. `install()` writes `package.loaded.panto`; `handle.restore()`
-- removes it again.
--
-- The surface mirrors the real one, covering what the extension actually calls:
--
--   panto.ext.session_info()        -> { session_id, session_dir, model, reasoning }
--   panto.ext.resolve_model(q)      -> cfg | nil, "resolve_model: ..."
--   panto.ext.models(query)         -> catalog answer
--   panto.ext.agent                 -- the primary, borrowed: :tools(), :conversation()
--   panto.ext.on(name, handler)     -- lifecycle subscriptions
--   panto.ext.register_tool / register_command
--   panto.ext.json                  -- panto's JSON codec (dkjson stands in here)
--   panto.agent{ config =, store =, session_id =, conversation = } -> agent
--   panto.file_system_jsonl_store{ dir = } / panto.null_store()
--   agent:conversation() / :session_id() / :tools() / :set_tools(decls) / :run_async(opts)
--   store:resolve(id) / :load(id)
--   conv:messages() / :message_metadata(i) / :add_system_message(text, { metadata = })
--   job:next_event() / :result() / :request_cancel() / :close()
--
-- Scripting. Every child takes its outcome from a queue: `handle.queue(outcome)`
-- is matched first-in-first-out, `handle.queue_for(label, outcome)` is matched
-- against the profile name the child was seeded with (its manifest metadata; a
-- resumed child carries no fresh manifest, so it always draws from the FIFO
-- queue). An outcome describes the settled turn — `{ status =, output =,
-- error =, structured_json =, resumable =, id =, settle = }`. A child with
-- nothing queued completes with a generated id and a generic output, which
-- keeps tests that only care about what was asked for short.
--
-- Settling. There is no event loop here, so a fake job settles on the Nth call
-- to `job:result()`: `settle = N`, lowest first, which is the same ordering key
-- the previous fake used for "first" awaits. Real wake pipes are opened and
-- polled (the job machinery refuses to start a child without one), but nothing
-- ever runs the loop and nothing writes a wake byte; the specs call await from
-- a plain script, which cannot park, so awaiting drains its jobs in place and
-- those polls are what turn into progress. `request_cancel` settles the next
-- poll as cancelled, exactly like the binding's pump.
--
-- Everything the host was asked to do is recorded on the returned handle:
-- `spawns` (one record per child agent, in creation order, carrying the
-- resolve_model arguments, the store directory, the seeded system messages, the
-- tool declarations it was given and the run_async options), `runs` (run_async
-- options only, so a gated child is visibly not started yet), `jobs`,
-- `resolves`, `tools`, `commands`, `models_queries`, `subscriptions`, `mkdirs`,
-- `max_live` and `polls`. Assertions read those tables directly.

local M = {}

-- ---------------------------------------------------------------------------
-- luv stand-in, installed once when this module loads
-- ---------------------------------------------------------------------------

-- `fs_mkdir` is recorded rather than performed, so a spec can assert which
-- directories a child store would need without touching the disk. Everything
-- else — pipes included — delegates to the real luv, so a spec exercises the
-- same wake-pipe arming production does. This has to happen at load time, not
-- inside install(), because a module that resolves luv once at load would
-- otherwise capture the real one.
local made_dirs = {}
do
    local ok, real = pcall(require, "luv")
    if ok and type(real) == "table" then
        package.loaded.luv = setmetatable({
            fs_mkdir = function(path)
                made_dirs[#made_dirs + 1] = path
                return true
            end,
        }, { __index = real })
    end
end

-- A job polled this many times without settling means the drain loop is
-- spinning; failing beats hanging the whole suite.
local POLL_LIMIT = 5000

local DEFAULT_SESSION = {
    session_id = "0198-primary",
    session_dir = "/tmp/panto-spec-sessions/--Users-travis-Code-panto-subagents--",
    model = "anthropic:sonnet",
    reasoning = "medium",
}

local DEFAULT_OVERVIEW = {
    model = "anthropic:sonnet",
    reasoning = "medium",
    providers = {
        { name = "anthropic", style = "messages", models = 7 },
        { name = "openai", style = "responses", models = 12 },
    },
}

-- Stands in for the lightuserdata re-registration tag a real decl carries. The
-- identity is what matters: a child must be given the primary's own tags.
M.SOURCE = setmetatable({}, { __name = "fake.tool_source" })

-- The primary's tools: two a child inherits, and the four it must not.
-- agent:tools() sorts by name, so this is the order the rock sees them in.
local DEFAULT_TOOL_NAMES = {
    "bash",
    "read_file",
    "subagents.lua",
    "subagents.models",
    "subagents.run",
    "subagents.workflow",
}

local function default_tools()
    local decls = {}
    for index, name in ipairs(DEFAULT_TOOL_NAMES) do
        decls[index] = {
            name = name,
            description = name .. " does a thing",
            schema = { type = "object", properties = {} },
            _source = M.SOURCE,
        }
    end
    return decls
end

-- panto installs its own JSON codec as `panto.ext.json`; dkjson is the closest
-- stand-in available to a bare `lua` process. Absent, the field is simply
-- missing, exactly as it would be on a host too old to provide it.
local function json_codec()
    local ok, dkjson = pcall(require, "dkjson")
    if not ok or type(dkjson) ~= "table" then
        return nil
    end
    return {
        encode = function(value)
            return dkjson.encode(value)
        end,
        decode = function(text)
            local value, _, err = dkjson.decode(text)
            if err then
                error(err, 0)
            end
            return value
        end,
    }
end

local function copy(source, fallback)
    local out = {}
    for key, value in pairs(source or fallback or {}) do
        out[key] = value
    end
    return out
end

-- ---------------------------------------------------------------------------
-- Conversations
-- ---------------------------------------------------------------------------

local conv_mt = {}
conv_mt.__index = conv_mt
conv_mt.__name = "fake.conversation"

-- scripted = array of { role =, text = | blocks =, metadata = }. `record`, when
-- given, is the child record seeded system messages are mirrored onto. A
-- scripted message may instead carry `metadata_error =`, which is how the
-- binding reports a stored record it cannot decode: by raising.
local function new_conversation(scripted, record)
    local messages = {}
    for index, message in ipairs(scripted or {}) do
        messages[index] = {
            role = message.role or "user",
            blocks = message.blocks or { { type = "text", text = message.text or "" } },
            metadata = message.metadata,
            metadata_error = message.metadata_error,
        }
    end
    return setmetatable({ _messages = messages, _record = record }, conv_mt)
end

-- Faithful to the binding: metadata is not part of a message table, so the only
-- way to reach it is message_metadata(i).
function conv_mt:messages()
    local out = {}
    for index, message in ipairs(self._messages) do
        local blocks = {}
        for block_index, block in ipairs(message.blocks) do
            blocks[block_index] = copy(block)
        end
        out[index] = { role = message.role, blocks = blocks }
    end
    return out
end

function conv_mt:message_metadata(index)
    local message = self._messages[tonumber(index) or 0]
    if message == nil then
        return nil
    end
    if message.metadata_error then
        error("panto: " .. message.metadata_error, 2)
    end
    return message.metadata
end

function conv_mt:add_system_message(text, opts)
    if type(text) ~= "string" then
        error("panto: add_system_message expects a string", 2)
    end
    local metadata = opts and opts.metadata
    if metadata ~= nil and type(metadata) ~= "table" then
        error("panto: metadata must be a table", 2)
    end
    self._messages[#self._messages + 1] = {
        role = "system",
        blocks = { { type = "text", text = text } },
        metadata = metadata,
    }
    local record = self._record
    if record then
        record.system_messages[#record.system_messages + 1] = { text = text, metadata = metadata }
        local manifest = metadata and metadata.subagents
        if type(manifest) == "table" and type(manifest.agent) == "string" then
            record.label = manifest.agent
        end
    end
end

-- ---------------------------------------------------------------------------
-- Session stores
-- ---------------------------------------------------------------------------

local store_mt = {}
store_mt.__index = store_mt
store_mt.__name = "fake.store"

function store_mt:resolve(id)
    local session = self._harness.sessions[id]
    if session == nil then
        return nil
    end
    return { id = id, message_count = #session, api_style = "messages" }
end

function store_mt:load(id)
    local session = self._harness.sessions[id]
    if session == nil then
        return nil
    end
    return new_conversation(session)
end

-- ---------------------------------------------------------------------------
-- Jobs
-- ---------------------------------------------------------------------------

local job_mt = {}
job_mt.__index = job_mt
job_mt.__name = "fake.job"

local function new_job(cfg)
    return setmetatable({
        _settle_after = math.max(1, math.floor(tonumber(cfg.settle) or 1)),
        _events = cfg.events or {},
        _finish = cfg.finish,
        _harness = cfg.harness,
        _polls = 0,
        _settled = false,
        _result = nil,
        _cancel_requested = false,
        _closed = false,
    }, job_mt)
end

function job_mt:next_event()
    if self._closed or #self._events == 0 then
        return nil
    end
    return table.remove(self._events, 1)
end

function job_mt:result()
    if self._closed then
        -- close() frees the settled result in the binding, so a caller that
        -- reads it back afterwards sees exactly this.
        return nil
    end
    self._polls = self._polls + 1
    if self._polls > POLL_LIMIT then
        error(string.format(
            "fake job: polled %d times without settling; the drain loop is making no progress on a job with no wake pipe",
            POLL_LIMIT), 0)
    end
    if self._harness then
        self._harness.polls = self._harness.polls + 1
    end
    if self._settled then
        return self._result
    end
    if self._cancel_requested then
        self._result = { status = "cancelled", error = "cancelled" }
    elseif self._polls >= self._settle_after then
        self._result = self._finish(self)
    else
        return nil
    end
    self._settled = true
    if self._harness then
        self._harness.live = self._harness.live - 1
    end
    return self._result
end

function job_mt:request_cancel()
    self._cancel_requested = true
end

function job_mt:close()
    if self._closed then
        return
    end
    self._closed = true
    self._result = nil
    self._events = {}
end

-- job(spec) -> a standalone fake job, for specs that drive the job machinery
-- directly rather than through a child. spec = { result =, events =, settle = }.
function M.job(spec)
    spec = spec or {}
    return new_job({
        settle = spec.settle,
        events = spec.events,
        finish = function()
            return spec.result or { status = "completed", text = "done" }
        end,
    })
end

-- ---------------------------------------------------------------------------
-- Outcomes: which scripted turn a child gets, and what it settles to
-- ---------------------------------------------------------------------------

local function next_outcome(h, label)
    local bucket = label and h.labelled[label]
    if bucket and #bucket > 0 then
        return table.remove(bucket, 1)
    end
    if #h.queued > 0 then
        return table.remove(h.queued, 1)
    end
    return {}
end

-- Bound as late as possible: the profile label only exists once the child has
-- been seeded with its manifest system message.
local function outcome_for(h, record)
    if record.outcome == nil then
        record.outcome = next_outcome(h, record.label)
    end
    return record.outcome
end

local function child_id(h, record)
    if record.session_id then
        return record.session_id
    end
    if record.id == nil then
        record.id = outcome_for(h, record).id or ("child-" .. record.index)
    end
    return record.id
end

-- The settled turn, in the shape agent:run_async reports it.
local function settled_result(h, record)
    local outcome = outcome_for(h, record)
    local status = outcome.status or "completed"
    local result = { status = status }
    local one_shot = record.run and record.run.dispatch_tools == false
    if status == "completed" then
        if one_shot then
            result.text = outcome.output
            if outcome.structured_json then
                local decl = record.tool_decls and record.tool_decls[1]
                result.tool_calls = { {
                    id = "call-" .. record.index,
                    name = outcome.tool_name or (decl and decl.name) or "emit_result",
                    input = outcome.structured_json,
                } }
            end
        else
            result.text = outcome.output or ("output of " .. tostring(record.label or child_id(h, record)))
        end
    else
        result.error = outcome.error or ("the child " .. status)
    end

    -- The durable file: a child that never reached its first assistant message
    -- has none, which is what makes it unresumable.
    local id = child_id(h, record)
    if outcome.resumable == false then
        h.sessions[id] = nil
    else
        h.sessions[id] = h.sessions[id] or {}
    end
    return result
end

-- ---------------------------------------------------------------------------
-- Agents
-- ---------------------------------------------------------------------------

local agent_mt = {}
agent_mt.__index = agent_mt
agent_mt.__name = "fake.agent"

function agent_mt:conversation()
    return self._conv
end

function agent_mt:session_id()
    return child_id(self._harness, self._record)
end

function agent_mt:tools()
    local out = {}
    for index, decl in ipairs(self._record.tool_decls or {}) do
        out[index] = decl
    end
    return out
end

function agent_mt:set_tools(decls)
    if self._borrowed then
        error("panto: set_tools is not supported on a borrowed agent", 2)
    end
    if type(decls) ~= "table" then
        error("panto: set_tools expects an array of declarations", 2)
    end
    local names = {}
    for index, decl in ipairs(decls) do
        if type(decl) ~= "table" or type(decl.name) ~= "string" then
            error("panto: tool declaration " .. index .. " has no name", 2)
        end
        names[index] = decl.name
    end
    self._record.tool_decls = decls
    self._record.tools = names
end

function agent_mt:run_async(options)
    if self._borrowed then
        return nil, "run_async: not supported on a borrowed agent"
    end
    if type(options) ~= "table" then
        error("panto: run_async expects a table", 2)
    end
    if options.prompt == nil and options.blocks == nil then
        return nil, "run_async: pass a prompt or blocks"
    end
    if options.metadata ~= nil and type(options.metadata) ~= "table" then
        return nil, "run_async: metadata must be a table"
    end

    local h, record = self._harness, self._record
    record.run = {
        prompt = options.prompt,
        blocks = options.blocks,
        metadata = options.metadata,
        dispatch_tools = options.dispatch_tools,
        wake_fd = options.wake_fd,
    }
    record.prompt = options.prompt
    h.runs[#h.runs + 1] = record.run

    local job = new_job({
        settle = outcome_for(h, record).settle,
        events = outcome_for(h, record).events,
        harness = h,
        finish = function()
            return settled_result(h, record)
        end,
    })
    h.jobs[#h.jobs + 1] = job
    record.job = job
    h.live = h.live + 1
    if h.live > h.max_live then
        h.max_live = h.live
    end
    return job
end

-- ---------------------------------------------------------------------------
-- install
-- ---------------------------------------------------------------------------

-- install(opts) -> handle
--
-- opts = {
--   session          = { session_id =, session_dir =, model =, reasoning = },
--   models_response  = table | function(query),
--   primary_tools    = array of decl tables (defaults to DEFAULT_TOOL_NAMES),
--   primary_messages = array of { role =, text = } for the primary conversation,
--   sessions         = { [id] = array of scripted messages } already on disk,
--   unknown_models   = { ["provider:model"] = true } resolve_model rejects,
-- }
function M.install(opts)
    opts = opts or {}

    for index = #made_dirs, 1, -1 do
        made_dirs[index] = nil
    end

    -- A partial `session` overrides only the fields it names.
    local session = copy(DEFAULT_SESSION)
    for key, value in pairs(opts.session or {}) do
        session[key] = value
    end

    local handle = {
        session = session,
        models_response = opts.models_response or DEFAULT_OVERVIEW,
        primary_tools = opts.primary_tools or default_tools(),
        sessions = opts.sessions or {},
        unknown_models = opts.unknown_models or {},
        spawns = {},
        runs = {},
        jobs = {},
        resolves = {},
        stores = {},
        tools = {},
        tools_by_name = {},
        commands = {},
        commands_by_name = {},
        models_queries = {},
        subscriptions = {},
        on_by_name = {},
        mkdirs = made_dirs,
        queued = {},
        labelled = {},
        polls = 0,
        live = 0,
        max_live = 0,
    }

    function handle.queue(outcome)
        handle.queued[#handle.queued + 1] = outcome or {}
    end

    function handle.queue_for(label, outcome)
        local bucket = handle.labelled[label]
        if bucket == nil then
            bucket = {}
            handle.labelled[label] = bucket
        end
        bucket[#bucket + 1] = outcome or {}
    end

    -- A child session that already exists on disk, for the resume cases.
    function handle.add_session(id, messages)
        handle.sessions[id] = messages or {}
    end

    function handle.made_dir(path)
        for _, made in ipairs(handle.mkdirs) do
            if made == path then
                return true
            end
        end
        return false
    end

    -- resolve_model hands back an opaque config; this is how a spec gets from
    -- the config a child was built with back to the query that produced it.
    local resolved_from = setmetatable({}, { __mode = "k" })

    local ext = {}

    function ext.session_info()
        return copy(handle.session)
    end

    function ext.resolve_model(query)
        if type(query) ~= "table" then
            return nil, "resolve_model: expected a table of arguments"
        end
        local request = {
            model = query.model,
            reasoning = query.reasoning,
            tool_choice = query.tool_choice,
        }
        handle.resolves[#handle.resolves + 1] = request
        if type(query.model) ~= "string" or query.model == "" then
            return nil, "resolve_model: model must be a 'provider:model' string"
        end
        if handle.unknown_models[query.model] then
            return nil, string.format("resolve_model: unknown model '%s'", query.model)
        end
        local cfg = {
            model = query.model,
            reasoning = query.reasoning,
            style = "messages",
            wire_model = query.model:match(":(.+)$") or query.model,
        }
        resolved_from[cfg] = request
        return cfg
    end

    function ext.models(query)
        handle.models_queries[#handle.models_queries + 1] = query or {}
        local response = handle.models_response
        if type(response) == "function" then
            return response(query or {})
        end
        return response
    end

    function ext.on(name, fn)
        handle.subscriptions[#handle.subscriptions + 1] = { name = name, fn = fn }
        handle.on_by_name[name] = fn
    end

    -- Fire a subscribed lifecycle handler, the way the host would.
    function handle.emit(name, event)
        for _, subscription in ipairs(handle.subscriptions) do
            if subscription.name == name then
                subscription.fn(event)
            end
        end
    end

    function ext.register_tool(tool)
        handle.tools[#handle.tools + 1] = tool
        handle.tools_by_name[tool.name] = tool
    end

    function ext.register_command(command)
        handle.commands[#handle.commands + 1] = command
        handle.commands_by_name[command.name] = command
    end

    ext.json = json_codec()

    -- The primary agent, borrowed: readable tools and conversation, nothing else.
    ext.agent = setmetatable({
        _borrowed = true,
        _harness = handle,
        _record = { index = 0, tool_decls = handle.primary_tools, system_messages = {}, tools = {} },
        _conv = new_conversation(opts.primary_messages),
    }, agent_mt)

    local function new_store(kind, arg)
        local dir = nil
        if type(arg) == "table" then
            dir = arg.dir
        elseif type(arg) == "string" then
            dir = arg
        end
        if kind == "fs" and type(dir) ~= "string" then
            error("panto.file_system_jsonl_store: missing dir", 2)
        end
        handle.stores[#handle.stores + 1] = dir
        return setmetatable({ dir = dir, kind = kind, _harness = handle }, store_mt)
    end

    local function new_agent(options)
        if type(options) ~= "table" then
            error("panto.agent expects a table", 2)
        end
        local record = {
            index = #handle.spawns + 1,
            config = options.config,
            store_dir = type(options.store) == "table" and options.store.dir or nil,
            session_id = options.session_id,
            resumed = options.session_id ~= nil,
            system_messages = {},
            tools = {},
            resolve = resolved_from[options.config],
        }
        if record.resolve then
            record.model = record.resolve.model
            record.reasoning = record.resolve.reasoning
            record.tool_choice = record.resolve.tool_choice
        end
        handle.spawns[record.index] = record

        local conv = options.conversation
        if conv == nil then
            conv = new_conversation(nil, record)
        else
            conv._record = record
        end
        return setmetatable({ _harness = handle, _record = record, _conv = conv }, agent_mt)
    end

    handle.ext = ext
    package.loaded.panto = {
        ext = ext,
        agent = new_agent,
        file_system_jsonl_store = function(arg)
            return new_store("fs", arg)
        end,
        null_store = function()
            return new_store("null", nil)
        end,
    }

    function handle.restore()
        package.loaded.panto = nil
    end

    return handle
end

return M