const std = @import("std");
const Atomic = std.atomic.Value;
const core = @import("core.zig");
const tracy = core.tracy;
const RingQueueU = core.RingQueueU;
const ArrayListUnmanaged = std.ArrayListUnmanaged;

const mutex_job_queue = core.BuildOption("mutex_job_queue");

pub const TaskInstanceOptions = struct {
    threadId: u32,
    threadCount: u32,
};

pub const JobManager = struct {
    allocator: std.mem.Allocator,

    // need a mutex for the jobQueue... todo later
    jobQueue: RingQueueU(ThreadContext),
    mutex: std.Thread.Mutex = .{},
    jobQueueConcurrent: core.ConcurrentQueueU(ThreadContext),
    workers: []*JobWorker,
    numCpus: usize,
    numWorkers: u32 = 0,

    pub fn create(allocator: std.mem.Allocator) !*@This() {
        var self = try allocator.create(@This());

        self.* = JobManager{
            .allocator = allocator,
            .numCpus = std.Thread.getCpuCount() catch 4,
            .jobQueue = RingQueueU(ThreadContext).init(allocator, 4096) catch unreachable,
            .jobQueueConcurrent = core.ConcurrentQueueU(ThreadContext).initCapacity(allocator, 4096) catch unreachable,
            .workers = undefined,
        };

        self.workers = self.allocator.alloc(*JobWorker, @max(self.numCpus - 2, 4)) catch unreachable;
        self.numWorkers = @intCast(self.workers.len);
        core.engine_log("allocating worker count : {d}", .{self.workers.len});

        var i: usize = 0;
        while (i < self.workers.len) : (i += 1) {
            self.workers[i] = JobWorker.init(self.allocator, i, true) catch unreachable;
            self.workers[i].workerThreadNumber = i;
            self.workers[i].manager = self;
        }

        return self;
    }

    pub fn runLocally(self: *@This(), capture: anytype, task: TaskInstanceOptions) !void {
        const Lambda = @TypeOf(capture);
        const ctx = try ThreadContext.new(self.allocator, Lambda, capture, task);

        //pub fn init(allocator: std.mem.Allocator, workerNumber: usize, detached:bool) !*@This() {
        const worker = try JobWorker.init(self.allocator, 0xff, false); // todo move these things into a pool of non-detached workers that we can just grab whenever
        defer worker.deinit();
        worker.currentJobContext = ctx;
        worker.run();
    }

    pub fn newJob(self: *@This(), capture: anytype, task: TaskInstanceOptions) !void {
        const Lambda = @TypeOf(capture);
        const ctx = try ThreadContext.new(self.allocator, Lambda, capture, task);

        if (mutex_job_queue) {
            self.mutex.lock();
            try self.jobQueue.push(ctx);
            self.mutex.unlock();
        } else {
            try self.jobQueueConcurrent.push(ctx);
        }

        self.bump();
    }

    pub fn bump(self: *@This()) void {
        var shouldBump: bool = false;

        if (mutex_job_queue) {
            shouldBump = self.jobQueue.count() > 0;
        } else {
            shouldBump = self.jobQueueConcurrent.count() > 0;
        }

        if (shouldBump) {
            for (self.workers) |worker| {
                if (!worker.isBusy()) {
                    worker.wake();
                    break;
                }
            }
        }
    }

    pub fn destroy(self: *@This()) void {
        for (self.workers) |worker| {
            worker.deinit();
        }
        self.allocator.free(self.workers);

        self.clearJobs();

        self.jobQueue.deinit(self.allocator);
        self.jobQueueConcurrent.deinit(self.allocator);
        self.allocator.destroy(self);
    }

    pub fn clearJobs(self: *@This()) void {
        var jobCtx: ?ThreadContext = null;

        if (mutex_job_queue) {
            self.mutex.lock();
            jobCtx = self.jobQueue.pop();
            self.mutex.unlock();
        } else {
            jobCtx = self.jobQueueConcurrent.pop();
        }

        while (jobCtx) |*c| {
            c.deinit();

            if (mutex_job_queue) {
                self.mutex.lock();
                jobCtx = self.jobQueue.pop();
                self.mutex.unlock();
            } else {
                jobCtx = self.jobQueueConcurrent.pop();
            }
        }
    }
};

pub const JobWorker = struct {
    detached: bool = true, // most threads are detached, completions are handled via callbacks.
    currentJobContext: ?ThreadContext = null,
    workerThread: ?std.Thread = null,
    futex: Atomic(u32) = Atomic(u32).init(0),
    current: u32 = 0,
    shouldDie: Atomic(bool) = Atomic(bool).init(false),
    busy: Atomic(bool) = Atomic(bool).init(false),
    allocator: std.mem.Allocator,
    workerThreadNumber: usize = 0, // identifier for debugging purposes, never changes.
    manager: ?*JobManager = null,
    scratchArena: std.heap.ArenaAllocator,

    pub fn wake(self: *JobWorker) void {
        // this needs to be re-evaluated, it's nice for system performance but
        // a 1-2ms worst case wake time is kind of unacceptable.
        std.Thread.Futex.wake(&self.futex, 1);
    }

    pub fn init(allocator: std.mem.Allocator, workerNumber: usize, detached: bool) !*@This() {
        const self = try allocator.create(JobWorker);

        self.* = .{
            .allocator = allocator,
            .detached = detached,
            .workerThreadNumber = workerNumber,
            .scratchArena = std.heap.ArenaAllocator.init(allocator),
        };

        if (detached) {
            self.workerThread = try std.Thread.spawn(.{}, @This().workerThreadFunc, .{self});
        }

        return self;
    }

    pub fn isBusy(self: *@This()) bool {
        return self.busy.load(.acquire);
    }

    fn run(self: *@This()) void {
        self.busy.store(true, .seq_cst);

        if(self.manager)|manager|
        {
            manager.bump();
        }

        var ctx = self.currentJobContext.?;
        ctx.workerInfo = self;
        ctx.func(ctx.capture, &ctx);
        self.busy.store(false, .seq_cst);

        ctx.deinit();
        if (self.scratchArena.reset(.retain_capacity)) {}
        self.currentJobContext = null;
    }

    pub fn workerThreadFunc(self: *@This()) void {
        const printed = std.fmt.allocPrintSentinel(self.allocator, "WorkerThread_{d}", .{self.workerThreadNumber}, 0) catch unreachable;
        tracy.InitThread();
        tracy.SetThreadName(@as([*:0]u8, @ptrCast(printed.ptr)));

        self.allocator.free(printed);

        var wakeGrabCount: u32 = 1;

        while (!self.shouldDie.load(.acquire)) {
            if (self.currentJobContext != null) {
                wakeGrabCount = 0;
                self.run();
            } else {
                std.Thread.Futex.wait(&self.futex, self.current);
            }

            if (self.manager) |manager| {
                if (mutex_job_queue) {
                    self.manager.?.mutex.lock();
                    if (self.manager.?.jobQueue.count() > 0) {
                        self.currentJobContext = self.manager.?.jobQueue.pop().?;
                    }
                    self.manager.?.mutex.unlock();
                } else {
                    self.currentJobContext = manager.jobQueueConcurrent.pop();
                    wakeGrabCount = 1;
                }
            }
        }

        if (self.currentJobContext) |*ctx| {
            ctx.deinit();
        }
    }

    pub fn deinit(self: *@This()) void {
        self.shouldDie.store(true, .seq_cst);
        self.wake();
        if(self.workerThread)|workerThread|
        {
            workerThread.join();
        }

        self.scratchArena.deinit();
        self.allocator.destroy(self);
    }
};

const Src = std.builtin.SourceLocation;
pub fn BarrierCV(src: Src) type {
    return struct {
        pub const Src = src;

        pub var count: u32 = 0;
        pub var generation: u32 = 0;
        pub var mutex: std.Thread.Mutex = .{};
        pub var cv: std.Thread.Condition = .{};

        pub fn sync(threadId: u32, barrierCount: u32) !void {
            _ = threadId;

            mutex.lock();
            defer mutex.unlock();

            count += 1;

            if (count == barrierCount) {
                // 2 second timeout
                count = 0;
                generation += 1;
                cv.broadcast();
            } else {
                try cv.timedWait(&mutex, 1000 * 1000 * 1000 * 2);
            }
        }
    };
}

pub fn BarrierSpin(src: Src) type {
    return struct {
        pub const Src = src;

        pub var count: std.atomic.Value(u32) = std.atomic.Value(u32).init(0);
        pub var park: std.atomic.Value(bool) = std.atomic.Value(bool).init(false);

        pub fn sync(threadId: u32, barrierCount: u32) !void {
            _ = threadId;

            const c = count.fetchAdd(1, .release);
            if(c + 1 == barrierCount)
            {
                park.store(true, .release);
                while(count.load(.acquire) > 1) {}
                _ = count.fetchSub(1, .release);
                park.store(false, .release);
            }
            else 
            {
                while(park.load(.acquire) == false) {}
                _ = count.fetchSub(1, .seq_cst);
            }
        }
    };
}

pub const Barrier = BarrierSpin;

pub const ThreadContext = struct {
    const Self = @This();

    funcName: []const u8 = "unnamed",
    allocator: std.mem.Allocator, //todo, backed arena allocator would be sick for this.
    func: *const fn (*anyopaque, *ThreadContext) void, // todo, add an error for job funcs
    capture: *anyopaque = undefined,
    threadId: u32,
    threadIdCount: u32,
    destroyFunc: *const fn (*anyopaque, std.mem.Allocator) void,
    workerInfo: ?*JobWorker = null,

    pub fn new(
        allocator: std.mem.Allocator,
        comptime CaptureType: type,
        capture: CaptureType,
        task: TaskInstanceOptions,
    ) !ThreadContext {
        if (!@hasDecl(CaptureType, "func")) {
            return error.NoValidLambda;
        }

        const Wrap = struct {
            pub fn wrappedFunc(pointer: *anyopaque, context: *ThreadContext) void {
                var ptr = @as(*CaptureType, @ptrCast(@alignCast(pointer)));
                ptr.func(context);
            }

            pub fn wrappedDestroy(ptr: *anyopaque, alloc: std.mem.Allocator) void {
                const p = @as(*CaptureType, @ptrCast(@alignCast(ptr)));
                alloc.destroy(p);
            }
        };

        const self = Self{
            .allocator = allocator,
            .threadId = task.threadId,
            .threadIdCount = task.threadCount,
            .funcName = @typeName(CaptureType),
            .func = Wrap.wrappedFunc,
            .destroyFunc = Wrap.wrappedDestroy,
            .capture = try allocator.create(CaptureType),
        };
        const ptr = @as(*CaptureType, @ptrCast(@alignCast(self.capture)));
        ptr.* = capture;
        return self;
    }

    // helper functions

    pub fn scratch(self: *@This()) std.mem.Allocator {
        return self.workerInfo.?.scratchArena.allocator();
    }

    pub fn splitSlice(self: @This(), comptime T: type, slice: []T) struct{ slice: []T, startIndex: usize } {
        const r =  core.p2.splitSliceWork(T, slice, self.threadId, self.threadIdCount);
        return .{.slice = r.slice, .startIndex = r.startIndex};

    }

    pub fn barrier(self: *@This(), comptime src: std.builtin.SourceLocation) !void {
        try Barrier(src).sync(self.threadId, self.threadIdCount);
    }

    pub fn deinit(self: *Self) void {
        self.destroyFunc(self.capture, self.allocator);
    }
};