const std = @import("std");
const builtin = @import("builtin");

pub const max_output_bytes: usize = 1024 * 1024;

pub const Status = enum { success, @"error" };

pub const PendingToolCall = struct {
    name: []const u8,
    arguments_json: []const u8,
};

pub const ExecuteOptions = struct {
    io: std.Io,
    cwd: []const u8,
    allow_bash: bool = false,
    environ_map: ?*const std.process.Environ.Map = null,
    agent_id: ?[]const u8 = null,
    conversation_id: ?[]const u8 = null,
    device_id: ?[]const u8 = null,
    api_base_url: ?[]const u8 = null,
    api_key: ?[]const u8 = null,
    /// External abort signal (e.g. `abort_message`). When set, a running Bash
    /// command returns a cancelled result promptly.
    cancel_token: ?*const std.atomic.Value(bool) = null,
};

const cancel_poll_interval_ms: i64 = 20;

fn cancelledByToken(token: ?*const std.atomic.Value(bool)) bool {
    return if (token) |t| t.load(.acquire) else false;
}

pub const ToolResult = struct {
    status: Status,
    tool_return: []u8,
    stdout: ?[]u8 = null,
    stderr: ?[]u8 = null,

    pub fn deinit(self: ToolResult, allocator: std.mem.Allocator) void {
        allocator.free(self.tool_return);
        if (self.stdout) |out| allocator.free(out);
        if (self.stderr) |err| allocator.free(err);
    }
};

const Args = struct {
    root: std.json.Parsed(std.json.Value),

    fn deinit(self: Args) void {
        self.root.deinit();
    }
};

pub fn executeTool(allocator: std.mem.Allocator, call: PendingToolCall, options: ExecuteOptions) !ToolResult {
    var args = try parseArgs(allocator, call.arguments_json);
    defer args.deinit();

    if (std.ascii.eqlIgnoreCase(call.name, "Read")) return readTool(allocator, args.root.value, options);
    if (std.ascii.eqlIgnoreCase(call.name, "List") or std.ascii.eqlIgnoreCase(call.name, "LS")) return listTool(allocator, args.root.value, options);
    if (std.ascii.eqlIgnoreCase(call.name, "Glob")) return globTool(allocator, args.root.value, options);
    if (std.ascii.eqlIgnoreCase(call.name, "Grep")) return grepTool(allocator, args.root.value, options);
    if (std.ascii.eqlIgnoreCase(call.name, "Bash")) return bashTool(allocator, args.root.value, options);

    return errorResult(allocator, "unsupported tool");
}

fn parseArgs(allocator: std.mem.Allocator, json_text: []const u8) !Args {
    return .{ .root = try std.json.parseFromSlice(std.json.Value, allocator, json_text, .{}) };
}

fn readTool(allocator: std.mem.Allocator, root: std.json.Value, options: ExecuteOptions) !ToolResult {
    const file_path = getString(root, "file_path") orelse return errorResult(allocator, "Read requires file_path");
    const offset = getOptionalUsize(root, "offset") orelse 0;
    const limit = getOptionalUsize(root, "limit") orelse max_output_bytes;
    const resolved_path = try resolvePath(allocator, options.cwd, file_path);
    defer allocator.free(resolved_path);
    var file = std.Io.Dir.cwd().openFile(options.io, resolved_path, .{}) catch return errorResult(allocator, "failed to open file");
    defer file.close(options.io);

    const bounded_limit = @min(limit, max_output_bytes);
    const out = try allocator.alloc(u8, bounded_limit);
    errdefer allocator.free(out);
    const read_len = try file.readPositionalAll(options.io, out, offset);
    return .{ .status = .success, .tool_return = try allocator.realloc(out, read_len) };
}

fn listTool(allocator: std.mem.Allocator, root: std.json.Value, options: ExecuteOptions) !ToolResult {
    const path = getString(root, "path") orelse ".";
    const resolved_path = try resolvePath(allocator, options.cwd, path);
    defer allocator.free(resolved_path);
    var dir = std.Io.Dir.openDirAbsolute(options.io, resolved_path, .{ .iterate = true }) catch return errorResult(allocator, "failed to open directory");
    defer dir.close(options.io);

    var buf = std.ArrayList(u8).empty;
    var it = dir.iterate();
    while (try it.next(options.io)) |entry| {
        try appendBounded(allocator, &buf, entry.name);
        try appendBounded(allocator, &buf, if (entry.kind == .directory) "/\n" else "\n");
        if (buf.items.len >= max_output_bytes) break;
    }
    return .{ .status = .success, .tool_return = try buf.toOwnedSlice(allocator) };
}

fn globTool(allocator: std.mem.Allocator, root: std.json.Value, options: ExecuteOptions) !ToolResult {
    const pattern = getString(root, "pattern") orelse return errorResult(allocator, "Glob requires pattern");
    const path = getString(root, "path") orelse ".";
    const resolved_path = try resolvePath(allocator, options.cwd, path);
    defer allocator.free(resolved_path);
    var start = std.Io.Dir.openDirAbsolute(options.io, resolved_path, .{ .iterate = true }) catch return errorResult(allocator, "failed to open glob path");
    defer start.close(options.io);

    var buf = std.ArrayList(u8).empty;
    try walkMatch(allocator, options.io, start, path, pattern, &buf, false, "");
    return .{ .status = .success, .tool_return = try buf.toOwnedSlice(allocator) };
}

fn grepTool(allocator: std.mem.Allocator, root: std.json.Value, options: ExecuteOptions) !ToolResult {
    const pattern = getString(root, "pattern") orelse return errorResult(allocator, "Grep requires pattern");
    const path = getString(root, "path") orelse ".";
    const glob = getString(root, "glob") orelse "*";
    const resolved_path = try resolvePath(allocator, options.cwd, path);
    defer allocator.free(resolved_path);
    var start = std.Io.Dir.openDirAbsolute(options.io, resolved_path, .{ .iterate = true }) catch return errorResult(allocator, "failed to open grep path");
    defer start.close(options.io);

    var buf = std.ArrayList(u8).empty;
    try appendBounded(allocator, &buf, "literal substring Grep PoC; regex flags are not supported\n");
    try walkMatch(allocator, options.io, start, path, glob, &buf, true, pattern);
    return .{ .status = .success, .tool_return = try buf.toOwnedSlice(allocator) };
}

fn bashTool(allocator: std.mem.Allocator, root: std.json.Value, options: ExecuteOptions) !ToolResult {
    if (!options.allow_bash) return errorResult(allocator, "Bash denied: allow_bash must be true");
    const command = getString(root, "command") orelse return errorResult(allocator, "Bash requires command");
    _ = getString(root, "description");
    const timeout_ms = getOptionalUsize(root, "timeout");

    const argv = if (builtin.os.tag == .windows)
        &[_][]const u8{ "cmd.exe", "/C", command }
    else
        &[_][]const u8{ "/bin/sh", "-lc", command };

    var environ = std.process.Environ.Map.init(allocator);
    defer environ.deinit();
    if (options.environ_map) |parent| {
        for (parent.keys(), parent.values()) |key, value| try environ.put(key, value);
    }
    try environ.put("USER_CWD", options.cwd);
    if (options.agent_id) |agent_id| {
        try environ.put("LETTA_AGENT_ID", agent_id);
        try environ.put("AGENT_ID", agent_id);
    }
    if (options.conversation_id) |conversation_id| {
        try environ.put("LETTA_CONVERSATION_ID", conversation_id);
        try environ.put("CONVERSATION_ID", conversation_id);
    }
    if (options.device_id) |device_id| try environ.put("LETTA_RUNTIME_ENVIRONMENT_DEVICE_ID", device_id);
    if (options.api_base_url) |api_base_url| try environ.put("LETTA_BASE_URL", api_base_url);
    if (options.api_key) |api_key| try environ.put("LETTA_API_KEY", api_key);

    const result = runProcessWithCancel(allocator, options.io, .{
        .argv = argv,
        .cwd = .{ .path = options.cwd },
        .environ_map = &environ,
        .stdout_limit = .limited(max_output_bytes),
        .stderr_limit = .limited(max_output_bytes),
        .timeout = if (timeout_ms) |ms| .{ .duration = .{ .raw = std.Io.Duration.fromMilliseconds(@intCast(ms)), .clock = .awake } } else .none,
    }, options.cancel_token) catch |err| switch (err) {
        error.Cancelled => return cancelledBashResult(allocator),
        else => return err,
    };
    errdefer allocator.free(result.stdout);
    errdefer allocator.free(result.stderr);

    const ok = switch (result.term) {
        .exited => |code| code == 0,
        else => false,
    };
    if (cancelledByToken(options.cancel_token)) return cancelledBashResultWithOutput(allocator, result.stdout, result.stderr, result.term);
    const ret = try std.fmt.allocPrint(allocator, "exit: {}\nstdout:\n{s}\nstderr:\n{s}", .{ result.term, result.stdout, result.stderr });
    return .{
        .status = if (ok) .success else .@"error",
        .tool_return = ret,
        .stdout = result.stdout,
        .stderr = result.stderr,
    };
}

fn runProcessWithCancel(allocator: std.mem.Allocator, io: std.Io, options: std.process.RunOptions, cancel_token: ?*const std.atomic.Value(bool)) !std.process.RunResult {
    if (cancelledByToken(cancel_token)) return error.Cancelled;
    const Selection = union(enum) {
        process: anyerror!std.process.RunResult,
        cancelled: anyerror!void,
    };
    const Tasks = struct {
        fn run(alloc: std.mem.Allocator, target_io: std.Io, run_options: std.process.RunOptions) anyerror!std.process.RunResult {
            return std.process.run(alloc, target_io, run_options);
        }

        fn waitForCancel(target_io: std.Io, token: *const std.atomic.Value(bool)) anyerror!void {
            while (!token.load(.acquire)) {
                try std.Io.sleep(target_io, .fromMilliseconds(cancel_poll_interval_ms), .awake);
            }
            return error.Cancelled;
        }
    };
    var results: [2]Selection = undefined;
    var select = std.Io.Select(Selection).init(io, &results);
    select.async(.process, Tasks.run, .{ allocator, io, options });
    if (cancel_token) |token| select.async(.cancelled, Tasks.waitForCancel, .{ io, token });
    defer select.cancelDiscard();
    const first = try select.await();
    return switch (first) {
        .process => |result| try result,
        .cancelled => |result| {
            try result;
            return error.Cancelled;
        },
    };
}

fn cancelledBashResult(allocator: std.mem.Allocator) !ToolResult {
    return .{ .status = .@"error", .tool_return = try allocator.dupe(u8, "cancelled") };
}

fn cancelledBashResultWithOutput(allocator: std.mem.Allocator, stdout: []u8, stderr: []u8, term: std.process.Child.Term) !ToolResult {
    errdefer allocator.free(stdout);
    errdefer allocator.free(stderr);
    const ret = try std.fmt.allocPrint(allocator, "cancelled\nexit: {}\nstdout:\n{s}\nstderr:\n{s}", .{ term, stdout, stderr });
    return .{ .status = .@"error", .tool_return = ret, .stdout = stdout, .stderr = stderr };
}

fn walkMatch(
    allocator: std.mem.Allocator,
    io: std.Io,
    dir: std.Io.Dir,
    display_root: []const u8,
    pattern: []const u8,
    buf: *std.ArrayList(u8),
    grep: bool,
    needle: []const u8,
) !void {
    var it_dir = dir;
    var it = it_dir.iterate();
    while (try it.next(io)) |entry| {
        if (buf.items.len >= max_output_bytes) return;
        const rel = if (std.mem.eql(u8, display_root, "."))
            try allocator.dupe(u8, entry.name)
        else
            try std.fs.path.join(allocator, &.{ display_root, entry.name });
        defer allocator.free(rel);

        if (entry.kind == .directory) {
            var child = it_dir.openDir(io, entry.name, .{ .iterate = true }) catch continue;
            defer child.close(io);
            try walkMatch(allocator, io, child, rel, pattern, buf, grep, needle);
        } else if (wildcardMatch(pattern, rel) or wildcardMatch(pattern, entry.name)) {
            if (grep) {
                var file = it_dir.openFile(io, entry.name, .{}) catch continue;
                defer file.close(io);
                var scratch = try allocator.alloc(u8, max_output_bytes);
                defer allocator.free(scratch);
                const read_len = file.readPositionalAll(io, scratch, 0) catch continue;
                if (std.mem.indexOf(u8, scratch[0..read_len], needle) != null) {
                    try appendBounded(allocator, buf, rel);
                    try appendBounded(allocator, buf, "\n");
                }
            } else {
                try appendBounded(allocator, buf, rel);
                try appendBounded(allocator, buf, "\n");
            }
        }
    }
}

fn wildcardMatch(pattern: []const u8, text: []const u8) bool {
    var p: usize = 0;
    var t: usize = 0;
    var star: ?usize = null;
    var match_i: usize = 0;
    while (t < text.len) {
        if (p < pattern.len and (pattern[p] == '?' or pattern[p] == text[t])) {
            p += 1;
            t += 1;
        } else if (p < pattern.len and pattern[p] == '*') {
            star = p;
            match_i = t;
            p += 1;
        } else if (star) |s| {
            p = s + 1;
            match_i += 1;
            t = match_i;
        } else {
            return false;
        }
    }
    while (p < pattern.len and pattern[p] == '*') p += 1;
    return p == pattern.len;
}

fn appendBounded(allocator: std.mem.Allocator, buf: *std.ArrayList(u8), text: []const u8) !void {
    if (buf.items.len >= max_output_bytes) return;
    const available = max_output_bytes - buf.items.len;
    const take = @min(available, text.len);
    try buf.appendSlice(allocator, text[0..take]);
}

fn truncateWithNotice(allocator: std.mem.Allocator, owned: []u8, limit: usize) ![]u8 {
    const notice = "\n[truncated at 1 MiB]\n";
    const keep = @min(limit, owned.len);
    var out = try allocator.alloc(u8, keep + notice.len);
    @memcpy(out[0..keep], owned[0..keep]);
    @memcpy(out[keep..], notice);
    allocator.free(owned);
    return out;
}

fn errorResult(allocator: std.mem.Allocator, msg: []const u8) !ToolResult {
    return .{ .status = .@"error", .tool_return = try allocator.dupe(u8, msg) };
}

fn getString(root: std.json.Value, key: []const u8) ?[]const u8 {
    if (root != .object) return null;
    const value = root.object.get(key) orelse return null;
    return switch (value) {
        .string => |s| s,
        else => null,
    };
}

fn getOptionalUsize(root: std.json.Value, key: []const u8) ?usize {
    if (root != .object) return null;
    const value = root.object.get(key) orelse return null;
    return switch (value) {
        .integer => |n| if (n >= 0) @intCast(n) else null,
        .float => |f| if (f >= 0) @intFromFloat(f) else null,
        else => null,
    };
}

fn resolvePath(allocator: std.mem.Allocator, cwd: []const u8, path: []const u8) ![]u8 {
    if (std.fs.path.isAbsolute(path)) return allocator.dupe(u8, path);
    return std.fs.path.resolve(allocator, &.{ cwd, if (path.len == 0) "." else path });
}

fn tmpPathAlloc(allocator: std.mem.Allocator, io: std.Io, dir: std.Io.Dir) ![]u8 {
    var buf: [std.fs.max_path_bytes]u8 = undefined;
    const len = try dir.realPath(io, &buf);
    return allocator.dupe(u8, buf[0..len]);
}

test "Read reads bounded file under configured cwd" {
    const allocator = std.testing.allocator;
    const io = std.testing.io;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    try tmp.dir.writeFile(io, .{ .sub_path = "note.txt", .data = "hello tool executor" });

    const cwd_path = try tmpPathAlloc(allocator, io, tmp.dir);
    defer allocator.free(cwd_path);
    var result = try executeTool(allocator, .{ .name = "Read", .arguments_json = "{\"file_path\":\"note.txt\",\"offset\":6,\"limit\":4}" }, .{ .io = io, .cwd = cwd_path });
    defer result.deinit(allocator);

    try std.testing.expectEqual(Status.success, result.status);
    try std.testing.expectEqualStrings("tool", result.tool_return);
}

test "List lists directory under configured cwd" {
    const allocator = std.testing.allocator;
    const io = std.testing.io;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    try tmp.dir.writeFile(io, .{ .sub_path = "a.txt", .data = "a" });
    try tmp.dir.createDir(io, "sub", .default_dir);

    const cwd_path = try tmpPathAlloc(allocator, io, tmp.dir);
    defer allocator.free(cwd_path);
    var result = try executeTool(allocator, .{ .name = "List", .arguments_json = "{\"path\":\".\"}" }, .{ .io = io, .cwd = cwd_path });
    defer result.deinit(allocator);

    try std.testing.expectEqual(Status.success, result.status);
    try std.testing.expect(std.mem.indexOf(u8, result.tool_return, "a.txt\n") != null);
    try std.testing.expect(std.mem.indexOf(u8, result.tool_return, "sub/\n") != null);
}

test "Bash injects the current Portal runtime into letta teleport" {
    if (builtin.os.tag == .windows) return error.SkipZigTest;
    const allocator = std.testing.allocator;
    const io = std.testing.io;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    const cwd_path = try tmpPathAlloc(allocator, io, tmp.dir);
    defer allocator.free(cwd_path);

    var parent_env = std.process.Environ.Map.init(allocator);
    defer parent_env.deinit();
    try parent_env.put("PATH", "/usr/bin:/bin");
    var result = try executeTool(allocator, .{
        .name = "Bash",
        .arguments_json = "{\"command\":\"printf '%s|%s|%s|%s|%s|%s' \\\"$LETTA_AGENT_ID\\\" \\\"$LETTA_CONVERSATION_ID\\\" \\\"$LETTA_RUNTIME_ENVIRONMENT_DEVICE_ID\\\" \\\"$LETTA_BASE_URL\\\" \\\"$USER_CWD\\\" \\\"$LETTA_API_KEY\\\"\",\"description\":\"inspect teleport context\"}",
    }, .{
        .io = io,
        .cwd = cwd_path,
        .allow_bash = true,
        .environ_map = &parent_env,
        .agent_id = "agent-runtime",
        .conversation_id = "conversation-runtime",
        .device_id = "portal-device",
        .api_base_url = "https://api.example.test",
        .api_key = "secret-test-key",
    });
    defer result.deinit(allocator);

    try std.testing.expectEqual(Status.success, result.status);
    const expected = try std.fmt.allocPrint(allocator, "agent-runtime|conversation-runtime|portal-device|https://api.example.test|{s}|secret-test-key", .{cwd_path});
    defer allocator.free(expected);
    try std.testing.expectEqualStrings(expected, result.stdout.?);
}

test "Bash cancellation returns cancelled result" {
    if (builtin.os.tag == .windows) return error.SkipZigTest;
    const allocator = std.testing.allocator;
    const io = std.testing.io;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    const cwd_path = try tmpPathAlloc(allocator, io, tmp.dir);
    defer allocator.free(cwd_path);
    var token: std.atomic.Value(bool) = .init(false);
    const Tasks = struct {
        fn cancelLater(target_io: std.Io, t: *std.atomic.Value(bool)) !void {
            try std.Io.sleep(target_io, .fromMilliseconds(20), .awake);
            t.store(true, .release);
        }
    };
    var group: std.Io.Group = .init;
    group.async(io, Tasks.cancelLater, .{ io, &token });
    defer group.cancel(io);
    var result = try executeTool(allocator, .{
        .name = "Bash",
        .arguments_json = "{\"command\":\"sleep 30\",\"description\":\"stall for cancellation\",\"timeout\":60000}",
    }, .{
        .io = io,
        .cwd = cwd_path,
        .allow_bash = true,
        .cancel_token = &token,
    });
    defer result.deinit(allocator);
    try std.testing.expectEqual(Status.@"error", result.status);
    try std.testing.expect(std.mem.indexOf(u8, result.tool_return, "cancelled") != null);
}

test "Bash with uncancelled token preserves success behavior" {
    if (builtin.os.tag == .windows) return error.SkipZigTest;
    const allocator = std.testing.allocator;
    const io = std.testing.io;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    const cwd_path = try tmpPathAlloc(allocator, io, tmp.dir);
    defer allocator.free(cwd_path);
    var token: std.atomic.Value(bool) = .init(false);
    var result = try executeTool(allocator, .{
        .name = "Bash",
        .arguments_json = "{\"command\":\"printf ok\",\"description\":\"success with token\"}",
    }, .{
        .io = io,
        .cwd = cwd_path,
        .allow_bash = true,
        .cancel_token = &token,
    });
    defer result.deinit(allocator);
    try std.testing.expectEqual(Status.success, result.status);
    try std.testing.expectEqualStrings("ok", result.stdout.?);
}

test "Bash denied unless allow_bash is true" {
    const allocator = std.testing.allocator;
    const io = std.testing.io;
    var tmp = std.testing.tmpDir(.{});
    defer tmp.cleanup();
    const cwd_path = try tmpPathAlloc(allocator, io, tmp.dir);
    defer allocator.free(cwd_path);

    var result = try executeTool(allocator, .{ .name = "Bash", .arguments_json = "{\"command\":\"echo nope\",\"description\":\"denied test\"}" }, .{ .io = io, .cwd = cwd_path, .allow_bash = false });
    defer result.deinit(allocator);

    try std.testing.expectEqual(Status.@"error", result.status);
    try std.testing.expect(std.mem.indexOf(u8, result.tool_return, "allow_bash") != null);
}