//! The Agent owns the conversation-driving loop: provider streaming +
//! tool dispatch.
//!
//! In phase 1/2 this was a thin pass-through to the provider. In phase 3
//! it grows the tool-call loop: after each provider streaming step, the
//! agent inspects the assistant message for ToolUse blocks, dispatches
//! the registered handlers (in parallel when there are multiple), and
//! appends a user message containing the ToolResult blocks back into the
//! conversation. The loop continues until a turn arrives with no ToolUse
//! blocks.

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

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

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

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

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

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

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

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

    /// Drive the conversation forward until the model stops calling tools.
    ///
    /// A single `runStep` invocation may call the provider multiple times
    /// if the model chains tool calls. Each provider call streams a new
    /// assistant message into `conv`; if that message contains ToolUse
    /// blocks the agent dispatches them concurrently, appends a user
    /// message of ToolResult blocks, and loops. The loop terminates when
    /// the provider's most recent response has no ToolUse blocks.
    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: if the provider committed an assistant
            // message with zero content blocks, something went wrong
            // upstream that wasn't surfaced as a provider error (e.g. a
            // mid-stream provider error that an older codepath swallowed,
            // or a model that genuinely returned nothing). Either way the
            // turn made no observable progress — surface it instead of
            // silently dropping back to the prompt.
            if (last.content.items.len == 0) return error.EmptyAssistantResponse;

            if (!hasToolUseBlock(last)) return;

            try self.dispatchToolCalls(conv, last);
            // Loop: feed the ToolResult message back to the provider.
        }
    }

    /// Returns true if the message contains at least one ToolUse block.
    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` concurrently, then
    /// append a single user Message containing all ToolResult blocks to
    /// `conv` in the same order the tool calls appeared.
    fn dispatchToolCalls(
        self: *Agent,
        conv: *conversation.Conversation,
        assistant_msg: conversation.Message,
    ) !void {
        // Count tool uses for sizing.
        var n: usize = 0;
        for (assistant_msg.content.items) |block| {
            if (block == .ToolUse) n += 1;
        }
        std.debug.assert(n > 0);

        const tasks = try self.allocator.alloc(ToolCallTask, n);
        defer self.allocator.free(tasks);

        // Populate tasks. We borrow ID/name slices from the conversation —
        // the assistant message stays in `conv` throughout dispatch, so
        // these slices remain valid until we copy them into the new
        // ToolResultBlock.
        {
            var i: usize = 0;
            for (assistant_msg.content.items) |block| {
                if (block != .ToolUse) continue;
                const tu = block.ToolUse;
                tasks[i] = .{
                    .agent = self,
                    .tool_use_id = tu.id,
                    .tool_name = tu.name,
                    .input = tu.input.items,
                    .result = null,
                    .err = null,
                };
                i += 1;
            }
        }

        // Spawn one thread per tool call. `std.Thread.spawn` is cheap
        // (sub-millisecond on Linux/macOS) compared to typical tool
        // latency, and `Tool.invoke` is contractually thread-safe, so we
        // fan out without a pool.
        const threads = try self.allocator.alloc(Thread, n);
        defer self.allocator.free(threads);

        var spawned: usize = 0;
        var joined = false;
        errdefer {
            // Join any in-flight threads so they don't outlive `tasks`.
            if (!joined) for (threads[0..spawned]) |t| t.join();
            for (tasks) |*task| {
                if (task.result) |r| self.allocator.free(r);
            }
        }

        for (tasks, 0..) |*task, idx| {
            threads[idx] = try Thread.spawn(.{}, runToolTask, .{task});
            spawned += 1;
        }
        for (threads[0..spawned]) |t| t.join();
        joined = true;

        // Build the user ToolResult message. From here on we own all
        // result byte slices; transfer them into ToolResultBlocks.
        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, n);

        // If any task failed, prefer to abort the turn — but first move
        // every successful result into a block so it gets freed by the
        // standard cleanup path, and free errored ones eagerly (there
        // are none to move). The errdefer above handles teardown.
        var first_err: ?anyerror = null;
        for (tasks) |*task| {
            if (task.err) |e| {
                first_err = e;
                continue;
            }
            const result_bytes = task.result.?;
            task.result = null; // ownership transferred below

            const id_copy = try self.allocator.dupe(u8, task.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;

        // Wrap the ToolResult blocks into a user Message and append.
        try conv.messages.append(self.allocator, .{
            .role = .user,
            .content = content,
        });
    }
};

/// Per-tool-call work item passed into a worker thread.
const ToolCallTask = struct {
    agent: *Agent,
    tool_use_id: []const u8, // borrowed from assistant_msg
    tool_name: []const u8, // borrowed from assistant_msg
    input: []const u8, // borrowed from assistant_msg

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

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

fn runToolTask(task: *ToolCallTask) void {
    const tool = task.agent.registry.lookup(task.tool_name) orelse {
        task.err = error.UnknownTool;
        return;
    };
    const out = tool.vtable.invoke(tool.ctx, task.input, task.agent.allocator) catch |e| {
        task.err = e;
        return;
    };
    task.result = out;
}

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

const testing = std.testing;

/// A stub Provider that, on each call to `streamStep`, appends a
/// pre-canned assistant message to the conversation. Used to drive the
/// agent's tool-call loop without any HTTP plumbing.
const StubProvider = struct {
    allocator: Allocator,
    scripted: []const ScriptedTurn,
    next: usize = 0,

    const ScriptedTurn = struct {
        /// Blocks to append as the next assistant message. The producer
        /// owns these — the stub clones them per turn so the conversation
        /// can take ownership.
        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 {}
};

/// Simple in-test tool: returns `prefix ++ input`. Used in dispatch tests.
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 .{
            .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);
    }
};

/// Tool that records the thread it ran on, then participates in a
/// rendezvous: every invocation must reach the barrier before any can
/// return. If dispatch is sequential, the first invocation would deadlock
/// (only one tool runs at a time, never reaching the threshold) — so this
/// test only passes when invocations run truly concurrently.
///
/// The barrier is bounded by a spin-with-yield with a wall-time ceiling
/// of 5 seconds; failure to reach quorum surfaces as an `error.BarrierTimeout`.
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 .{
            .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);
        }

        // Spin-with-yield until everyone has arrived. ~5s ceiling at the
        // typical yield granularity is plenty for a 3-way barrier; on a
        // truly single-threaded dispatch this loop never resolves.
        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);
    }
};

/// Tool whose invoke always errors. Used to verify the turn aborts.
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 .{
            .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, _: ?provider_mod.BlockMeta) 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 {}
};

test "registerTool and lookup via registry" {
    var stub = StubProvider{ .allocator = testing.allocator, .scripted = &.{} };
    var agent = Agent.init(testing.allocator, 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 agent = Agent.init(testing.allocator, 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 agent = Agent.init(allocator, 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);

    // user, assistant(tool_use), user(tool_result), assistant(text)
    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;

    // Use a barrier: each tool must wait until all three have arrived
    // before returning. If dispatch were sequential, the first tool
    // would hit its iteration ceiling and `error.BarrierTimeout`. Reaching
    // the barrier proves all three ran concurrently.
    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 agent = Agent.init(allocator, 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);

    // Each tool produced one ToolResult, in original 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("b", tr_msg.content.items[1].ToolResult.tool_use_id);
    try testing.expectEqualStrings("c", tr_msg.content.items[2].ToolResult.tool_use_id);

    // And the three calls happened on three distinct threads.
    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 = "" } },
        } },
        // Second turn should never run.
        .{ .blocks = &.{.{ .Text = "should-not-see" }} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var agent = Agent.init(allocator, 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));

    // Conversation has user + assistant(tool_use). No ToolResult message
    // was appended because the dispatch errored before append.
    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 agent = Agent.init(allocator, 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 agent = Agent.init(allocator, 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" {
    // Mirrors the real-world failure mode where a provider silently ends the
    // turn with no content blocks — e.g. a mid-stream error that an older
    // codepath swallowed. The agent must surface the failure so the user
    // doesn't see the prompt come back with no explanation.
    const allocator = testing.allocator;

    const scripted = [_]StubProvider.ScriptedTurn{
        .{ .blocks = &.{} },
    };
    var stub = StubProvider{ .allocator = allocator, .scripted = &scripted };
    var agent = Agent.init(allocator, 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));
}