const std = @import("std");
const mem = std.mem;
const testing = std.testing;
const Allocator = mem.Allocator;

pub fn ComponentStore(comptime TComponent: type) type {
    return struct {
        const Self = @This();

        nextIndex: usize = 0,
        allocator: *Allocator,
        entityIDToIndex: std.AutoArrayHashMap(usize, usize),
        indexToEntityID: std.ArrayList(usize),
        components: std.ArrayList(TComponent),

        pub fn init(allocator: *Allocator) Self {
            return Self {
                .allocator = allocator,
                .entityIDToIndex = std.AutoArrayHashMap(usize, usize).init(allocator),
                .indexToEntityID = std.ArrayList(usize).init(allocator),
                .components = std.ArrayList(TComponent).init(allocator),
            };
        }

        pub fn deinit(self: *Self) void {
            self.entityIDToIndex.deinit();
            self.indexToEntityID.deinit();
            self.components.deinit();
        }

        pub fn has(self: Self, entityID: usize) bool {
            return self.entityIDToIndex.contains(entityID);
        }

        pub fn set(self: *Self, entityID: usize, component: TComponent) void {
            if (!self.entityIDToIndex.contains(entityID)) {
                const index = self.nextIndex;
                self.nextIndex += 1;

                // FIXME: actually handle OOM
                self.components.append(component) catch unreachable;
                self.entityIDToIndex.put(entityID, index) catch unreachable;
                self.indexToEntityID.append(entityID) catch unreachable;
            }
            else
            {
                self.components.items[self.entityIDToIndex.get(entityID).?] = component;
            }
        }

        pub fn count(self: Self) usize {
            return self.components.items.len;
        }

        pub fn get(self: Self, entityID: usize) !TComponent {
            var index = self.entityIDToIndex.get(entityID) orelse error.NoComponentError;
            return self.components.items[try index];
        }

        pub fn remove(self: *Self, entityID: usize) void {
            if (self.entityIDToIndex.contains(entityID)) {
                var storageIndex = self.entityIDToIndex.get(entityID).?;
                var lastEntity = self.indexToEntityID.items[self.nextIndex - 1];

                self.entityIDToIndex.put(lastEntity, storageIndex) catch unreachable;

                var removedComponent = self.components.swapRemove(storageIndex);
                var removedEntityID = self.indexToEntityID.swapRemove(storageIndex);
                var removedEntry = self.entityIDToIndex.remove(entityID);

                self.nextIndex -= 1;
            }
        }

        pub fn all_components(self: Self) []TComponent {
            return self.components.items;
        }

        pub fn all_entities(self: Self) []usize {
            return self.indexToEntityID.items();
        }
    };
}

const TestComponent = struct {
    number: i32,
};

test "component store set item" {
    const component = TestComponent {
        .number = 3
    };

    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();

    const allocator = &arena.allocator;

    var component_store = ComponentStore(TestComponent).init(allocator);
    component_store.set(12, component);

    testing.expect(component_store.count() == 1);
    testing.expect(component_store.has(12));
    testing.expect((try component_store.get(12)).number == component.number);
}

test "iterate items" {
    const component = TestComponent {
        .number = 3
    };

    const componentTwo = TestComponent {
        .number = 10
    };

    const componentThree = TestComponent {
        .number = 15
    };

    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();

    const allocator = &arena.allocator;

    var component_store = ComponentStore(TestComponent).init(allocator);
    component_store.set(3, component);
    component_store.set(5, componentTwo);
    component_store.set(8, componentThree);

    var all_components = component_store.all_components();

    testing.expect(all_components[0].number == 3);
    testing.expect(all_components[1].number == 10);
    testing.expect(all_components[2].number == 15);
}

test "has item" {
    const component = TestComponent {
        .number = 3
    };

    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();

    const allocator = &arena.allocator;

    var component_store = ComponentStore(TestComponent).init(allocator);
    component_store.set(5, component);

    testing.expect(component_store.has(5));
    testing.expect(!component_store.has(6));
}

test "remove item" {
    const component = TestComponent {
        .number = 5
    };

    const componentTwo = TestComponent {
        .number = 87
    };

    const componentThree = TestComponent {
        .number = 124
    };

    var arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena.deinit();

    const allocator = &arena.allocator;

    var component_store = ComponentStore(TestComponent).init(allocator);
    component_store.set(5, component);
    component_store.set(10, componentTwo);
    component_store.set(15, componentThree);

    component_store.remove(10);

    testing.expect(component_store.has(5));
    testing.expect(component_store.has(15));
    testing.expect(!component_store.has(10));

    testing.expect((try component_store.get(5)).number == 5);
    testing.expect((try component_store.get(15)).number == 124);
}