--!optimize 2
--!native

local jecs = require(game:GetService("ReplicatedStorage").ecs)

type World = jecs.WorldShim
type Entity<T = any> = jecs.Entity<T>

local function panic(str)
    -- We don't want to interrupt the loop when we error
    task.spawn(error, str)
end

local function Scheduler(world, ...)
    local systems = { ... }
    local systemsNames = {}
    local N = #systems
    local system
    local dt

    for i, module in systems do
        local sys = require(module)
        systems[i] = sys
        local file, line = debug.info(2, "sl")
        systemsNames[sys] = `{file}->::{line}::->{debug.info(sys, "n")}`
    end

    local function run()
        local name = systemsNames[system]

        debug.profilebegin(name)
        debug.setmemorycategory(name)
        system(world, dt)
        debug.profileend()
    end

    local function loop(sinceLastFrame)
        debug.profilebegin("loop()")

        for i = N, 1, -1 do
            system = systems[i]

            dt = sinceLastFrame

            local didNotYield, why = xpcall(function()
                for _ in run do end
            end, debug.traceback)

            if didNotYield then
				continue
			end

			if string.find(why, "thread is not yieldable") then
				N -= 1
				local name = table.remove(systems, i)
				panic("Not allowed to yield in the systems."
    				.. "\n"
    				.. `System: {name} has been ejected`
				)
			else
			    panic(why)
			end
        end

        debug.profileend()
        debug.resetmemorycategory()
    end

    return loop
end

type Tracker<T> = { track: (world: World, fn: (changes: {
        added: () -> () -> (number, T),
        removed: () -> () -> number,
        changed: () -> () -> (number, T, T)
    }) -> ()) -> ()
}

type Entity<T = any> = number & { __nominal_type_dont_use: T }

local function diff(a, b)
    local size = 0
    for k, v in a do
        if b[k] ~= v then
            return true
        end
        size += 1
    end
    for k, v in b do
        size -= 1
    end

    if size ~= 0 then
        return true
    end

    return false
end

local function ChangeTracker<T>(world, T: Entity<T>): Tracker<T>
    local PreviousT = jecs.pair(jecs.Rest, T)
    local add = {}
    local added
    local removed
    local is_trivial

    local function changes_added()
        added = true
        local q = world:query(T):without(PreviousT):drain()
        return function()
            local id, data = q.next()
            if not id then
                return nil
            end

            is_trivial = typeof(data) ~= "table"

            add[id] = data

            return id, data
        end
    end

    local function changes_changed()
        local q = world:query(T, PreviousT):drain()

        return function()
            local id, new, old = q.next()
            while true do
                if not id then
                    return nil
                end

                if not is_trivial then
                    if diff(new, old) then
                        break
                    end
                elseif new ~= old then
                    break
                end

                id, new, old = q.next()
            end

            local record = world.entityIndex.sparse[id]
            local archetype = record.archetype
            local column = archetype.records[PreviousT].column
            local data = if is_trivial then new else table.clone(new)
            archetype.columns[column][record.row] = data

            return id, old, new
        end
    end

    local function changes_removed()
        removed = true

        local q = world:query(PreviousT):without(T):drain()
        return function()
            local id = q.next()
            if id then
                world:remove(id, PreviousT)
            end
            return id
        end
    end

    local changes = {
        added = changes_added,
        changed = changes_changed,
        removed = changes_removed,
    }

    local function track(fn)
        added = false
        removed = false

        fn(changes)

        if not added then
            for _ in changes_added() do
            end
        end

        if not removed then
            for _ in changes_removed() do
            end
        end

        for e, data in add do
            world:set(e, PreviousT, if is_trivial then data else table.clone(data))
        end
    end

    local tracker = { track = track }

    return tracker
end

local bt
do
    local SUCCESS = 0
    local FAILURE = 1
    local RUNNING = 2

    local function SEQUENCE(nodes)
        return function(...)
            for _, node in nodes do
                local status = node(...)
                if status == FAILURE or status == RUNNING then
                    return status
                end
            end
            return SUCCESS
        end
    end
    local function FALLBACK(nodes)
        return function(...)
            for _, node in nodes do
                local status = node(...)
                if status == SUCCESS or status == RUNNING then
                    return status
                end
            end
            return FAILURE
        end
    end
    bt = {
        SEQUENCE = SEQUENCE,
        FALLBACK = FALLBACK,
        RUNNING = RUNNING
    }
end

local function interval(s)
    local pin

    local function throttle()
        if not pin then
            pin = os.clock()
        end

        local elapsed = os.clock() - pin > s
        if elapsed then
            pin = os.clock()
        end

        return elapsed
    end
    return throttle
end

return {
    Scheduler = Scheduler,
    ChangeTracker = ChangeTracker,
    interval = interval,
    BehaviorTree = bt
}