//! Registry of tools owned by an `Agent`.
//!
//! Two kinds of registration coexist:
//!
//!   - A single `Tool`: a thread-safe, self-contained handler. The
//!     registry holds one entry keyed by `tool.decl.name`.
//!   - A `ToolSource`: a batch-dispatched runtime that owns many tools.
//!     The registry holds one entry per declared tool, all pointing back
//!     at the same source (different `tool_index` per entry).
//!
//! Iteration yields the per-tool metadata as a uniform `ToolView` so
//! callers (chiefly: provider request serializers) don't need to know
//! which flavor of registration each tool came from.
//!
//! Iteration is not synchronized — callers must avoid mutating the
//! registry during iteration. In the current agent loop this is naturally
//! true: the provider iterates once at request-build time, and tool
//! registration only happens at agent setup.

const std = @import("std");
const Allocator = std.mem.Allocator;
const tool_mod = @import("tool.zig");
const tool_source_mod = @import("tool_source.zig");

const Tool = tool_mod.Tool;
const ToolSource = tool_source_mod.ToolSource;
const ToolDecl = tool_source_mod.ToolDecl;

/// Tagged registry value. The registry stores one of these per *tool
/// name*. ToolSources expand to one entry per declared tool, each with a
/// distinct `tool_index`.
pub const Entry = union(enum) {
    single: Tool,
    source: SourceRef,

    pub const SourceRef = struct {
        source: *ToolSource,
        /// Index into `source.tools`.
        tool_index: usize,
    };
};

/// Read-only view of a tool's metadata, uniform across `Tool` and
/// `ToolSource` registrations. Returned by registry iteration and
/// lookup.
pub const ToolView = struct {
    decl: ToolDecl,
    /// Which entry this view came from. Carries enough information to
    /// dispatch the call (single Tool vs source-backed).
    entry: Entry,

    pub fn name(self: ToolView) []const u8 {
        return self.decl.name;
    }
};

pub const ToolRegistry = struct {
    /// Per-tool-name entries.
    entries: std.StringHashMap(Entry),
    /// Heap-allocated sources, kept in a list so `deinit` can tear each
    /// down exactly once even though many entries reference a single
    /// source.
    sources: std.array_list.Managed(*ToolSource),
    allocator: Allocator,

    pub fn init(allocator: Allocator) ToolRegistry {
        return .{
            .entries = std.StringHashMap(Entry).init(allocator),
            .sources = std.array_list.Managed(*ToolSource).init(allocator),
            .allocator = allocator,
        };
    }

    /// Tear down the registry. Each single `Tool`'s `vtable.deinit` is
    /// invoked once. Each `ToolSource`'s `vtable.deinit` is invoked once
    /// (not once per declared tool).
    pub fn deinit(self: *ToolRegistry) void {
        var it = self.entries.iterator();
        while (it.next()) |entry| {
            switch (entry.value_ptr.*) {
                .single => |t| t.vtable.deinit(t.ctx, self.allocator),
                .source => {},
            }
        }
        self.entries.deinit();

        for (self.sources.items) |src| {
            src.vtable.deinit(src.ctx, self.allocator);
            self.allocator.destroy(src);
        }
        self.sources.deinit();
    }

    /// Register a single tool. The registry takes ownership.
    ///
    /// Returns `error.DuplicateTool` if a tool with the same name is
    /// already registered (whether from a single Tool or from a source).
    /// In the duplicate case the caller's tool is NOT taken over; the
    /// caller is responsible for tearing it down.
    pub fn register(self: *ToolRegistry, tool: Tool) !void {
        const gop = try self.entries.getOrPut(tool.decl.name);
        if (gop.found_existing) return error.DuplicateTool;
        gop.value_ptr.* = .{ .single = tool };
    }

    /// Register a tool source. The registry takes ownership of `src` —
    /// it is heap-copied into the registry's source list and freed at
    /// deinit.
    ///
    /// Returns `error.DuplicateTool` if any of the source's declared
    /// tools collides with an existing registration. On collision the
    /// source is NOT taken over (caller still owns it and must tear it
    /// down) and any tools that *had* been inserted before the collision
    /// are rolled back.
    pub fn registerSource(self: *ToolRegistry, src: ToolSource) !void {
        // First pass: check for any collision before committing.
        for (src.tools) |decl| {
            if (self.entries.contains(decl.name)) return error.DuplicateTool;
        }

        // Allocate the persistent heap copy of the source. From this
        // point forward, on any failure we must free the allocation and
        // roll back any entries we inserted.
        const heap = try self.allocator.create(ToolSource);
        errdefer self.allocator.destroy(heap);
        heap.* = src;

        var inserted: usize = 0;
        errdefer {
            // Roll back any inserts we made before the failure.
            for (src.tools[0..inserted]) |decl| {
                _ = self.entries.remove(decl.name);
            }
        }

        for (src.tools, 0..) |decl, i| {
            const gop = try self.entries.getOrPut(decl.name);
            if (gop.found_existing) return error.DuplicateTool;
            gop.value_ptr.* = .{ .source = .{ .source = heap, .tool_index = i } };
            inserted = i + 1;
        }

        try self.sources.append(heap);
    }

    /// Remove a single-tool registration by name. Calls the tool's
    /// `vtable.deinit`. No-op if the name is not registered or if it
    /// belongs to a source (sources are removed as a unit; not yet
    /// exposed).
    pub fn unregister(self: *ToolRegistry, name: []const u8) void {
        const entry_ptr = self.entries.getPtr(name) orelse return;
        switch (entry_ptr.*) {
            .single => |t| {
                _ = self.entries.remove(name);
                t.vtable.deinit(t.ctx, self.allocator);
            },
            .source => {}, // ignore — sources tear down at registry deinit
        }
    }

    /// Look up a tool by name. Returns a uniform `ToolView`. Pointer
    /// invariants are the same as `std.StringHashMap.getPtr`: invalidated
    /// by subsequent register/unregister calls.
    pub fn lookup(self: *const ToolRegistry, name: []const u8) ?ToolView {
        const entry = self.entries.get(name) orelse return null;
        return makeView(entry);
    }

    pub fn count(self: *const ToolRegistry) usize {
        return self.entries.count();
    }

    pub fn iterator(self: *const ToolRegistry) Iterator {
        return .{ .inner = self.entries.iterator() };
    }

    pub const Iterator = struct {
        inner: std.StringHashMap(Entry).Iterator,

        pub fn next(self: *Iterator) ?ToolView {
            const entry = self.inner.next() orelse return null;
            return makeView(entry.value_ptr.*);
        }
    };

    fn makeView(entry: Entry) ToolView {
        return switch (entry) {
            .single => |t| .{ .decl = t.decl, .entry = entry },
            .source => |sr| .{ .decl = sr.source.tools[sr.tool_index], .entry = entry },
        };
    }
};

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

const testing = std.testing;

/// A trivial in-test Tool implementation backed by a single owned counter
/// allocation. Used to verify ownership/deinit behavior.
const TestTool = struct {
    invocations: u32 = 0,
    name_owned: []u8,
    desc_owned: []u8,
    schema_owned: []u8,

    fn create(allocator: Allocator, name: []const u8) !Tool {
        const self = try allocator.create(TestTool);
        errdefer allocator.destroy(self);

        const name_owned = try allocator.dupe(u8, name);
        errdefer allocator.free(name_owned);
        const desc_owned = try allocator.dupe(u8, "test tool");
        errdefer allocator.free(desc_owned);
        const schema_owned = try allocator.dupe(u8, "{}");
        errdefer allocator.free(schema_owned);

        self.* = .{
            .name_owned = name_owned,
            .desc_owned = desc_owned,
            .schema_owned = schema_owned,
        };
        return .{
            .decl = .{
                .name = self.name_owned,
                .description = self.desc_owned,
                .schema_json = self.schema_owned,
            },
            .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: *TestTool = @ptrCast(@alignCast(ctx));
        self.invocations += 1;
        return try allocator.dupe(u8, input);
    }

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

/// A minimal source backing N tools. Each tool name maps to a configured
/// response prefix; invoke_batch returns "<prefix>:<input>" for each
/// call. Tracks the batch sizes it was called with for inspection.
const TestSource = struct {
    name_owned: []u8,
    decls: []ToolDecl,
    /// Allocations backing every `decl`'s strings. Freed at deinit.
    allocations: std.array_list.Managed([]u8),
    batch_sizes: std.array_list.Managed(usize),
    allocator: Allocator,

    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 allocations = std.array_list.Managed([]u8).init(allocator);
        errdefer {
            for (allocations.items) |s| allocator.free(s);
            allocations.deinit();
        }

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

        const decls = try allocator.alloc(ToolDecl, tool_names.len);
        errdefer allocator.free(decls);

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

        self.* = .{
            .name_owned = name_owned,
            .decls = decls,
            .allocations = allocations,
            .batch_sizes = std.array_list.Managed(usize).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));
        try self.batch_sizes.append(calls.len);
        for (calls, 0..) |call, i| {
            const buf = std.fmt.allocPrint(
                allocator,
                "{s}:{s}",
                .{ call.tool_name, call.input },
            ) catch |e| {
                results[i] = .{ .err = e };
                continue;
            };
            results[i] = .{ .ok = buf };
        }
    }

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

test "register, lookup, count" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "echo"));
    try reg.register(try TestTool.create(allocator, "ls"));

    try testing.expectEqual(@as(usize, 2), reg.count());
    try testing.expect(reg.lookup("echo") != null);
    try testing.expect(reg.lookup("ls") != null);
    try testing.expect(reg.lookup("missing") == null);
    try testing.expectEqualStrings("echo", reg.lookup("echo").?.decl.name);
}

test "duplicate registration returns error and leaves original in place" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "echo"));

    // The second tool isn't taken over on duplicate; tear it down ourselves.
    var dup = try TestTool.create(allocator, "echo");
    try testing.expectError(error.DuplicateTool, reg.register(dup));
    dup.vtable.deinit(dup.ctx, allocator);

    try testing.expectEqual(@as(usize, 1), reg.count());
}

test "unregister calls deinit and removes" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "tmp"));
    try testing.expectEqual(@as(usize, 1), reg.count());

    reg.unregister("tmp");
    try testing.expectEqual(@as(usize, 0), reg.count());
    try testing.expect(reg.lookup("tmp") == null);

    // No-op on missing.
    reg.unregister("never_existed");
}

test "iterator visits every tool" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "a"));
    try reg.register(try TestTool.create(allocator, "b"));
    try reg.register(try TestTool.create(allocator, "c"));

    var saw_a = false;
    var saw_b = false;
    var saw_c = false;

    var it = reg.iterator();
    while (it.next()) |t| {
        if (std.mem.eql(u8, t.decl.name, "a")) saw_a = true;
        if (std.mem.eql(u8, t.decl.name, "b")) saw_b = true;
        if (std.mem.eql(u8, t.decl.name, "c")) saw_c = true;
    }
    try testing.expect(saw_a and saw_b and saw_c);
}

test "deinit frees all remaining tools" {
    // If this leaks, the testing allocator will catch it.
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    try reg.register(try TestTool.create(allocator, "x"));
    try reg.register(try TestTool.create(allocator, "y"));
    reg.deinit();
}

test "registerSource exposes every declared tool by name" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    const src = try TestSource.create(allocator, "panto-lua", &.{ "alpha", "beta", "gamma" });
    try reg.registerSource(src);

    try testing.expectEqual(@as(usize, 3), reg.count());
    const v = reg.lookup("beta") orelse return error.NotFound;
    try testing.expectEqualStrings("beta", v.decl.name);
    try testing.expect(v.entry == .source);
}

test "registerSource: collision with existing single tool aborts and rolls back" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.register(try TestTool.create(allocator, "shared"));

    // Build a source that includes the colliding name. We must tear it
    // down ourselves on failure.
    var src = try TestSource.create(allocator, "src", &.{ "first", "shared", "third" });
    try testing.expectError(error.DuplicateTool, reg.registerSource(src));
    src.vtable.deinit(src.ctx, allocator);

    // No partial state from the source remains.
    try testing.expectEqual(@as(usize, 1), reg.count());
    try testing.expect(reg.lookup("first") == null);
    try testing.expect(reg.lookup("third") == null);
}

test "registerSource: collision between two sources" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.registerSource(try TestSource.create(allocator, "a", &.{ "foo", "bar" }));

    var s = try TestSource.create(allocator, "b", &.{ "baz", "foo" });
    try testing.expectError(error.DuplicateTool, reg.registerSource(s));
    s.vtable.deinit(s.ctx, allocator);

    try testing.expectEqual(@as(usize, 2), reg.count());
}

test "source view exposes per-tool metadata uniformly" {
    const allocator = testing.allocator;
    var reg = ToolRegistry.init(allocator);
    defer reg.deinit();

    try reg.registerSource(try TestSource.create(allocator, "lua", &.{ "x", "y" }));
    try reg.register(try TestTool.create(allocator, "z"));

    try testing.expectEqual(@as(usize, 3), reg.count());

    // Every entry has the canonical fields populated.
    var it = reg.iterator();
    var n: usize = 0;
    while (it.next()) |v| : (n += 1) {
        try testing.expect(v.decl.name.len > 0);
        try testing.expect(v.decl.description.len > 0);
        try testing.expect(v.decl.schema_json.len > 0);
    }
    try testing.expectEqual(@as(usize, 3), n);
}