//! The Agent owns the conversation-driving loop: provider streaming +
//! tool dispatch.
//!
//! On each turn, after the provider streams an assistant message, the
//! agent inspects it for ToolUse blocks. If any are present, the agent:
//!
//!   1. Groups them by their *owning registration* in the registry — a
//!      single `Tool` is its own group; every `ToolSource`-backed tool
//!      whose name maps to the same source forms one group.
//!   2. Spawns one concurrent task per group via `std.Io.Group`.
//!      A single-`Tool` group runs the tool's `invoke` once; a
//!      `ToolSource` group calls the source's `invoke_batch` with all
//!      of its calls at once. We use `Group.concurrent` (not `async`)
//!      because tool invocations may block on I/O and we need real
//!      concurrency, not just expressed asynchrony.
//!   3. Awaits the group. ToolResult blocks are assembled in the
//!      *original* call order (i.e. the order the LLM emitted them).
//!   4. Appends a user message containing the ToolResult blocks back
//!      into the conversation and loops.
//!
//! The "thread-safe" promise for single `Tool` registrations is
//! unchanged. For `ToolSource`-backed tools, the source's runtime
//! receives all of its calls on one thread per turn, so it can keep a
//! single-threaded interpreter (Lua, Python, ...) without further
//! synchronization.

const std = @import("std");
const Allocator = std.mem.Allocator;
const Io = std.Io;

const provider_mod = @import("provider.zig");
const conversation = @import("conversation.zig");
const tool_mod = @import("tool.zig");
const tool_source_mod = @import("tool_source.zig");
const tool_registry_mod = @import("tool_registry.zig");

pub const Tool = tool_mod.Tool;
pub const ToolSource = tool_source_mod.ToolSource;
pub const ToolRegistry = tool_registry_mod.ToolRegistry;

const Entry = tool_registry_mod.Entry;

pub const Agent = struct {
    provider: provider_mod.Provider,
    allocator: Allocator,
    io: Io,
    registry: ToolRegistry,

    pub fn init(allocator: Allocator, io: Io, prov: provider_mod.Provider) Agent {
        return .{
            .provider = prov,
            .allocator = allocator,
            .io = io,
            .registry = ToolRegistry.init(allocator),
        };
    }

    pub fn deinit(self: *Agent) void {
        self.registry.deinit();
        self.provider.deinit();
    }

    /// Register a single tool. The agent's registry takes ownership.
    pub fn registerTool(self: *Agent, tool: Tool) !void {
        try self.registry.register(tool);
    }

    /// Register a tool source. The agent's registry takes ownership.
    pub fn registerToolSource(self: *Agent, src: ToolSource) !void {
        try self.registry.registerSource(src);
    }

    /// Remove a tool by name. No-op if not registered or if the name
    /// belongs to a source.
    pub fn unregisterTool(self: *Agent, name: []const u8) void {
        self.registry.unregister(name);
    }

    /// Drive the conversation forward until the model stops calling tools.
    pub fn runStep(
        self: *Agent,
        conv: *conversation.Conversation,
        receiver: *provider_mod.Receiver,
    ) !void {
        while (true) {
            try self.provider.streamStep(conv, &self.registry, receiver);

            const last = conv.messages.items[conv.messages.items.len - 1];
            std.debug.assert(last.role == .assistant);

            // Defense-in-depth: a provider that silently committed an
            // empty assistant message means the turn made no observable
            // progress. Surface it instead of looping back to the prompt.
            if (last.content.items.len == 0) return error.EmptyAssistantResponse;

            if (!hasToolUseBlock(last)) return;

            try self.dispatchToolCalls(conv, last);
        }
    }

    fn hasToolUseBlock(msg: conversation.Message) bool {
        for (msg.content.items) |block| {
            if (block == .ToolUse) return true;
        }
        return false;
    }

    /// Dispatch every ToolUse block in `assistant_msg`. Groups by owning
    /// registration; one OS thread per group; results assembled in the
    /// original call order.
    fn dispatchToolCalls(
        self: *Agent,
        conv: *conversation.Conversation,
        assistant_msg: conversation.Message,
    ) !void {
        // Build the flat call list (in original order) and group calls
        // by owning registration.
        var calls: std.array_list.Managed(FlatCall) = .init(self.allocator);
        defer calls.deinit();

        for (assistant_msg.content.items) |block| {
            if (block != .ToolUse) continue;
            const tu = block.ToolUse;
            const entry = self.registry.lookup(tu.name) orelse {
                // Unknown tool: abort the turn with a clear error.
                return error.UnknownTool;
            };
            try calls.append(.{
                .tool_use_id = tu.id,
                .tool_name = tu.name,
                .input = tu.input.items,
                .entry = entry.entry,
                .result = null,
                .err = null,
            });
        }
        std.debug.assert(calls.items.len > 0);

        // Partition into groups. A group's `kind` determines how it
        // runs; the `member_indices` are positions into `calls` (the
        // original call order) so we can write back results without
        // re-ordering.
        var groups: std.array_list.Managed(Group) = .init(self.allocator);
        defer {
            for (groups.items) |*g| g.deinit(self.allocator);
            groups.deinit();
        }
        try buildGroups(self.allocator, calls.items, &groups);

        // Spawn one concurrent task per group via `std.Io.Group`.
        // Single-tool groups run the tool's vtable; source groups run
        // the source's `invoke_batch`. We use `concurrent` rather than
        // `async` because tool work may block on I/O — under a
        // single-threaded `Io` `async` would deadlock; `concurrent`
        // forces real concurrency (or `error.ConcurrencyUnavailable`).
        var task_group: Io.Group = .init;
        // `cancel` is idempotent with `await`; if anything below this
        // point errors before we successfully `await`, this releases
        // the group's resources.
        defer task_group.cancel(self.io);
        errdefer {
            for (calls.items) |*c| {
                if (c.result) |r| self.allocator.free(r);
            }
        }

        for (groups.items) |*g| {
            try task_group.concurrent(self.io, runGroup, .{ self, g, calls.items });
        }
        // `error.Canceled` here means cancellation propagated into this
        // dispatch from above; surface it like any other error.
        try task_group.await(self.io);

        // Assemble ToolResult blocks in original call order. If any
        // call errored, prefer to abort the turn — but only after the
        // standard errdefer above has freed remaining results.
        var content: std.ArrayList(conversation.ContentBlock) = .empty;
        errdefer {
            for (content.items) |*b| b.deinit(self.allocator);
            content.deinit(self.allocator);
        }
        try content.ensureTotalCapacity(self.allocator, calls.items.len);

        var first_err: ?anyerror = null;
        for (calls.items) |*c| {
            if (c.err) |e| {
                first_err = e;
                continue;
            }
            const result_bytes = c.result orelse {
                // Internal error: every successful call should have left
                // bytes behind. Treat as MissingToolResult.
                first_err = error.MissingToolResult;
                continue;
            };
            c.result = null; // ownership transferred below

            const id_copy = try self.allocator.dupe(u8, c.tool_use_id);
            errdefer self.allocator.free(id_copy);

            var content_buf: conversation.TextualBlock = .empty;
            errdefer content_buf.deinit(self.allocator);
            try content_buf.appendSlice(self.allocator, result_bytes);
            self.allocator.free(result_bytes);

            content.appendAssumeCapacity(.{ .ToolResult = .{
                .tool_use_id = id_copy,
                .content = content_buf,
            } });
        }

        if (first_err) |e| return e;

        try conv.messages.append(self.allocator, .{
            .role = .user,
            .content = content,
        });
    }
};

/// One ToolUse, as flattened into the agent's dispatch list. `result`
/// and `err` are filled in by the worker; exactly one is non-null on
/// successful task completion.
const FlatCall = struct {
    tool_use_id: []const u8, // borrowed from assistant_msg
    tool_name: []const u8, // borrowed from assistant_msg
    input: []const u8, // borrowed from assistant_msg
    entry: Entry,

    /// Owned result bytes from `Tool.invoke` or `ToolSource.invoke_batch`.
    /// Allocated with the agent's allocator. Transferred into a
    /// ToolResultBlock on success.
    result: ?[]u8,

    /// If non-null, the call failed and the turn must abort.
    err: ?anyerror,
};

/// One dispatch group. Either a single Tool invocation, or a batch of
/// calls headed to one ToolSource.
const Group = union(enum) {
    single: SingleGroup,
    source: SourceGroup,

    pub const SingleGroup = struct {
        tool: Tool,
        /// Index into the flat calls array.
        call_index: usize,
    };

    pub const SourceGroup = struct {
        source: *ToolSource,
        /// Indices into the flat calls array. Owned by the group.
        member_indices: []usize,
    };

    fn deinit(self: *Group, allocator: Allocator) void {
        switch (self.*) {
            .single => {},
            .source => |sg| allocator.free(sg.member_indices),
        }
    }
};

/// Partition the flat call list into groups. Order of groups is
/// arbitrary; order within a `source` group preserves the original
/// call order so that batch results can be written back positionally.
fn buildGroups(
    allocator: Allocator,
    calls: []const FlatCall,
    out: *std.array_list.Managed(Group),
) !void {
    // Map from source pointer to the index of its group in `out`.
    // Buffers per source, accumulated then frozen into slices.
    var pending: std.AutoHashMap(*ToolSource, std.array_list.Managed(usize)) =
        .init(allocator);
    defer {
        var it = pending.valueIterator();
        while (it.next()) |l| l.deinit();
        pending.deinit();
    }

    for (calls, 0..) |c, i| {
        switch (c.entry) {
            .single => |t| try out.append(.{ .single = .{ .tool = t, .call_index = i } }),
            .source => |sr| {
                const gop = try pending.getOrPut(sr.source);
                if (!gop.found_existing) {
                    gop.value_ptr.* = std.array_list.Managed(usize).init(allocator);
                }
                try gop.value_ptr.append(i);
            },
        }
    }

    // Freeze each pending list into a source-group entry. We move
    // ownership of the indices into `Group.source.member_indices`.
    var pit = pending.iterator();
    while (pit.next()) |entry| {
        const src = entry.key_ptr.*;
        const indices = try entry.value_ptr.toOwnedSlice();
        try out.append(.{ .source = .{ .source = src, .member_indices = indices } });
    }
}

/// Worker entry point. Runs one group to completion, populating
/// `calls[i].result` or `calls[i].err` for each member call.
///
/// Return type is `void`, which coerces to `Io.Cancelable!void` as
/// required by `Group.concurrent`. Tool errors are reported via
/// `FlatCall.err`, not by returning from this function.
fn runGroup(agent: *Agent, group: *Group, calls: []FlatCall) void {
    switch (group.*) {
        .single => |sg| {
            const i = sg.call_index;
            const c = &calls[i];
            const out = sg.tool.vtable.invoke(sg.tool.ctx, c.input, agent.allocator) catch |e| {
                c.err = e;
                return;
            };
            c.result = out;
        },
        .source => |sg| runSourceGroup(agent, sg, calls),
    }
}

fn runSourceGroup(agent: *Agent, sg: Group.SourceGroup, calls: []FlatCall) void {
    const n = sg.member_indices.len;

    const batch_calls = agent.allocator.alloc(tool_source_mod.Call, n) catch |e| {
        for (sg.member_indices) |i| calls[i].err = e;
        return;
    };
    defer agent.allocator.free(batch_calls);

    const batch_results = agent.allocator.alloc(tool_source_mod.CallResult, n) catch |e| {
        for (sg.member_indices) |i| calls[i].err = e;
        return;
    };
    defer agent.allocator.free(batch_results);

    for (sg.member_indices, 0..) |idx, j| {
        batch_calls[j] = .{
            .tool_name = calls[idx].tool_name,
            .input = calls[idx].input,
        };
        batch_results[j] = .{ .err = error.SourceDroppedCall };
    }

    sg.source.vtable.invoke_batch(
        sg.source.ctx,
        batch_calls,
        batch_results,
        agent.allocator,
    ) catch |e| {
        // Whole-batch failure: free any partial successes the source
        // already wrote, then mark every member as failed.
        for (batch_results) |r| switch (r) {
            .ok => |b| agent.allocator.free(b),
            .err => {},
        };
        for (sg.member_indices) |i| calls[i].err = e;
        return;
    };

    // Per-call success/error.
    for (sg.member_indices, 0..) |i, j| {
        switch (batch_results[j]) {
            .ok => |b| calls[i].result = b,
            .err => |e| calls[i].err = e,
        }
    }
}

// -----------------------------------------------------------------------------
// Tests
// -----------------------------------------------------------------------------

const testing = std.testing;

const StubProvider = struct {
    allocator: Allocator,
    scripted: []const ScriptedTurn,
    next: usize = 0,

    const ScriptedTurn = struct {
        blocks: []const TestBlock,
    };

    const TestBlock = union(enum) {
        Text: []const u8,
        ToolUse: struct {
            id: []const u8,
            name: []const u8,
            input: []const u8,
        },
    };

    fn provider(self: *StubProvider) provider_mod.Provider {
        return .{ .ptr = self, .vtable = &vt };
    }

    const vt: provider_mod.ProviderVTable = .{
        .streamStep = vtStreamStep,
        .deinit = vtDeinit,
    };

    fn vtStreamStep(
        ptr: *anyopaque,
        conv: *conversation.Conversation,
        _: *const ToolRegistry,
        _: *provider_mod.Receiver,
    ) anyerror!void {
        const self: *StubProvider = @ptrCast(@alignCast(ptr));
        if (self.next >= self.scripted.len) return error.NoMoreScriptedTurns;
        const turn = self.scripted[self.next];
        self.next += 1;

        var blocks: std.ArrayList(conversation.ContentBlock) = .empty;
        errdefer {
            for (blocks.items) |*b| b.deinit(self.allocator);
            blocks.deinit(self.allocator);
        }
        for (turn.blocks) |tb| {
            switch (tb) {
                .Text => |s| {
                    try blocks.append(self.allocator, .{
                        .Text = try conversation.textualBlockFromSlice(self.allocator, s),
                    });
                },
                .ToolUse => |tu| {
                    const id = try self.allocator.dupe(u8, tu.id);
                    errdefer self.allocator.free(id);
                    const name = try self.allocator.dupe(u8, tu.name);
                    errdefer self.allocator.free(name);
                    var input_buf: conversation.TextualBlock = .empty;
                    errdefer input_buf.deinit(self.allocator);
                    try input_buf.appendSlice(self.allocator, tu.input);
                    try blocks.append(self.allocator, .{ .ToolUse = .{
                        .id = id,
                        .name = name,
                        .input = input_buf,
                    } });
                },
            }
        }
        const moved = try blocks.toOwnedSlice(self.allocator);
        defer self.allocator.free(moved);
        try conv.addAssistantMessage(moved);
    }

    fn vtDeinit(_: *anyopaque) void {}
};

const EchoTool = struct {
    prefix_owned: []u8,
    name_owned: []u8,

    fn create(allocator: Allocator, name: []const u8, prefix: []const u8) !Tool {
        const self = try allocator.create(EchoTool);
        errdefer allocator.destroy(self);
        self.name_owned = try allocator.dupe(u8, name);
        errdefer allocator.free(self.name_owned);
        self.prefix_owned = try allocator.dupe(u8, prefix);
        return .{
            .decl = .{
                .name = self.name_owned,
                .description = "echo",
                .schema_json = "{}",
            },
            .ctx = self,
            .vtable = &vt,
        };
    }

    const vt: Tool.VTable = .{ .invoke = invoke, .deinit = deinit };

    fn invoke(ctx: *anyopaque, input: []const u8, allocator: Allocator) anyerror![]u8 {
        const self: *EchoTool = @ptrCast(@alignCast(ctx));
        return try std.fmt.allocPrint(allocator, "{s}{s}", .{ self.prefix_owned, input });
    }

    fn deinit(ctx: *anyopaque, allocator: Allocator) void {
        const self: *EchoTool = @ptrCast(@alignCast(ctx));
        allocator.free(self.name_owned);
        allocator.free(self.prefix_owned);
        allocator.destroy(self);
    }
};

const BarrierTool = struct {
    name_owned: []u8,
    barrier: *Barrier,

    const Barrier = struct {
        target: u32,
        arrived: std.atomic.Value(u32) = .init(0),
        thread_ids: [4]std.atomic.Value(u64) = .{
            .init(0), .init(0), .init(0), .init(0),
        },
    };

    fn create(allocator: Allocator, name: []const u8, barrier: *Barrier) !Tool {
        const self = try allocator.create(BarrierTool);
        errdefer allocator.destroy(self);
        self.name_owned = try allocator.dupe(u8, name);
        self.barrier = barrier;
        return .{
            .decl = .{
                .name = self.name_owned,
                .description = "barrier",
                .schema_json = "{}",
            },
            .ctx = self,
            .vtable = &vt,
        };
    }

    const vt: Tool.VTable = .{ .invoke = invoke, .deinit = deinit };

    fn invoke(ctx: *anyopaque, _: []const u8, allocator: Allocator) anyerror![]u8 {
        const self: *BarrierTool = @ptrCast(@alignCast(ctx));
        const arrived = self.barrier.arrived.fetchAdd(1, .acq_rel);
        if (arrived < self.barrier.thread_ids.len) {
            self.barrier.thread_ids[arrived].store(std.Thread.getCurrentId(), .release);
        }

        var i: usize = 0;
        while (self.barrier.arrived.load(.acquire) < self.barrier.target) : (i += 1) {
            if (i > 50_000) return error.BarrierTimeout;
            std.Thread.yield() catch {};
        }
        return try allocator.dupe(u8, "done");
    }

    fn deinit(ctx: *anyopaque, allocator: Allocator) void {
        const self: *BarrierTool = @ptrCast(@alignCast(ctx));
        allocator.free(self.name_owned);
        allocator.destroy(self);
    }
};

const FailingTool = struct {
    name_owned: []u8,

    fn create(allocator: Allocator, name: []const u8) !Tool {
        const self = try allocator.create(FailingTool);
        errdefer allocator.destroy(self);
        self.name_owned = try allocator.dupe(u8, name);
        return .{
            .decl = .{
                .name = self.name_owned,
                .description = "fails",
                .schema_json = "{}",
            },
            .ctx = self,
            .vtable = &vt,
        };
    }

    const vt: Tool.VTable = .{ .invoke = invoke, .deinit = deinit };

    fn invoke(_: *anyopaque, _: []const u8, _: Allocator) anyerror![]u8 {
        return error.ToolExploded;
    }

    fn deinit(ctx: *anyopaque, allocator: Allocator) void {
        const self: *FailingTool = @ptrCast(@alignCast(ctx));
        allocator.free(self.name_owned);
        allocator.destroy(self);
    }
};

const NoopReceiver = struct {
    fn make() provider_mod.Receiver {
        return .{ .ptr = @constCast(@ptrCast(&dummy)), .vtable = &vt };
    }
    var dummy: u8 = 0;
    const vt: provider_mod.ReceiverVTable = .{
        .onMessageStart = noop1,
        .onBlockStart = noop2,
        .onContentDelta = noop3,
        .onBlockComplete = noop4,
        .onMessageComplete = noop5,
        .onError = noop6,
    };
    fn noop1(_: *anyopaque, _: conversation.MessageRole) anyerror!void {}
    fn noop2(_: *anyopaque, _: provider_mod.ContentBlockType, _: usize) anyerror!void {}
    fn noop3(_: *anyopaque, _: usize, _: []const u8) anyerror!void {}
    fn noop4(_: *anyopaque, _: usize, _: conversation.ContentBlock) anyerror!void {}
    fn noop5(_: *anyopaque, _: conversation.Message) anyerror!void {}
    fn noop6(_: *anyopaque, _: anyerror) void {}
};

/// A configurable ToolSource for testing the grouped-dispatch path.
/// Stores every batch it receives so tests can assert "calls X and Y
/// arrived in the same batch on the same thread".
const TestSource = struct {
    name_owned: []u8,
    decls: []tool_source_mod.ToolDecl,
    decl_strings: std.array_list.Managed([]u8),
    /// Sequence of (thread_id, [tool_name; n]) per batch received.
    /// Only mutated inside `invoke_batch`. Because libpanto guarantees
    /// at most one outstanding `invoke_batch` per source at any time
    /// (one batch per turn per source), no synchronization is needed.
    batches: std.array_list.Managed(Batch),
    allocator: Allocator,

    const Batch = struct {
        thread_id: u64,
        names: std.array_list.Managed([]u8),
    };

    fn create(
        allocator: Allocator,
        source_name: []const u8,
        tool_names: []const []const u8,
    ) !ToolSource {
        const self = try allocator.create(TestSource);
        errdefer allocator.destroy(self);

        var strings = std.array_list.Managed([]u8).init(allocator);
        errdefer {
            for (strings.items) |s| allocator.free(s);
            strings.deinit();
        }

        const name_owned = try allocator.dupe(u8, source_name);
        try strings.append(name_owned);

        const decls = try allocator.alloc(tool_source_mod.ToolDecl, tool_names.len);
        errdefer allocator.free(decls);
        for (tool_names, 0..) |tn, i| {
            const n = try allocator.dupe(u8, tn);
            try strings.append(n);
            const d = try allocator.dupe(u8, "test src tool");
            try strings.append(d);
            const s = try allocator.dupe(u8, "{}");
            try strings.append(s);
            decls[i] = .{ .name = n, .description = d, .schema_json = s };
        }

        self.* = .{
            .name_owned = name_owned,
            .decls = decls,
            .decl_strings = strings,
            .batches = std.array_list.Managed(Batch).init(allocator),
            .allocator = allocator,
        };

        return ToolSource{
            .name = self.name_owned,
            .tools = self.decls,
            .ctx = self,
            .vtable = &vt,
        };
    }

    const vt: ToolSource.VTable = .{
        .invoke_batch = invokeBatch,
        .deinit = deinitSrc,
    };

    fn invokeBatch(
        ctx: *anyopaque,
        calls: []const tool_source_mod.Call,
        results: []tool_source_mod.CallResult,
        allocator: Allocator,
    ) anyerror!void {
        const self: *TestSource = @ptrCast(@alignCast(ctx));
        var batch: Batch = .{
            .thread_id = std.Thread.getCurrentId(),
            .names = std.array_list.Managed([]u8).init(self.allocator),
        };
        for (calls) |c| {
            const copy = try self.allocator.dupe(u8, c.tool_name);
            try batch.names.append(copy);
        }
        try self.batches.append(batch);

        for (calls, 0..) |c, i| {
            results[i] = .{
                .ok = std.fmt.allocPrint(
                    allocator,
                    "{s}->{s}",
                    .{ c.tool_name, c.input },
                ) catch |e| {
                    results[i] = .{ .err = e };
                    continue;
                },
            };
        }
    }

    fn deinitSrc(ctx: *anyopaque, _: Allocator) void {
        const self: *TestSource = @ptrCast(@alignCast(ctx));
        for (self.decl_strings.items) |s| self.allocator.free(s);
        self.decl_strings.deinit();
        for (self.batches.items) |*b| {
            for (b.names.items) |n| self.allocator.free(n);
            b.names.deinit();
        }
        self.batches.deinit();
        self.allocator.free(self.decls);
        self.allocator.destroy(self);
    }
};

/// A source that always fails the whole batch by returning an error
/// from invoke_batch (rather than recording per-call errors). Used to
/// verify libpanto's whole-batch-failure path.
const FailingSource = struct {
    name_owned: []u8,
    decls: []tool_source_mod.ToolDecl,
    decl_strings: std.array_list.Managed([]u8),
    allocator: Allocator,

    fn create(allocator: Allocator, source_name: []const u8, tool_names: []const []const u8) !ToolSource {
        const self = try allocator.create(FailingSource);
        errdefer allocator.destroy(self);

        var strings = std.array_list.Managed([]u8).init(allocator);
        errdefer {
            for (strings.items) |s| allocator.free(s);
            strings.deinit();
        }

        const name_owned = try allocator.dupe(u8, source_name);
        try strings.append(name_owned);

        const decls = try allocator.alloc(tool_source_mod.ToolDecl, tool_names.len);
        errdefer allocator.free(decls);
        for (tool_names, 0..) |tn, i| {
            const n = try allocator.dupe(u8, tn);
            try strings.append(n);
            const d = try allocator.dupe(u8, "fails");
            try strings.append(d);
            const s = try allocator.dupe(u8, "{}");
            try strings.append(s);
            decls[i] = .{ .name = n, .description = d, .schema_json = s };
        }

        self.* = .{
            .name_owned = name_owned,
            .decls = decls,
            .decl_strings = strings,
            .allocator = allocator,
        };
        return ToolSource{ .name = self.name_owned, .tools = self.decls, .ctx = self, .vtable = &vt };
    }

    const vt: ToolSource.VTable = .{ .invoke_batch = invokeBatch, .deinit = deinitSrc };

    fn invokeBatch(
        _: *anyopaque,
        _: []const tool_source_mod.Call,
        _: []tool_source_mod.CallResult,
        _: Allocator,
    ) anyerror!void {
        return error.SourceExploded;
    }

    fn deinitSrc(ctx: *anyopaque, _: Allocator) void {
        const self: *FailingSource = @ptrCast(@alignCast(ctx));
        for (self.decl_strings.items) |s| self.allocator.free(s);
        self.decl_strings.deinit();
        self.allocator.free(self.decls);
        self.allocator.destroy(self);
    }
};

test "registerTool and lookup via registry" {
    var stub = StubProvider{ .allocator = testing.allocator, .scripted = &.{} };
    var threaded: std.Io.Threaded = .init(testing.allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(testing.allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerTool(try EchoTool.create(testing.allocator, "echo", "ECHO:"));
    try testing.expectEqual(@as(usize, 1), agent.registry.count());
    try testing.expect(agent.registry.lookup("echo") != null);
}

test "duplicate registerTool returns error" {
    var stub = StubProvider{ .allocator = testing.allocator, .scripted = &.{} };
    var threaded: std.Io.Threaded = .init(testing.allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(testing.allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerTool(try EchoTool.create(testing.allocator, "echo", "A:"));

    var dup = try EchoTool.create(testing.allocator, "echo", "B:");
    try testing.expectError(error.DuplicateTool, agent.registerTool(dup));
    dup.vtable.deinit(dup.ctx, testing.allocator);
}

test "runStep dispatches a tool call and loops to a final text turn" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "tc_1", .name = "echo", .input = "hello" } },
        } },
        .{ .blocks = &.{
            .{ .Text = "ok" },
        } },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerTool(try EchoTool.create(allocator, "echo", "ECHO:"));

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("call a tool");

    var recv = NoopReceiver.make();
    try agent.runStep(&conv, &recv);

    try testing.expectEqual(@as(usize, 4), conv.messages.items.len);

    try testing.expectEqual(conversation.MessageRole.assistant, conv.messages.items[1].role);
    try testing.expectEqual(@as(usize, 1), conv.messages.items[1].content.items.len);
    try testing.expectEqualStrings("tc_1", conv.messages.items[1].content.items[0].ToolUse.id);

    try testing.expectEqual(conversation.MessageRole.user, conv.messages.items[2].role);
    try testing.expectEqual(@as(usize, 1), conv.messages.items[2].content.items.len);
    const tr = conv.messages.items[2].content.items[0].ToolResult;
    try testing.expectEqualStrings("tc_1", tr.tool_use_id);
    try testing.expectEqualStrings("ECHO:hello", tr.content.items);

    try testing.expectEqual(conversation.MessageRole.assistant, conv.messages.items[3].role);
    try testing.expectEqualStrings("ok", conv.messages.items[3].content.items[0].Text.items);
}

test "runStep dispatches multiple tool calls in parallel" {
    const allocator = testing.allocator;

    var barrier: BarrierTool.Barrier = .{ .target = 3 };

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "a", .name = "barrierA", .input = "" } },
            .{ .ToolUse = .{ .id = "b", .name = "barrierB", .input = "" } },
            .{ .ToolUse = .{ .id = "c", .name = "barrierC", .input = "" } },
        } },
        .{ .blocks = &.{
            .{ .Text = "done" },
        } },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerTool(try BarrierTool.create(allocator, "barrierA", &barrier));
    try agent.registerTool(try BarrierTool.create(allocator, "barrierB", &barrier));
    try agent.registerTool(try BarrierTool.create(allocator, "barrierC", &barrier));

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("go");

    var recv = NoopReceiver.make();
    try agent.runStep(&conv, &recv);

    const tr_msg = conv.messages.items[2];
    try testing.expectEqual(@as(usize, 3), tr_msg.content.items.len);
    try testing.expectEqualStrings("a", tr_msg.content.items[0].ToolResult.tool_use_id);
    try testing.expectEqualStrings("b", tr_msg.content.items[1].ToolResult.tool_use_id);
    try testing.expectEqualStrings("c", tr_msg.content.items[2].ToolResult.tool_use_id);

    const t0 = barrier.thread_ids[0].load(.acquire);
    const t1 = barrier.thread_ids[1].load(.acquire);
    const t2 = barrier.thread_ids[2].load(.acquire);
    try testing.expect(t0 != 0 and t1 != 0 and t2 != 0);
    try testing.expect(t0 != t1 and t1 != t2 and t0 != t2);
}

test "runStep propagates tool errors and aborts the turn" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "x", .name = "boom", .input = "" } },
        } },
        .{ .blocks = &.{.{ .Text = "should-not-see" }} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerTool(try FailingTool.create(allocator, "boom"));

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("break it");

    var recv = NoopReceiver.make();
    try testing.expectError(error.ToolExploded, agent.runStep(&conv, &recv));

    try testing.expectEqual(@as(usize, 2), conv.messages.items.len);
}

test "runStep errors UnknownTool when the model calls something unregistered" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "z", .name = "ghost", .input = "" } },
        } },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("call a ghost");

    var recv = NoopReceiver.make();
    try testing.expectError(error.UnknownTool, agent.runStep(&conv, &recv));
}

test "runStep with no tool calls returns after one provider step" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{.{ .Text = "hi" }} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("hello");

    var recv = NoopReceiver.make();
    try agent.runStep(&conv, &recv);

    try testing.expectEqual(@as(usize, 2), conv.messages.items.len);
    try testing.expectEqualStrings("hi", conv.messages.items[1].content.items[0].Text.items);
}

test "runStep surfaces EmptyAssistantResponse when provider commits an empty message" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("hi");

    var recv = NoopReceiver.make();
    try testing.expectError(error.EmptyAssistantResponse, agent.runStep(&conv, &recv));
}

// ------------ ToolSource tests ------------

test "runStep delivers all source-backed calls in one batch on one thread" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "a", .name = "lua_x", .input = "1" } },
            .{ .ToolUse = .{ .id = "b", .name = "lua_y", .input = "2" } },
            .{ .ToolUse = .{ .id = "c", .name = "lua_x", .input = "3" } },
        } },
        .{ .blocks = &.{.{ .Text = "done" }} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerToolSource(try TestSource.create(allocator, "panto-lua", &.{ "lua_x", "lua_y" }));

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("go");

    var recv = NoopReceiver.make();
    try agent.runStep(&conv, &recv);

    // Locate the source and inspect its observed batches.
    const view = agent.registry.lookup("lua_x") orelse return error.NotFound;
    const src_ptr = view.entry.source.source;
    const test_src: *TestSource = @ptrCast(@alignCast(src_ptr.ctx));

    try testing.expectEqual(@as(usize, 1), test_src.batches.items.len);
    const b = test_src.batches.items[0];
    try testing.expectEqual(@as(usize, 3), b.names.items.len);
    try testing.expectEqualStrings("lua_x", b.names.items[0]);
    try testing.expectEqualStrings("lua_y", b.names.items[1]);
    try testing.expectEqualStrings("lua_x", b.names.items[2]);

    // ToolResults arrived in the original call order.
    const tr_msg = conv.messages.items[2];
    try testing.expectEqual(@as(usize, 3), tr_msg.content.items.len);
    try testing.expectEqualStrings("a", tr_msg.content.items[0].ToolResult.tool_use_id);
    try testing.expectEqualStrings("lua_x->1", tr_msg.content.items[0].ToolResult.content.items);
    try testing.expectEqualStrings("b", tr_msg.content.items[1].ToolResult.tool_use_id);
    try testing.expectEqualStrings("lua_y->2", tr_msg.content.items[1].ToolResult.content.items);
    try testing.expectEqualStrings("c", tr_msg.content.items[2].ToolResult.tool_use_id);
    try testing.expectEqualStrings("lua_x->3", tr_msg.content.items[2].ToolResult.content.items);
}

test "runStep: distinct sources run on distinct threads in parallel" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "a", .name = "src_a_t", .input = "" } },
            .{ .ToolUse = .{ .id = "b", .name = "src_b_t", .input = "" } },
        } },
        .{ .blocks = &.{.{ .Text = "done" }} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerToolSource(try TestSource.create(allocator, "src_a", &.{"src_a_t"}));
    try agent.registerToolSource(try TestSource.create(allocator, "src_b", &.{"src_b_t"}));

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("go");

    var recv = NoopReceiver.make();
    try agent.runStep(&conv, &recv);

    const view_a = agent.registry.lookup("src_a_t") orelse return error.NotFound;
    const view_b = agent.registry.lookup("src_b_t") orelse return error.NotFound;
    const sa: *TestSource = @ptrCast(@alignCast(view_a.entry.source.source.ctx));
    const sb: *TestSource = @ptrCast(@alignCast(view_b.entry.source.source.ctx));

    try testing.expectEqual(@as(usize, 1), sa.batches.items.len);
    try testing.expectEqual(@as(usize, 1), sb.batches.items.len);
    // The two sources ran on distinct OS threads.
    try testing.expect(sa.batches.items[0].thread_id != sb.batches.items[0].thread_id);
}

test "runStep: source whole-batch error aborts the turn" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "a", .name = "fa", .input = "" } },
            .{ .ToolUse = .{ .id = "b", .name = "fb", .input = "" } },
        } },
        .{ .blocks = &.{.{ .Text = "never" }} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerToolSource(try FailingSource.create(allocator, "fs", &.{ "fa", "fb" }));

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("kaboom");

    var recv = NoopReceiver.make();
    try testing.expectError(error.SourceExploded, agent.runStep(&conv, &recv));

    // Conversation stops at user + assistant(tool_use). No ToolResult appended.
    try testing.expectEqual(@as(usize, 2), conv.messages.items.len);
}

test "runStep: mixed single Tools and source-backed tools coexist in one turn" {
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{
            .{ .ToolUse = .{ .id = "a", .name = "single", .input = "X" } },
            .{ .ToolUse = .{ .id = "b", .name = "src_t1", .input = "Y" } },
            .{ .ToolUse = .{ .id = "c", .name = "src_t2", .input = "Z" } },
        } },
        .{ .blocks = &.{.{ .Text = "done" }} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var threaded: std.Io.Threaded = .init(allocator, .{});
    defer threaded.deinit();
    const io = threaded.io();
    var agent = Agent.init(allocator, io, stub.provider());
    defer agent.deinit();

    try agent.registerTool(try EchoTool.create(allocator, "single", "S:"));
    try agent.registerToolSource(try TestSource.create(allocator, "src", &.{ "src_t1", "src_t2" }));

    var conv = conversation.Conversation.init(allocator);
    defer conv.deinit();
    try conv.addUserMessage("go");

    var recv = NoopReceiver.make();
    try agent.runStep(&conv, &recv);

    const tr_msg = conv.messages.items[2];
    try testing.expectEqual(@as(usize, 3), tr_msg.content.items.len);
    try testing.expectEqualStrings("S:X", tr_msg.content.items[0].ToolResult.content.items);
    try testing.expectEqualStrings("src_t1->Y", tr_msg.content.items[1].ToolResult.content.items);
    try testing.expectEqualStrings("src_t2->Z", tr_msg.content.items[2].ToolResult.content.items);
}