const std = @import("std");
const utils = @import("utils.zig");

const ArrayList = std.ArrayList;
const ArrayListUnmanaged = std.ArrayListUnmanaged;

const asserts = utils.asserts;

const Atomic = std.atomic.Value;

pub const ConcurrentQueueError = error{
    CorruptedState,
    QueueIsEmpty,
    QueueIsFull,
};

pub fn ConcurrentQueueU(comptime T: type) type {
    return ConcurrentQueueUnmanagedAdvanced(T, .{});
}

pub fn ConcurrentQueueAssert(comptime T: type) type {
    return ConcurrentQueueUnmanagedAdvanced(T, .{ .allowAsserts = true });
}

// lock-free concurrent queue, fixed capacity,
// will never resize.
//
// new version based on
// https://github.com/rigtorp/MPMCQueue

pub const ConcurrentStatus = packed struct(usize) {
    alive: bool = false,
    generation: u63 = 0,
};

pub fn ConcurrentQueueUnmanagedAdvanced(comptime T: type, comptime opts: struct {
    allowAsserts: bool = false,
    debug: bool = false,
}) type {
    return struct {
        data: []align(64) T align(64),
        status: []align(64) Atomic(ConcurrentStatus) align(64), // 1 = valid, 0 = invalid,
        pushId: Atomic(usize) align(64),
        popId: Atomic(usize) align(64),

        // tail points to next free slot
        // head points to the next one to pop
        pub fn initCapacity(allocator: std.mem.Allocator, cap: usize) !@This() {
            const new = @This(){
                .data = try allocator.alignedAlloc(T, .@"64", cap),
                .status = try allocator.alignedAlloc(Atomic(ConcurrentStatus), .@"64", cap),
                .pushId = Atomic(usize).init(0),
                .popId = Atomic(usize).init(0),
            };

            for (new.status) |*s| {
                s.* = Atomic(ConcurrentStatus).init(std.mem.zeroes(ConcurrentStatus));
            }

            return new;
        }

        pub fn push(self: *@This(), value: T) !void {
            const pushId = self.pushId.fetchAdd(1, .acq_rel);

            const slot = @mod(pushId, self.data.len);

            const newStatus = ConcurrentStatus{
                .generation = @intCast(@divTrunc(pushId, self.data.len)),
                .alive = true,
            };

            if (opts.debug) std.log.warn("pushing {d}", .{pushId});
            while (true) {
                const status = self.status[slot].load(.acquire);
                if (status.generation == newStatus.generation and status.alive == false) {
                    break;
                }
                std.atomic.spinLoopHint();
            }

            if (opts.debug) std.log.warn("pushed {d}", .{pushId});
            self.status[slot].store(newStatus, .release);
            self.data[slot] = value;
        }

        // pop the value from the queue, moves the head forward
        pub fn pop(self: *@This()) ?T { // things are empty
            if (self.count() == 0)
                return null;

            const popId = self.popId.fetchAdd(1, .acq_rel);

            const slot = @mod(popId, self.data.len);
            const expectedStatus = ConcurrentStatus{
                .generation = @intCast(@divTrunc(popId, self.data.len)),
                .alive = true,
            };

            if (opts.debug) std.log.warn("popping {d}", .{popId});
            while (true) {
                const status = self.status[slot].load(.acquire);
                if (status.generation == expectedStatus.generation and status.alive) {
                    break;
                }
                std.atomic.spinLoopHint();
            }
            if (opts.debug) std.log.warn("popped {d}", .{popId});

            self.status[slot].store(.{ .alive = false, .generation = expectedStatus.generation +% 1 }, .release);
            return self.data[slot];
        }

        pub fn count(self: @This()) usize {
            const head = self.pushId.load(.acquire);
            const tail = self.popId.load(.acquire);

            if (opts.debug)
                std.log.warn("count {d}", .{head - tail});
            if (opts.allowAsserts) {
                asserts(tail <= head, "tail > head in concurrent queue, this shouldnt ever happen", .{}, "concurrent queue assert");
            }

            return head - tail;
        }

        pub fn capacity(self: @This()) usize {
            return self.data.len;
        }

        pub fn deinit(self: *@This(), allocator: std.mem.Allocator) void {
            allocator.free(self.data);
            allocator.free(self.status);
        }
    };
}

test "concurrent queue basic correctness test" {
    const Info = struct {
        x: u32 = 0,
    };

    const allocator = std.testing.allocator;

    var y = try ConcurrentQueueUnmanagedAdvanced(Info, .{ .allowAsserts = true, .debug = true }).initCapacity(allocator, 420);
    defer y.deinit(allocator);

    try y.push(.{});
    try y.push(.{ .x = 1 });
    try y.push(.{ .x = 2 });
    try y.push(.{ .x = 1 });

    try utils.assertf(y.count() == 4, "expected there to be {d} elements in queue, we saw {d}", .{ 4, y.count() });
    _ = y.pop();
    _ = y.pop();
    _ = y.pop();
    _ = y.pop();

    try utils.assertf(y.count() == 0, "expected there to be {d} elements in queue, we saw {d}", .{ 0, y.count() });

    var x = try ConcurrentQueueU(Info).initCapacity(allocator, 12);
    defer x.deinit(allocator);

    try x.push(.{ .x = 0 });
    try x.push(.{ .x = 1 });
    try x.push(.{ .x = 2 });
    try x.push(.{ .x = 3 });

    try x.push(.{ .x = 4 });
    try x.push(.{ .x = 5 });
    try x.push(.{ .x = 6 });
    try x.push(.{ .x = 7 });

    try x.push(.{ .x = 8 });
    try x.push(.{ .x = 9 });
    try x.push(.{ .x = 10 });

    // const maybeError = x.push(.{ .x = 11 });

    // try utils.assertf(maybeError == ConcurrentQueueError.QueueIsFull, "Expected queue to have an error", .{});
}

test "concurrent queue multiple producer single consumer" {
    // 1. create multiple threads
    const threadCount = 24;

    const Payload = struct {
        x: i64 = 0,
        arb: [4096]u8 = undefined,
    };

    const QueueType = ConcurrentQueueUnmanagedAdvanced(Payload, .{ .debug = false, .allowAsserts = true });

    const Wrap = struct {
        pub fn threadFunc(queueRef: *QueueType, id: i64, exitSignal: *Atomic(bool), pushedCountResults: *Atomic(i64)) void {
            var pushedCount: i64 = 0;
            while (!exitSignal.load(.monotonic)) {
                queueRef.push(.{
                    .x = id + pushedCount,
                }) catch unreachable;
                pushedCount += 1;
            }

            _ = pushedCountResults.fetchAdd(pushedCount, .seq_cst);
        }
    };

    var threads: [threadCount]std.Thread = undefined;

    var testQueue = try QueueType.initCapacity(std.testing.allocator, 4096 * 4);
    defer testQueue.deinit(std.testing.allocator);

    var exitSignalAtomic = Atomic(bool).init(false);
    var pushedCountResults = Atomic(i64).init(0);

    for (0..threadCount) |i| {
        threads[i] = try std.Thread.spawn(.{}, Wrap.threadFunc, .{ &testQueue, @as(i64, @intCast(i * 10000)), &exitSignalAtomic, &pushedCountResults });
    }

    // 5 second message pump test
    var oldTime: f64 = test_getTime();

    // 10 second test, 5 seconds of input, 5 seconds of drain
    const startTime: f64 = 2.2;
    var timeLeft: f64 = startTime;

    var poppedCount: i64 = 0;

    var signaled: bool = false;

    while (timeLeft > 0 or testQueue.count() > 0) {
        const newTime = test_getTime();
        const deltaTime = newTime - oldTime;

        if (timeLeft - deltaTime < 0.0 and !signaled) {
            signaled = true;
            exitSignalAtomic.store(true, .seq_cst);
        }

        timeLeft -= deltaTime;
        oldTime = newTime;
        if (testQueue.pop()) |x| {
            _ = x;
            poppedCount += 1;
        }
        std.atomic.spinLoopHint();
    }

    for (0..threadCount) |i| {
        threads[i].join();
    }

    std.debug.print("popped {d} entries in {d} seconds payloadSize: {d} dataRate: {d:.3} MiB/s time/event {d:.3} us \n", .{
        poppedCount,
        startTime,
        @sizeOf(Payload),
        @as(f64, @floatFromInt(@as(usize, @intCast(poppedCount)) * @sizeOf(Payload))) / startTime / 1024 / 1024,
        startTime / @as(f64, @floatFromInt(@as(usize, @intCast(poppedCount)))) * 1000 * 1000,
    });

    try utils.assertf(
        poppedCount == pushedCountResults.load(.seq_cst),
        "mismatched, we popped {d} records while the workers pushed {d}",
        .{
            poppedCount,
            pushedCountResults.load(.seq_cst),
        },
    );
}

fn test_getTime() f64 {
    return @as(f64, @floatFromInt(std.time.milliTimestamp())) / 1000;
}