--!strict
--!optimize 2

export type SystemPhase = "pre" | "update" | "post"
export type SystemFn = (world: World, dt: number) -> ()
export type ComponentTable = { [number]: any }

export type World = {
    _nextEntityId: number,
    _alive: { [number]: boolean },
    _components: { [string]: ComponentTable },
    _systemsPre: { SystemFn },
    _systemsUpdate: { SystemFn },
    _systemsPost: { SystemFn },

    Spawn: (self: World) -> number,
    Despawn: (self: World, entityId: number) -> (),

    HasEntity: (self: World, entityId: number) -> boolean,

    GetComponentTable: (self: World, name: string) -> ComponentTable,
    AddComponent: (self: World, entityId: number, name: string, data: any) -> (),
    GetComponent: (self: World, entityId: number, name: string) -> any?,
    RemoveComponent: (self: World, entityId: number, name: string) -> (),

    RegisterSystem: (self: World, phase: SystemPhase, system: SystemFn) -> (),
    Step: (self: World, dt: number) -> (),

    ForEach: (self: World, name: string, fn: (entityId: number, data: any) -> ()) -> (),
}

local ECSWorld = {}

local function runSystems(list: { SystemFn }, world: World, dt: number)
    local count = #list
    for i = 1, count do
        list[i](world, dt)
    end
end

local function spawnEntity(self: World): number
    local id = self._nextEntityId + 1
    self._nextEntityId = id
    self._alive[id] = true
    return id
end

local function despawnEntity(self: World, entityId: number)
    if not self._alive[entityId] then
        return
    end
    self._alive[entityId] = nil
    local components = self._components
    for _, tableForComponent in pairs(components) do
        tableForComponent[entityId] = nil
    end
end

local function hasEntity(self: World, entityId: number): boolean
    return self._alive[entityId] == true
end

local function getComponentTable(self: World, name: string): ComponentTable
    local components = self._components
    local tbl = components[name]
    if tbl ~= nil then
        return tbl
    end
    local newTbl: ComponentTable = {}
    components[name] = newTbl
    return newTbl
end

local function addComponent(self: World, entityId: number, name: string, data: any)
    if not self._alive[entityId] then
        return
    end
    local tbl = getComponentTable(self, name)
    tbl[entityId] = data
end

local function getComponent(self: World, entityId: number, name: string): any?
    local tbl = self._components[name]
    if tbl == nil then
        return nil
    end
    return tbl[entityId]
end

local function removeComponent(self: World, entityId: number, name: string)
    local tbl = self._components[name]
    if tbl == nil then
        return
    end
    tbl[entityId] = nil
end

local function registerSystem(self: World, phase: SystemPhase, system: SystemFn)
    if phase == "pre" then
        local list = self._systemsPre
        list[#list + 1] = system
    elseif phase == "update" then
        local list = self._systemsUpdate
        list[#list + 1] = system
    else
        local list = self._systemsPost
        list[#list + 1] = system
    end
end

local function step(self: World, dt: number)
    if dt <= 0 then
        return
    end
    runSystems(self._systemsPre, self, dt)
    runSystems(self._systemsUpdate, self, dt)
    runSystems(self._systemsPost, self, dt)
end

local function forEach(self: World, name: string, fn: (entityId: number, data: any) -> ())
    local tbl = self._components[name]
    if tbl == nil then
        return
    end
    for entityId, data in tbl do
        if self._alive[entityId] then
            fn(entityId, data)
        end
    end
end

function ECSWorld.new(): World
    local world: World = {
        _nextEntityId = 0,
        _alive = {},
        _components = {},
        _systemsPre = {},
        _systemsUpdate = {},
        _systemsPost = {},

        Spawn = spawnEntity,
        Despawn = despawnEntity,
        HasEntity = hasEntity,

        GetComponentTable = getComponentTable,
        AddComponent = addComponent,
        GetComponent = getComponent,
        RemoveComponent = removeComponent,

        RegisterSystem = registerSystem,
        Step = step,

        ForEach = forEach,
    }

    return world
end

return ECSWorld