const Runner = @This();

const session = @import("../session.zig");

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

const concurrent_task = 2;
queue: std.Io.Queue(TaskHandle),
queue_storage: [](TaskHandle),
io: std.Io,
task_group: std.Io.Group,
gpa: std.mem.Allocator,

task_id: std.atomic.Value(TaskHandle) = std.atomic.Value(TaskHandle).init(0),
tasks: TaskStore,

userdata: *anyopaque,
on_result: OnResultCallback,

const OnResultCallback = *const fn (*anyopaque, Task) void;
const TaskMap = std.AutoHashMapUnmanaged(TaskHandle, Task);

const TaskStore = struct {
    map: std.AutoHashMapUnmanaged(TaskHandle, Task) = .empty,
    lock: std.Io.RwLock = .init,
    io: std.Io,
    gpa: std.mem.Allocator,
    
    pub fn init(io: std.Io, gpa: std.mem.Allocator) TaskStore {
        return .{
            .io = io,
            .gpa = gpa,
        };
    }
    
    pub fn deinit(self: *TaskStore) void {
        self.map.deinit(self.gpa);
    }
    
    pub fn add(self: *TaskStore, task: Task) !void {
        self.lock.lockUncancelable(self.io);
        defer self.lock.unlock(self.io);
        try self.map.put(self.gpa, task.id, task);
    }
    
    pub fn get(self: *TaskStore, id: TaskHandle) *Task {
        return self.map.getPtr(id).?;
    }
    
    pub fn remove(self: *TaskStore, id: TaskHandle) void {
        self.lock.lockUncancelable(self.io);
        defer self.lock.unlock(self.io);
        _ = self.map.remove(id);
    }
};

const TaskHandle = u64;
fn nextTaskId(self: *Runner) u64 {
    return self.task_id.fetchAdd(1, .monotonic);
}

fn initTask(self: *Runner, tool_use: session.ToolUseStart) !TaskHandle {
    const task = Task.init(self.io, tool_use, self.nextTaskId());
    try self.tasks.add(task);
    return task.id;
}

const log = std.log.scoped(.tool_runner);
pub fn init(io: std.Io, gpa: std.mem.Allocator, userdata: *anyopaque, on_result: OnResultCallback) !Runner {
    const storage = try gpa.alloc(TaskHandle, concurrent_task * 10);
    return Runner{
        .io = io,
        .queue = .init(storage),
        .queue_storage = storage,
        .task_group = .init,
        .tasks = .init(io, gpa),
        .gpa = gpa,
        .userdata = userdata,
        .on_result = on_result,
    };
}

pub fn start(self: *Runner) !void {
    for (0..concurrent_task) |i| {
        try self.task_group.concurrent(self.io, consumer, .{self, i});
    }
}

pub fn stop(self: *Runner) void {
    self.task_group.cancel(self.io);
}

pub fn deinit(self: *Runner) void {
    self.queue.close(self.io);
    self.stop();
    self.gpa.free(self.queue_storage);
    self.tasks.deinit();
}

pub fn execute(self: *Runner, tool_call: session.ToolUseStart) !TaskHandle {
    const handle = try self.initTask(tool_call);
    log.info("Put task {d}", .{handle});
    try self.queue.putOne(self.io, handle);
    return handle;
}

pub fn interrupt(self: *Runner) void {
    self.task_group.cancel(self.io);
}

fn runTool(self: *Runner, handle: TaskHandle) void {
    var task = self.tasks.get(handle);

    const tc = task.tool_use;
    log.info("execute tool call: {s} {s}", .{ @tagName(tc.tool), tc.id });

    task.tool_use_result = .{
        .id = tc.id,
        .content = "",
        .timestamp = .{ .nanoseconds = 3 },
        .tool = tc.tool,
        .is_error = null,
    };
}

fn consumer(self: *Runner, num: usize) void {
    log.info("consumer started: {d}", .{num});
    while (self.queue.getOne(self.io)) |task_handle| {
        log.info("consumer {d} got task: {d}", .{num ,task_handle});
        var task = self.tasks.get(task_handle);
        if (task.canceled.load(.monotonic)) {
            self.tasks.remove(task.id);
            log.info("Consumer {d} skips canceled task {d}", .{num ,task.id});
            continue;
        }
        task.group.async(self.io, runTool, .{ self, task_handle });
        task.group.await(self.io) catch |err| {
            var t = self.tasks.get(task_handle);
            t.err = err;
            log.info("Consumer {d} got error {s} for task {d}", .{num ,@errorName(err), task.id});
            self.on_result(self.userdata, t.*);
            continue;
        };
        const t = self.tasks.get(task_handle);
        log.info("Consumer {d} got result for task {d}", .{num ,t.id});
        self.on_result(self.userdata, t.*);
    } else |_| {
        log.info("Consumer {d} stopped", .{num});
    }
}

const Task = struct {
    tool_use: session.ToolUseStart,
    tool_use_result: ?session.ToolUseCompleted = null,
    group: Io.Group,
    io: Io,
    canceled: std.atomic.Value(bool) = .init(false),
    id: TaskHandle,
    err: ?anyerror = null,

    fn init(io: Io, tool_use: session.ToolUseStart, id: TaskHandle) Task {
        return .{ .io = io, .tool_use = tool_use, .group = .init, .id = id };
    }

    fn deinit(self: *Task) void {
        self.cancel();
    }

    fn await(self: *Task) !session.ToolUseCompleted {
        try self.group.await(self.io);
        return self.tool_use_result orelse error.Canceled;
    }

    fn cancel(self: *Task) void {
        self.canceled.store(true, .monotonic);
        self.group.cancel(self.io);
    }

    pub fn jsonStringify(self: *const Task, jws: anytype) std.json.Stringify.Error!void {
        try jws.write(.{
            .id = self.id,
            .err = self.err,
            .tool_use = self.tool_use,
            .tool_use_result = self.tool_use_result,
        });
    }
};