-- The e2e fixture extension: one scripted provider protocol plus the handful of
-- tools a child may call. Loaded by e2e/run.lua into a throwaway config layer,
-- never by a real session.
--
-- The protocol interprets a prompt as a tiny script, so one fixture plays every
-- role the scenarios need — primary and child alike. The first LINE of the most
-- recent plain user prompt is the script; everything after it is data (a
-- workflow step's prompt arrives with its dependency sections appended, and
-- those must not be parsed). The script is split on ";;" into one step per
-- round, a round being one assistant message already in the conversation, and a
-- step may hold several "&&"-joined commands.
--
--   say <text>            answer with <text>
--   tool <name> <json>    call a tool (repeat with && for one batch)
--   results               answer with every tool result in the last user message
--   find <needle>         "found <needle>" / "missing <needle>" over the history
--   tag <label> <needle>  "<label>:found" / "<label>:missing"
--   sys <needle>          the same over request.system_prompt
--   tools                 answer with the offered tool names, comma-joined
--   body                  answer with the whole prompt, first line included
--   emit <name> <json>    a forced output-tool call (one-shot structured worker)
--   err <message>         fail the stream mid-turn
--
-- CONSTRAINT: a child's protocol calls run on the primary's non-yielding path,
-- so nothing in open/next/close may yield. Everything that has to park lives in
-- a TOOL handler instead — those run as coroutines, which is what `pair` below
-- uses to prove two children are genuinely in flight at once.

local panto = require("panto")
local uv = require("luv")

-- The rock's own dependencies (lyaml, toml2lua, dkjson) are required lazily
-- during a turn, and panto's luarocks tree has no reason to hold them; the
-- driver points this at the repo's ./.rocks tree instead. APPENDING keeps
-- panto's own luv — the one built against the host libuv — ahead of the copy
-- living in that tree.
local rocks = os.getenv("PANTO_E2E_ROCKS")
if rocks ~= nil and rocks ~= "" then
    package.path = table.concat({
        package.path,
        rocks .. "/share/lua/5.4/?.lua",
        rocks .. "/share/lua/5.4/?/init.lua",
    }, ";")
    package.cpath = package.cpath .. ";" .. rocks .. "/lib/lua/5.4/?.so"
end

-- Marks are appended to a per-scenario file so the driver can assert on
-- interleaving after the process is gone. Absent variable = tracing off.
local trace_path = os.getenv("PANTO_E2E_TRACE")

local function mark(text)
    if trace_path == nil or trace_path == "" then
        return
    end
    local fh = io.open(trace_path, "a")
    if fh == nil then
        return
    end
    fh:write(text, "\n")
    fh:close()
end

-- Park this coroutine on a real libuv timer. Tool handlers only.
local function sleep_ms(ms)
    local co = coroutine.running()
    local timer = uv.new_timer()
    uv.timer_start(timer, ms, 0, function()
        uv.timer_stop(timer)
        uv.close(timer)
        coroutine.resume(co)
    end)
    coroutine.yield()
end

-- ---------------------------------------------------------------------------
-- Script parsing
-- ---------------------------------------------------------------------------

local function split(text, sep)
    local out, start = {}, 1
    while true do
        local from, to = text:find(sep, start, true)
        if from == nil then
            out[#out + 1] = text:sub(start)
            return out
        end
        out[#out + 1] = text:sub(start, from - 1)
        start = to + 1
    end
end

local function trim(text)
    return (text:gsub("^%s+", ""):gsub("%s+$", ""))
end

-- Where the current turn starts: the most recent PLAIN user message. On a tool
-- round trip the last user message is the tool result, which carries no text
-- block, so this reaches past it; on a resumed conversation everything before
-- it belongs to an earlier turn and must not be counted as this one's.
local function prompt_index(request)
    for i = #(request.messages or {}), 1, -1 do
        local message = request.messages[i]
        if message.role == "user" then
            for _, block in ipairs(message.blocks or {}) do
                if block.type == "text" then
                    return i
                end
            end
        end
    end
    return 0
end

local function prompt_of(request)
    local index = prompt_index(request)
    for _, block in ipairs(index > 0 and request.messages[index].blocks or {}) do
        if block.type == "text" then
            return block.text
        end
    end
    return ""
end

-- Rounds are scoped to the current turn: one per assistant message the model
-- has already produced since this turn's prompt.
local function round_of(request)
    local seen = 0
    for i = prompt_index(request) + 1, #(request.messages or {}) do
        if request.messages[i].role == "assistant" then
            seen = seen + 1
        end
    end
    return seen + 1
end

-- Every tool result text in this turn's last user message that has any,
-- newline-joined. Earlier turns' results stay out of it.
local function tool_results_text(request)
    for i = #(request.messages or {}), prompt_index(request) + 1, -1 do
        local found = {}
        for _, block in ipairs(request.messages[i].blocks or {}) do
            if block.type == "tool_result" then
                for _, part in ipairs(block.content or {}) do
                    if part.type == "text" then
                        found[#found + 1] = part.text
                    end
                end
            end
        end
        if #found > 0 then
            return table.concat(found, "\n")
        end
    end
    return ""
end

-- Everything the turn can see except the system prompt: the transcript panto
-- serialized for us plus every block text in the active window.
local function haystack(request)
    local parts = { request.history_transcript or "" }
    for _, message in ipairs(request.messages or {}) do
        for _, block in ipairs(message.blocks or {}) do
            if type(block.text) == "string" then
                parts[#parts + 1] = block.text
            end
            for _, part in ipairs(block.content or {}) do
                if type(part.text) == "string" then
                    parts[#parts + 1] = part.text
                end
            end
        end
    end
    return table.concat(parts, "\n")
end

-- ---------------------------------------------------------------------------
-- Streams
-- ---------------------------------------------------------------------------

local function stream_of(steps)
    local index = 0
    return {
        next = function()
            index = index + 1
            local step = steps[index]
            if step == nil then
                return { type = "done" }
            end
            return step()
        end,
    }
end

local function text_stream(text)
    return stream_of({ function()
        return { type = "text_delta", text = text }
    end })
end

local function tool_call_steps(commands)
    local steps = {}
    for index, command in ipairs(commands) do
        local verb, rest = command:match("^(%S+)%s*(.*)$")
        if verb ~= "tool" and verb ~= "emit" then
            return nil, "step " .. index .. " is not a tool call: " .. command
        end
        local name, input_json = rest:match("^(%S+)%s+(.*)$")
        if name == nil then
            return nil, "step " .. index .. " has no tool input: " .. command
        end
        steps[index] = function()
            return {
                type = "tool_call",
                id = "fx" .. index,
                name = name,
                input_json = input_json,
            }
        end
    end
    return steps
end

local function build(request, step)
    local commands = {}
    for _, part in ipairs(split(step, "&&")) do
        local command = trim(part)
        if command ~= "" then
            commands[#commands + 1] = command
        end
    end
    if #commands == 0 then
        return text_stream("")
    end

    local first = commands[1]
    local verb, rest = first:match("^(%S+)%s*(.*)$")
    verb = verb or ""

    if verb == "tool" or verb == "emit" then
        local steps, err = tool_call_steps(commands)
        if steps == nil then
            return stream_of({ function()
                return { type = "error", kind = "invalid_request", message = err }
            end })
        end
        return stream_of(steps)
    end

    if verb == "say" then
        return text_stream(rest)
    end
    if verb == "results" then
        return text_stream(tool_results_text(request))
    end
    if verb == "body" then
        return text_stream(prompt_of(request))
    end
    if verb == "tools" then
        local names = {}
        for _, tool in ipairs(request.tools or {}) do
            names[#names + 1] = tool.name
        end
        table.sort(names)
        return text_stream(table.concat(names, ","))
    end
    if verb == "find" then
        local hit = haystack(request):find(rest, 1, true) ~= nil
        return text_stream((hit and "found " or "missing ") .. rest)
    end
    if verb == "tag" then
        local label, needle = rest:match("^(%S+)%s+(.*)$")
        local hit = label ~= nil and haystack(request):find(needle, 1, true) ~= nil
        return text_stream(tostring(label) .. (hit and ":found" or ":missing"))
    end
    if verb == "sys" then
        local hit = (request.system_prompt or ""):find(rest, 1, true) ~= nil
        return text_stream((hit and "sys-found " or "sys-missing ") .. rest)
    end
    if verb == "model" then
        return text_stream(string.format(
            "model=%s effort=%s", tostring(request.model), tostring(request.effort or "")))
    end
    if verb == "err" then
        return stream_of({ function()
            return { type = "error", kind = "terminal", message = rest }
        end })
    end

    return stream_of({ function()
        return { type = "error", kind = "invalid_request", message = "unknown verb " .. verb }
    end })
end

-- ---------------------------------------------------------------------------
-- Registration
-- ---------------------------------------------------------------------------
--
-- Every candidate source is evaluated before the allow/deny policy runs, so
-- registration belongs in activate(), never at file scope.

local function activate()
    panto.ext.register_protocol {
        name = "fixture",
        effort_levels = {
            { label = "tiny", detail = "quick" },
            { label = "deep", detail = "thorough" },
        },
        open = function(request)
            local script = prompt_of(request):match("^([^\n]*)") or ""
            local rounds = split(script, ";;")
            local round = round_of(request)
            -- Past the end of the script the turn reports its tool results and
            -- stops. That keeps a tool round trip to one scripted step, and it
            -- is what makes an unexpected extra round show up as a wrong answer
            -- instead of an endless loop.
            local step = trim(rounds[round] or "results")
            mark(string.format("open %s round=%d", tostring(request.model), round))
            return build(request, step)
        end,
    }

    panto.ext.register_tool {
        name = "echo",
        description = "echoes its input",
        schema = { type = "object", properties = { text = { type = "string" } } },
        handler = function(input)
            return "echo:" .. tostring(input and input.text)
        end,
    }

    panto.ext.register_tool {
        name = "whoami",
        description = "reports the calling agent's session id",
        schema = { type = "object" },
        handler = function()
            return panto.ext.agent:session_id()
        end,
    }

    panto.ext.register_tool {
        name = "pair",
        description = "rendezvous with a sibling child",
        schema = { type = "object", properties = { text = { type = "string" } } },
        handler = function(input)
            -- Mark the window this handler spends parked. Two windows can only
            -- interleave if two children are genuinely in flight together; a
            -- serialized pair would mark enter/leave/enter/leave instead.
            local tag = tostring(input and input.text)
            _G.e2e_paired = (_G.e2e_paired or 0) + 1
            mark("enter " .. tag)
            local spins = 0
            while (_G.e2e_paired or 0) < 2 and spins < 400 do
                sleep_ms(5)
                spins = spins + 1
            end
            sleep_ms(20)
            mark("leave " .. tag)
            return "pair:" .. tag
        end,
    }
end

return { name = "fixture", activate = activate }