--!native
--!strict
--!optimize 2

export type SystemPhase = "pre" | "update" | "post"

export type ComponentTable = { [number]: T }

type ComponentStore = { [string]: ComponentTable }
type ComponentEntityLists = { [string]: { number } }
type ComponentIndexMaps = { [string]: { [number]: number } }
type EntityComponentSets = { [number]: { [string]: boolean } }

export type SystemFn = (world: World, dt: number) -> ()

export type WorldStats = {
	entityCount: number,
	componentCounts: { [string]: number },
	systemCounts: {
		pre: number,
		update: number,
		post: number,
	},
	lastStepDt: number,
	totalSteps: number,
}

export type WorldOptions = {
	initialEntityCapacity: number?,
	initialComponentEntityCapacity: number?,
	scratchSize: number?,
	debugProfiling: boolean?,
	debugAssertions: boolean?,
}

export type DebugOptions = {
	profiling: boolean?,
	assertions: boolean?,
}

export type World = {
	_nextEntityId: number,
	_freeEntityIds: { number },

	_entityCapacity: number,
	_componentEntityCapacity: number,
	_scratchSize: number,

	_alive: { [number]: boolean },

	_components: ComponentStore,
	_componentEntities: ComponentEntityLists,
	_componentIndex: ComponentIndexMaps,

	_entityComponents: EntityComponentSets,
	_componentScratch: { string },

	_systemsPre: { SystemFn },
	_systemsUpdate: { SystemFn },
	_systemsPost: { SystemFn },
	_hasSystems: boolean,

	_debugProfiling: boolean,
	_debugAssertions: boolean,

	_lastStepDt: number,
	_totalSteps: number,

	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: T) -> (),
	GetComponent: (self: World, entityId: number, name: string) -> T?,
	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: T) -> ()) -> (),
	ForEach2: (
		self: World,
		nameA: string,
		nameB: string,
		fn: (entityId: number, dataA: A, dataB: B) -> ()
	) -> (),

	Clear: (self: World) -> (),

	SetDebugOptions: (self: World, options: DebugOptions) -> (),
	GetStats: (self: World) -> WorldStats,
}

type ECSWorldModule = {
	new: (options: WorldOptions?) -> World,
}

local ECSWorld: ECSWorldModule = {}

local DEFAULT_ENTITY_CAPACITY = 256
local DEFAULT_COMPONENT_ENTITY_CAPACITY = 64
local DEFAULT_SCRATCH_SIZE = 16

local debugLib = debug :: any

local profileBegin: ((string) -> ())? = nil
local profileEnd: (() -> ())? = nil

if debugLib and debugLib.profilebegin and debugLib.profileend then
	profileBegin = debugLib.profilebegin :: (string) -> ()
	profileEnd = debugLib.profileend :: () -> ()
end

local function runSystems(list: { SystemFn }, world: World, dt: number)
	local count = #list
	if count == 0 then
		return
	end

	for i = 1, count do
		list[i](world, dt)
	end
end

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

	local newTbl = {} :: ComponentTable
	components[name] = newTbl

	return newTbl
end

local function ensureComponentMeta(self: World, name: string)
	local entitiesByComponent = self._componentEntities
	local indexByComponent = self._componentIndex

	if entitiesByComponent[name] == nil then
		local list = table.create(self._componentEntityCapacity)
		table.clear(list)
		entitiesByComponent[name] = list :: { number }
	end

	if indexByComponent[name] == nil then
		indexByComponent[name] = {}
	end
end

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

	local data = dataAny :: ComponentTable
	if data[entityId] == nil then
		return
	end

	data[entityId] = nil

	local entitiesByComponent = self._componentEntities
	local indexByComponent = self._componentIndex

	local entityList = entitiesByComponent[name]
	local indexMap = indexByComponent[name]

	if entityList ~= nil and indexMap ~= nil then
		local idx = indexMap[entityId]
		if idx ~= nil then
			local lastIndex = #entityList
			local lastId = entityList[lastIndex]
			entityList[lastIndex] = nil

			if idx ~= lastIndex then
				entityList[idx] = lastId
				indexMap[lastId] = idx
			end

			indexMap[entityId] = nil
		end
	end

	local entityComponents = self._entityComponents
	local compSet = entityComponents[entityId]
	if compSet ~= nil then
		compSet[name] = nil
		if next(compSet) == nil then
			entityComponents[entityId] = nil
		end
	end
end

local function spawnEntity(self: World): number
	local free = self._freeEntityIds
	local freeCount = #free

	local id: number
	if freeCount > 0 then
		id = free[freeCount]
		free[freeCount] = nil
	else
		local nextId = self._nextEntityId + 1
		self._nextEntityId = nextId
		id = nextId

		if nextId > self._entityCapacity then
			self._entityCapacity = nextId
		end
	end

	self._alive[id] = true
	return id
end

local function despawnEntity(self: World, entityId: number)
	if not self._alive[entityId] then
		if self._debugAssertions then
			error("ECSWorld.Despawn: entity is not alive")
		end
		return
	end

	self._alive[entityId] = nil

	local entityComponents = self._entityComponents
	local compSet = entityComponents[entityId]
	if compSet ~= nil then
		local scratch = self._componentScratch
		local count = 0

		for name in compSet do
			count += 1
			scratch[count] = name
		end

		for i = 1, count do
			local name = scratch[i]
			scratch[i] = nil
			detachComponent(self, entityId, name)
		end
	end

	local free = self._freeEntityIds
	free[#free + 1] = entityId
end

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

local function addComponent(self: World, entityId: number, name: string, data: T)
	if not self._alive[entityId] then
		if self._debugAssertions then
			error("ECSWorld.AddComponent: entity is not alive")
		end
		return
	end

	local tbl = getComponentTable(self, name)
	tbl[entityId] = data

	ensureComponentMeta(self, name)

	local entityList = self._componentEntities[name] :: { number }
	local indexMap = self._componentIndex[name] :: { [number]: number }

	local entityComponents = self._entityComponents
	local compSet = entityComponents[entityId]
	if compSet == nil then
		local newSet: { [string]: boolean } = {}
		entityComponents[entityId] = newSet
		compSet = newSet
	end
	compSet[name] = true

	if indexMap[entityId] ~= nil then
		return
	end

	local newIndex = #entityList + 1
	entityList[newIndex] = entityId
	indexMap[entityId] = newIndex
end

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

local function removeComponent(self: World, entityId: number, name: string)
	detachComponent(self, entityId, name)
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
	elseif phase == "post" then
		local list = self._systemsPost
		list[#list + 1] = system
	else
		error(("ECSWorld.RegisterSystem: invalid phase '%s'"):format(tostring(phase)))
	end

	self._hasSystems = true
end

local function step(self: World, dt: number)
	if dt <= 0 or not self._hasSystems then
		return
	end

	self._lastStepDt = dt
	self._totalSteps += 1

	local profilingEnabled = self._debugProfiling and profileBegin ~= nil and profileEnd ~= nil

	if profilingEnabled then
		local pb = profileBegin :: (string) -> ()
		local pe = profileEnd :: () -> ()

		pb("ECSWorld::Step(pre)")
		runSystems(self._systemsPre, self, dt)
		pe()

		pb("ECSWorld::Step(update)")
		runSystems(self._systemsUpdate, self, dt)
		pe()

		pb("ECSWorld::Step(post)")
		runSystems(self._systemsPost, self, dt)
		pe()
	else
		runSystems(self._systemsPre, self, dt)
		runSystems(self._systemsUpdate, self, dt)
		runSystems(self._systemsPost, self, dt)
	end
end

local function forEach(self: World, name: string, fn: (entityId: number, data: T) -> ())
	local entityList = self._componentEntities[name]
	if entityList == nil then
		return
	end

	local dataTblAny = self._components[name]
	if dataTblAny == nil then
		return
	end

	local dataTbl = dataTblAny :: ComponentTable
	local alive = self._alive
	local count = #entityList
	local fnLocal = fn

	for i = 1, count do
		local entityId = entityList[i]
		if alive[entityId] then
			fnLocal(entityId, dataTbl[entityId])
		end
	end
end

local function forEach2(
	self: World,
	nameA: string,
	nameB: string,
	fn: (entityId: number, dataA: A, dataB: B) -> ()
)
	local entitiesA = self._componentEntities[nameA]
	local entitiesB = self._componentEntities[nameB]
	if entitiesA == nil or entitiesB == nil then
		return
	end

	local dataAAny = self._components[nameA]
	local dataBAny = self._components[nameB]
	if dataAAny == nil or dataBAny == nil then
		return
	end

	local alive = self._alive

	local listSmall, listLarge
	local dataSmallAny, dataLargeAny
	local nameSmall, nameLarge

	if #entitiesA <= #entitiesB then
		listSmall, listLarge = entitiesA, entitiesB
		dataSmallAny, dataLargeAny = dataAAny, dataBAny
		nameSmall, nameLarge = nameA, nameB
	else
		listSmall, listLarge = entitiesB, entitiesA
		dataSmallAny, dataLargeAny = dataBAny, dataAAny
		nameSmall, nameLarge = nameB, nameA
	end

	local count = #listSmall
	if count == 0 then
		return
	end

	local indexLarge = self._componentIndex[nameLarge]
	if indexLarge == nil then
		return
	end

	local fnLocal = fn

	if nameSmall == nameA then
		local dataA = dataSmallAny :: ComponentTable
		local dataB = dataLargeAny :: ComponentTable

		for i = 1, count do
			local entityId = listSmall[i]
			if alive[entityId] and indexLarge[entityId] ~= nil then
				fnLocal(entityId, dataA[entityId], dataB[entityId])
			end
		end
	else
		local dataB = dataSmallAny :: ComponentTable
		local dataA = dataLargeAny :: ComponentTable

		for i = 1, count do
			local entityId = listSmall[i]
			if alive[entityId] and indexLarge[entityId] ~= nil then
				fnLocal(entityId, dataA[entityId], dataB[entityId])
			end
		end
	end
end

local function clearWorld(self: World)
	self._nextEntityId = 0
	self._lastStepDt = 0
	self._totalSteps = 0

	table.clear(self._freeEntityIds)
	table.clear(self._alive)

	for _, tbl in self._components do
		table.clear(tbl)
	end

	for _, list in self._componentEntities do
		table.clear(list)
	end

	for _, indexMap in self._componentIndex do
		table.clear(indexMap)
	end

	table.clear(self._entityComponents)
	table.clear(self._componentScratch)
end

local function setDebugOptions(self: World, options: DebugOptions)
	local profiling = options.profiling
	if profiling ~= nil then
		self._debugProfiling = profiling
	end

	local assertions = options.assertions
	if assertions ~= nil then
		self._debugAssertions = assertions
	end
end

local function getStats(self: World): WorldStats
	local entityCount = 0
	for _ in self._alive do
		entityCount += 1
	end

	local componentCounts: { [string]: number } = {}
	for name, list in self._componentEntities do
		componentCounts[name] = #list
	end

	local systemCounts = {
		pre = #self._systemsPre,
		update = #self._systemsUpdate,
		post = #self._systemsPost,
	}

	return {
		entityCount = entityCount,
		componentCounts = componentCounts,
		systemCounts = systemCounts,
		lastStepDt = self._lastStepDt,
		totalSteps = self._totalSteps,
	}
end

function ECSWorld.new(options: WorldOptions?): World
	local opt = options or {}

	local entityCapacity = opt.initialEntityCapacity or DEFAULT_ENTITY_CAPACITY
	local componentEntityCapacity = opt.initialComponentEntityCapacity or DEFAULT_COMPONENT_ENTITY_CAPACITY
	local scratchSize = opt.scratchSize or DEFAULT_SCRATCH_SIZE

	if entityCapacity < 1 then
		entityCapacity = DEFAULT_ENTITY_CAPACITY
	end

	if componentEntityCapacity < 1 then
		componentEntityCapacity = DEFAULT_COMPONENT_ENTITY_CAPACITY
	end

	if scratchSize < 1 then
		scratchSize = DEFAULT_SCRATCH_SIZE
	end

	local debugProfiling = opt.debugProfiling == true
	local debugAssertions = opt.debugAssertions == true

	local freeEntityIds = table.create(entityCapacity)
	table.clear(freeEntityIds)

	local world: World = {
		_nextEntityId = 0,
		_freeEntityIds = freeEntityIds,

		_entityCapacity = entityCapacity,
		_componentEntityCapacity = componentEntityCapacity,
		_scratchSize = scratchSize,

		_alive = {},

		_components = {} :: ComponentStore,
		_componentEntities = {} :: ComponentEntityLists,
		_componentIndex = {} :: ComponentIndexMaps,

		_entityComponents = {} :: EntityComponentSets,
		_componentScratch = table.create(scratchSize) :: { string },

		_systemsPre = {},
		_systemsUpdate = {},
		_systemsPost = {},
		_hasSystems = false,

		_debugProfiling = debugProfiling,
		_debugAssertions = debugAssertions,

		_lastStepDt = 0,
		_totalSteps = 0,

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

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

		RegisterSystem = registerSystem,
		Step = step,

		ForEach = forEach,
		ForEach2 = forEach2,

		Clear = clearWorld,

		SetDebugOptions = setDebugOptions,
		GetStats = getStats,
	}

	return world
end

return ECSWorld