local jecs = require("@jecs")

type World = jecs.WorldShim

type Tracker<T> = {
	track: (
		world: World,
		fn: (
			changes: {
				added: () -> () -> (number, T),
				removed: () -> () -> number,
				changed: () -> () -> (number, T, 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

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

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

			add[id] = new

			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 Vector3
do
	Vector3 = {}
	Vector3.__index = Vector3

	function Vector3.new(x, y, z)
		x = x or 0
		y = y or 0
		z = z or 0
		return setmetatable({ X = x, Y = y, Z = z }, Vector3)
	end

	function Vector3.__add(left, right)
		return Vector3.new(left.X + right.X, left.Y + right.Y, left.Z + right.Z)
	end

	function Vector3.__mul(left, right)
		if typeof(right) == "number" then
			return Vector3.new(left.X * right, left.Y * right, left.Z * right)
		end
		return Vector3.new(left.X * right.X, left.Y * right.Y, left.Z * right.Z)
	end

	Vector3.one = Vector3.new(1, 1, 1)
	Vector3.zero = Vector3.new()
end

local world = jecs.World.new()
local Name = world:component()

local function named(ctr, name)
	local e = ctr(world)
	world:set(e, Name, name)
	return e
end
local function name(e)
	return world:get(e, Name)
end

local Position = named(world.component, "Position")

-- Create the ChangeTracker with the component type to track
local PositionTracker = ChangeTracker(world, Position)

local e1 = named(world.entity, "e1")
world:set(e1, Position, Vector3.new(10, 20, 30))

local e2 = named(world.entity, "e2")
world:set(e2, Position, Vector3.new(10, 20, 30))

PositionTracker.track(function(changes)
	-- You can iterate over different types of changes: Added, Changed, Removed

	-- added queries for every entity with a new Position component
	for e, p in changes.added() do
		print(`Added {e}: \{{p.X}, {p.Y}, {p.Z}}`)
	end

	-- changed queries for entities who's changed their data since
	-- last was it tracked
	for _ in changes.changed() do
		print([[This won't print because it is the first time
                we are tracking the Position component]])
	end

	-- removed queries for entities who's removed their Position component
	-- since last it was tracked
	for _ in changes.removed() do
		print([[This won't print because it is the first time
                we are tracking the Position component]])
	end
end)

world:set(e1, Position, Vector3.new(1, 1, 2) * 999)

PositionTracker.track(function(changes)
	for e, p in changes.added() do
		print([[This won't never print no Position component was added
                since last time we tracked]])
	end

	for e, old, new in changes.changed() do
		print(`{name(e)}'s Position changed from \{{old.X}, {old.Y}, {old.Z}\} to \{{new.X}, {new.Y}, {new.Z}\}`)
	end

	-- If you don't call e.g. changes.removed() then it will automatically drain its iterator and stage their changes.
	-- This ensures you will not have any off-by-one frame errors.
end)

world:remove(e2, Position)

PositionTracker.track(function(changes)
	for e in changes.removed() do
		print(`Position was removed from {name(e)}`)
	end
end)

-- Output:
--  Added 265: {10, 20, 30}
--  Added 264: {10, 20, 30}
--  e1's Position changed from {10, 20, 30} to {999, 999, 1998}
--  Position was removed from e2
