--!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