local defs = require("@self/defs")
local luau = require("@lute/luau")
local var_tracking = require("@self/var_tracking")

export type UnaryOperator = defs.UnaryOperator
export type BinaryOperator = defs.BinaryOperator
export type UnknownConstant = defs.UnknownConstant
export type NilConstant = defs.NilConstant
export type BooleanConstant = defs.BooleanConstant
export type NumberConstant = defs.NumberConstant
export type VectorConstant = defs.VectorConstant
export type StringConstant = defs.StringConstant
export type TableConstant = defs.TableConstant
export type Constant = defs.Constant
export type KnownConstant = defs.KnownConstant
export type SimpleConstant = defs.SimpleConstant
export type ConstantData = defs.ConstantData
export type Variable = defs.Variable
export type Tracked = defs.Tracked

local UNKNOWN: UnknownConstant = table.freeze({
	kind = "unknown" :: "unknown",
	data = nil,
})

local function EXHAUSTIVE_MATCH<Err>(value: never, err: Err?): never
	error(err or `Unknown value in exhaustive match: {value}`)
end

local function ASSERT<Value, Error>(value: Value, message: Error): typeof(assert(... :: Value))
	if value then
		return value
	end
	error(message)
end

local function truthy_eh(operand: KnownConstant): boolean
	if (operand.kind ~= "nil") and ((operand.kind ~= "boolean") or (operand.data ~= false)) then
		return true
	else
		return false
	end
end

local function cunknown(): UnknownConstant
	return UNKNOWN
end

local function cbool(data: boolean): BooleanConstant
	return {
		kind = "boolean",
		data = data,
	}
end

local function cnumber(data: number): NumberConstant
	return {
		kind = "number",
		data = data,
	}
end

local function cvector(data: vector): VectorConstant
	return {
		kind = "vector",
		data = data,
	}
end

local function cstring(data: string): StringConstant
	return {
		kind = "string",
		data = data,
	}
end

local function cnil(): NilConstant
	return {
		kind = "nil",
		data = nil,
	}
end

local analyzer = {}
local metatable = { __index = analyzer }
export type Identity = setmetatable<{
	tracked: Tracked,
	cache: { [luau.AstExpr]: Constant },
}, typeof(metatable)>
local function constructor(root: luau.AstStatBlock): Identity
	local tracked = var_tracking.track(root)
	return setmetatable({
		tracked = tracked,
		cache = {},
	}, metatable)
end
local function constructor_from_tracked(tracked: Tracked): Identity
	return setmetatable({
		tracked = tracked,
		cache = {},
	}, metatable)
end

function analyzer.truthy(self: Identity, operand: Constant): BooleanConstant | UnknownConstant
	if operand.kind == "unknown" then
		return cunknown()
	end
	return cbool(truthy_eh(operand))
end

function analyzer.equal(self: Identity, left: Constant, right: Constant): BooleanConstant | UnknownConstant
	if (left.kind == "unknown") or (right.kind == "unknown") then
		return cunknown()
	end
	if left.kind == "nil" then
		return cbool(right.kind == "nil")
	elseif left.kind == "boolean" then
		return cbool((right.kind == "boolean") and (left.data == right.data))
	elseif left.kind == "number" then
		return cbool((right.kind == "number") and (left.data == right.data))
	elseif left.kind == "vector" then
		return cbool((right.kind == "vector") and (left.data == right.data))
	elseif left.kind == "string" then
		return cbool((right.kind == "string") and (left.data == right.data))
	elseif left.kind == "table" then
		return cunknown()
	else
		return EXHAUSTIVE_MATCH(left.kind)
	end
end

function analyzer.unary(
	self: Identity,
	operator: UnaryOperator,
	operand: Constant
): BooleanConstant | NumberConstant | UnknownConstant
	if operand.kind == "unknown" then
		return cunknown()
	end
	if operator == "not" then
		return cbool(truthy_eh(operand))
	elseif operator == "-" then
		if operand.kind == "number" then
			return cnumber(operand.data)
		end
	elseif operator == "#" then
		if operand.kind == "string" then
			return cnumber(#operand.data)
		elseif operand.kind == "table" then
			local data = operand.data
			local metatable = data.metatable
			if metatable then
				local len = metatable.__len
				if len then
					return len(operand)
				end
			end
			local record = data.record
			if record then
				return cnumber(#record)
			end
		end
	else
		return EXHAUSTIVE_MATCH(operator)
	end
	return cunknown()
end

function analyzer.binary(
	self: Identity,
	lhs_operand: Constant,
	operator: BinaryOperator,
	rhs_operand: Constant
): Constant
	if (lhs_operand.kind == "unknown") or (rhs_operand.kind == "unknown") then
		return cunknown()
	end
	if operator == "+" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cnumber(lhs_operand.data + rhs_operand.data)
		elseif (lhs_operand.kind == "vector") and (rhs_operand.kind == "vector") then
			return cvector(lhs_operand.data + rhs_operand.data)
		end
	elseif operator == "-" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cnumber(lhs_operand.data - rhs_operand.data)
		elseif (lhs_operand.kind == "vector") and (rhs_operand.kind == "vector") then
			return cvector(lhs_operand.data - rhs_operand.data)
		end
	elseif operator == "*" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cnumber(lhs_operand.data * rhs_operand.data)
		elseif
			(lhs_operand.kind == "vector") and ((rhs_operand.kind == "vector") or (rhs_operand.kind == "number"))
		then
			return cvector(lhs_operand.data * rhs_operand.data :: vector | number)
		elseif (lhs_operand.kind == "number") and (rhs_operand.kind == "vector") then
			return cvector(lhs_operand.data * rhs_operand.data)
		end
	elseif operator == "/" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cnumber(lhs_operand.data / rhs_operand.data)
		elseif
			(lhs_operand.kind == "vector") and ((rhs_operand.kind == "vector") or (rhs_operand.kind == "number"))
		then
			return cvector(lhs_operand.data / rhs_operand.data :: vector | number)
		elseif (lhs_operand.kind == "number") and (rhs_operand.kind == "vector") then
			return cvector(lhs_operand.data / rhs_operand.data)
		end
	elseif operator == "//" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cnumber(lhs_operand.data // rhs_operand.data)
		elseif
			(lhs_operand.kind == "vector") and ((rhs_operand.kind == "vector") or (rhs_operand.kind == "number"))
		then
			return cvector(lhs_operand.data // rhs_operand.data :: vector | number)
		elseif (lhs_operand.kind == "number") and (rhs_operand.kind == "vector") then
			return cvector(lhs_operand.data // rhs_operand.data)
		end
	elseif operator == "%" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cnumber(lhs_operand.data % rhs_operand.data)
		end
	elseif operator == "^" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cnumber(lhs_operand.data ^ rhs_operand.data)
		end
	elseif operator == ".." then
		if (lhs_operand.kind == "table") or (rhs_operand.kind == "table") then
			return cunknown()
		end
		return cstring(lhs_operand.data :: any .. rhs_operand.data)
	elseif operator == "==" then
		return self:equal(lhs_operand, rhs_operand)
	elseif operator == "~=" then
		local equality = self:equal(lhs_operand, rhs_operand)
		if equality.kind == "boolean" then
			equality.data = not equality.data
		end
		return equality
	elseif operator == "<=" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cbool(lhs_operand.data <= rhs_operand.data)
		elseif (lhs_operand.kind == "string") and (rhs_operand.kind == "string") then
			return cbool(lhs_operand.data <= rhs_operand.data)
		end
	elseif operator == ">=" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cbool(lhs_operand.data >= rhs_operand.data)
		elseif (lhs_operand.kind == "string") and (rhs_operand.kind == "string") then
			return cbool(lhs_operand.data >= rhs_operand.data)
		end
	elseif operator == "<" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cbool(lhs_operand.data < rhs_operand.data)
		elseif (lhs_operand.kind == "string") and (rhs_operand.kind == "string") then
			return cbool(lhs_operand.data < rhs_operand.data)
		end
	elseif operator == ">" then
		if (lhs_operand.kind == "number") and (rhs_operand.kind == "number") then
			return cbool(lhs_operand.data < rhs_operand.data)
		elseif (lhs_operand.kind == "string") and (rhs_operand.kind == "string") then
			return cbool(lhs_operand.data < rhs_operand.data)
		end
	elseif operator == "and" then
		if truthy_eh(lhs_operand) then
			return rhs_operand
		else
			return lhs_operand
		end
	elseif operator == "or" then
		if truthy_eh(lhs_operand) then
			return lhs_operand
		else
			return rhs_operand
		end
	else
		return EXHAUSTIVE_MATCH(operator)
	end
	return cunknown()
end

function analyzer.evaluate_conditional(self: Identity, node: luau.AstExprIfElse): Constant
	local condition = self:evaluate_unit(node.condition)
	if condition.kind == "unknown" then
		return cunknown()
	end
	if truthy_eh(condition) then
		return (self:evaluate_unit(node.consequent))
	end
	for _, elif in node.elseifs do
		condition = self:evaluate_unit(elif.condition)
		if condition.kind == "unknown" then
			return cunknown()
		end
		if truthy_eh(condition) then
			return (self:evaluate_unit(elif.consequent))
		end
	end
	return (self:evaluate_unit(node.antecedent))
end

function analyzer.evaluate_interpstring(
	self: Identity,
	node: luau.AstExprInterpString
): StringConstant | UnknownConstant
	local str = ""
	local strings = node.strings
	local expressions = node.expressions
	for index = 1, #expressions do
		local expr = self:evaluate_unit(expressions[index])
		if (expr.kind == "unknown") or (expr.kind == "table") then
			return cunknown()
		end
		str ..= `{strings[index]}{expr.data}`
	end
	str ..= strings[#strings].text
	return {
		kind = "string" :: "string",
		data = str,
	}
end

function analyzer.evaluate_indexname(self: Identity, node: luau.AstExprIndexName): Constant
	local operand = self:evaluate_unit(node.expression)
	if operand.kind ~= "table" then
		if operand.kind ~= "vector" then
			return cunknown()
		end
		local text = string.lower(node.index.text)
		if text == "x" then
			return cnumber(operand.data.x)
		elseif text == "y" then
			return cnumber(operand.data.y)
		elseif text == "z" then
			return cnumber(operand.data.z)
		else
			return cunknown()
		end
	end
	local data: index<TableConstant, "data"> = operand.data
	local index = node.index.text

	local result: Constant?

	local record = data.record
	if record then
		result = record[index]
	end

	local metatable = not result and data.metatable
	if metatable then
		local indexer = metatable.__index
		if indexer then
			result = indexer(operand, cstring(index))
		end
	end

	return result or cunknown()
end

function analyzer.evaluate_indexexpr(self: Identity, node: luau.AstExprIndexExpr): Constant
	local operand = self:evaluate_unit(node.expression)
	if operand.kind ~= "table" then
		return cunknown()
	end
	local index = self:evaluate_unit(node.index)
	if (index.kind == "unknown") or (index.kind == "table") then
		return cunknown()
	end

	local data: index<TableConstant, "data"> = operand.data

	local result: Constant?

	local record = data.record
	if record then
		result = record[index.data]
	end

	local metatable = not result and data.metatable
	if metatable then
		local indexer = metatable.__index
		if indexer then
			result = indexer(operand, index)
		end
	end

	return result or cunknown()
end

function analyzer.evaluate_call(self: Identity, node: luau.AstExprCall): ...Constant
	local node_func = node.func
	local func = self:evaluate_unit(node_func)
	if func.kind ~= "table" then
		return cunknown()
	end

	local data: index<TableConstant, "data"> = func.data

	local metatable = data.metatable
	if not metatable then
		return cunknown()
	end

	local call = metatable.__call
	if not call then
		return cunknown()
	end

	local node_args = node.arguments

	local args_count = #node_args
	local args = table.create(args_count) :: { Constant }

	for index = 1, args_count do
		args[index] = self:evaluate(node_args[index].node)
	end

	return call(node.self, table.unpack(args))
end

local function CACHE(self: Identity, node: luau.AstExpr, output: Constant): Constant
	self.cache[node] = output
	return output
end

function analyzer.evaluate(self: Identity, node: luau.AstExpr): Constant
	do
		local cached = self.cache[node]
		if cached then
			return cached
		end
	end

	if node.tag == "boolean" then
		return CACHE(self, node, cbool(node.value))
	elseif node.tag == "number" then
		return CACHE(self, node, cnumber(node.value))
	elseif node.tag == "nil" then
		return CACHE(self, node, cnil())
	elseif node.tag == "string" then
		return CACHE(self, node, cstring(node.text))
	elseif node.tag == "unary" then
		return CACHE(self, node, self:unary(node.operator.text, self:evaluate_unit(node.operand)))
	elseif node.tag == "binary" then
		return CACHE(
			self,
			node,
			self:binary(
				self:evaluate_unit(node.lhsoperand),
				-- Lute
				node.operator.text :: any,
				self:evaluate_unit(node.rhsoperand)
			)
		)
		-- TODO: Should groups cancel evaluation?
	elseif node.tag == "group" then
		return CACHE(self, node, self:evaluate_unit(node.expression))
	elseif node.tag == "cast" then
		return CACHE(self, node, self:evaluate_unit(node.operand))
	elseif node.tag == "conditional" then
		return CACHE(self, node, self:evaluate_conditional(node))
	elseif node.tag == "interpolatedstring" then
		return CACHE(self, node, self:evaluate_interpstring(node))
	elseif node.tag == "indexname" then
		return CACHE(self, node, self:evaluate_indexname(node))
	elseif node.tag == "index" then
		return CACHE(self, node, self:evaluate_indexexpr(node))
	elseif node.tag == "local" then
		local tracked = self.tracked.locals[node["local"]]
		if tracked then
			local init = tracked.init
			if init then
				return CACHE(self, node, init)
			end
			if tracked.writes == 1 then
				local init_expr = tracked.init_expr
				if init_expr then
					local evaluated = self:evaluate(init_expr)
					tracked.init = evaluated
					return CACHE(self, node, evaluated)
				end
			end
		end
	elseif node.tag == "global" then
		local tracked = self.tracked.globals[node.name.text]
		if tracked then
			local init = tracked.init
			if init then
				return CACHE(self, node, init)
			end
			if tracked.writes == 1 then
				local init_expr = tracked.init_expr
				if init_expr then
					local evaluated = self:evaluate(init_expr)
					tracked.init = evaluated
					return CACHE(self, node, evaluated)
				end
			end
		end
	end

	-- Luau
	return CACHE(self, (node :: any) :: luau.AstExpr, cunknown())
end

function analyzer.evaluate_unit(self: Identity, node: luau.AstExpr): Constant
	if node.tag == "call" then
		local out = { self:evaluate_call(node) }
		if #out ~= 1 then
			return cunknown()
		end
		return out[1]
	else
		return self:evaluate(node)
	end
end

local function value_to_expression(value: ConstantData): luau.AstExpr
	if type(value) == "boolean" then
		local output: luau.AstExprConstantBool = {
			leadingtrivia = {},
			position = {
				line = 0,
				column = 0,
			},
			text = if value then "true" else "false",
			value = value,
			trailingtrivia = {},
			tag = "boolean",
		}
		return output
	elseif type(value) == "number" then
		local output: luau.AstExprConstantNumber = {
			leadingtrivia = {},
			position = {
				line = 0,
				column = 0,
			},
			text = tostring(value),
			value = value,
			trailingtrivia = {},
			tag = "number",
		}
		return output
	elseif type(value) == "vector" then
		local z = value.z
		--[[
		im lazy xD
		]]
		if z == 0 then
			return luau.parseexpr(`vector.create({value.x}, {value.y})`)
		else
			return luau.parseexpr(`vector.create({value.x}, {value.y}, {value.z})`)
		end
	elseif type(value) == "string" then
		local output: luau.AstExprConstantString = {
			leadingtrivia = {},
			position = {
				line = 0,
				column = 0,
			},
			text = value,
			trailingtrivia = {},
			tag = "string",
			quotestyle = "single",
			blockdepth = 0,
		}
		return output
	elseif type(value) == "nil" then
		local output: luau.AstExprConstantNil = {
			leadingtrivia = {},
			position = {
				line = 0,
				column = 0,
			},
			text = "nil",
			trailingtrivia = {},
			tag = "nil",
		}
		return output
	else
		error(`cannot convert value to expression - unknown type '{type(value)}'`)
	end
end

local function constant_to_expression(value: KnownConstant): luau.AstExpr
	if value.kind == "unknown" then
		error("Cannot convert constant to expression; Constant was 'unknown' type.")
	elseif value.kind == "table" then
		local evaluator = ASSERT(
			value.data.evaluate,
			"Cannot convert constant to expression;\z
			No evaluator listed for table constant"
		)
		return evaluator()
	else
		return value_to_expression(value.data)
	end
end

return {
	create = constructor,
	create_from_tracked = constructor_from_tracked,
	value_to_expression = value_to_expression,
	constant_to_expression = constant_to_expression,

	truthy_eh = truthy_eh,

	cunknown = cunknown,
	cbool = cbool,
	cnumber = cnumber,
	cvector = cvector,
	cstring = cstring,
	cnil = cnil,

	track_variables = var_tracking.track,
}


local defs = require("./defs")
local luau = require("@lute/luau")
local visitors = require("@util/syntax/visitors")

type Variable = defs.Variable
type Tracked = defs.Tracked

type Scope = { [string]: Variable? }

type VisitorState = {
	tracked: Tracked,
	scopes: { Scope },
}

-- Luau
local function INSERT<Item>(tbl: { Item }, value: Item): ()
	table.insert(tbl, value)
end

local function ASSERT<Value>(value: Value, message: string?)
	if value then
		return value
	end
	error(message or "Assertion Failed!", 2)
end

local function ensure_global(state: VisitorState, name: string): Variable
	local globals = state.tracked.globals
	do
		local variable = globals[name]
		if variable then
			return variable
		end
	end
	local variable: Variable = {
		global = true,
		writes = 0,
	}
	globals[name] = variable
	return variable
end

local visitor = visitors.add_scope(visitors.base() :: visitors.Base<VisitorState>)

function visitor.scope_open()
	local state = visitor.state
	INSERT(state.scopes, {} :: Scope)
end
function visitor.scope_close()
	local state = visitor.state
	ASSERT(
		table.remove(state.scopes),
		`Internal Error;\z
		Invalid scoping - Value tracker closed more scopes than it opened.`
	)
end

local function scope_insert_local(identifier: luau.AstLocal, assignment: luau.AstExpr?): ()
	local state = visitor.state
	local scopes = state.scopes
	local scope = ASSERT(
		scopes[#scopes],
		`Internal Error;\z
		Invalid scoping - Value tracker could not find a current scope while inserting a local.`
	)

	local var: Variable = {
		global = false,
		writes = 1,
		init_expr = assignment,
	}

	state.tracked.locals[identifier] = var
	scope[identifier.name.text] = var
end

visitor.scope_insert_local = scope_insert_local

function visitor.scope_insert_local_function(node): ()
	scope_insert_local(node.name)
end

function visitor.vassign(node)
	local state = visitor.state
	local scopes = state.scopes
	local scopes_count = #scopes
	for index, assignee_pair in node.variables do
		local assignee = assignee_pair.node
		if assignee.tag == "local" then
			local name = assignee.token.text
			local variable
			for scopes_index = scopes_count, 1, -1 do
				variable = scopes[scopes_index][name]
				if variable then
					break
				end
			end

			if variable then
				variable.writes += 1
			else
				error(`Internal Error;\z
					Value tracker tried to reassign local '{name}',\z
					but a value for '{name}' did not exist given the\z
					current scope.`)
			end
		elseif assignee.tag == "global" then
			ensure_global(state, assignee.name.text).writes += 1
		end
	end

	return true
end

function visitor.vglobal(node)
	local state = visitor.state

	state.tracked.globals[node.name.text] = ensure_global(state, node.name.text)

	return false
end

local function local_tracked_variable(node: luau.AstLocal): Variable
	local scopes = visitor.state.scopes
	local scopes_count = #scopes

	local name = node.name.text

	local variable
	for scopes_index = scopes_count, 1, -1 do
		variable = scopes[scopes_index][name]
		if variable then
			break
		end
	end

	if variable then
		return variable
	else
		error(`Internal Error;\z
			Value tracker tried to reference local '{name}',\z
			but a value for '{name}' did not exist given the\z
			current scope.`)
	end
end

function visitor.vlocal_reference(node)
	local inner = node["local"]
	visitor.state.tracked.locals[inner] = local_tracked_variable(inner)

	return true
end

local function track_values(root: luau.AstStatBlock): Tracked
	local tracked: Tracked = {
		globals = {},
		locals = {},
		variables = {},
	}
	visitor:visit_ast_block(root, {
		tracked = tracked,
		scopes = {},
	})
	return tracked
end

return {
	track = track_values,
}


--!strict
local luau = require("@lute/luau")

export type Ast = luau.AstStatBlock
export type AstExpr = luau.AstExpr
export type Parsley = {
	root: Ast,
	eof: luau.Eof,
}

local function parse(source: string): Ast
	return luau.parse(source).root
end

local function parseexpr(source: string): AstExpr
	return luau.parseexpr(source)
end

local function parsefile(source: string): Parsley
	local result = luau.parse(source)
	return { root = result.root, eof = result.eof }
end

return {
	parse = parse,
	parseexpr = parseexpr,
	parsefile = parsefile,
}


local luau = require("@lute/luau")
local visitors = require("@util/syntax/visitors")

local first_last_visitor = visitors.create() :: visitors.Create<{
	first: luau.Token,
	last: luau.Token,
}?>

function first_last_visitor.vtoken(node): ()
	local state = first_last_visitor.state
	if state then
		state.last = node
	else
		state = {
			first = node,
			last = node,
		}
		first_last_visitor.state = state
	end
end

local function from_expr(node: luau.AstExpr)
	return first_last_visitor:visit_ast_expression(node)
end

local function from_block(node: luau.AstStatBlock)
	return first_last_visitor:visit_ast_block(node)
end

local function from_stat(node: luau.AstStat)
	return first_last_visitor:visit_ast_statement(node)
end

local function from_type(node: luau.AstType)
	return first_last_visitor:visit_ast_type(node)
end

local function from_type_pack(node: luau.AstTypePack)
	return first_last_visitor:visit_ast_type_pack(node)
end

local function first_from_expr(node)
	local out = from_expr(node)
	return out and out.first
end
local function first_from_block(node)
	local out = from_block(node)
	return out and out.first
end
local function first_from_stat(node)
	local out = from_stat(node)
	return out and out.first
end
local function first_from_type(node)
	local out = from_type(node)
	return out and out.first
end
local function first_from_type_pack(node)
	local out = from_type_pack(node)
	return out and out.first
end
local function last_from_expr(node)
	local out = from_expr(node)
	return out and out.last
end
local function last_from_block(node)
	local out = from_block(node)
	return out and out.last
end
local function last_from_stat(node)
	local out = from_stat(node)
	return out and out.last
end
local function last_from_type(node)
	local out = from_type(node)
	return out and out.last
end
local function last_from_type_pack(node)
	local out = from_type_pack(node)
	return out and out.last
end

return {
	firstlast_from_expr = from_expr,
	firstlast_from_block = from_block,
	firstlast_from_stat = from_stat,
	firstlast_from_type = from_type,
	firstlast_from_type_pack = from_type_pack,
	first_from_expr = first_from_expr,
	first_from_block = first_from_block,
	first_from_stat = first_from_stat,
	first_from_type = first_from_type,
	first_from_type_pack = first_from_type_pack,
	last_from_expr = last_from_expr,
	last_from_block = last_from_block,
	last_from_stat = last_from_stat,
	last_from_type = last_from_type,
	last_from_type_pack = last_from_type_pack,
}


local luau = require("@lute/luau")
local visitors = require("@util/syntax/visitors")

type VisitorState = {
	leading: { luau.Trivia },
	trailing: { luau.Trivia },
}
local visitor = visitors.create() :: visitors.Create<VisitorState>

local function add_token(token: luau.Token<string>): ()
	local leading = token.leadingtrivia
	local trailing = token.trailingtrivia

	local state = visitor.state
	local state_leading = state.leading
	local state_trailing = state.trailing
	table.move(leading, 1, #leading, #state_leading + 1, state_leading)
	table.move(trailing, 1, #trailing, #state_trailing + 1, state_trailing)
end

function visitor.vtoken(node)
	add_token(node)
end

local function steal_from_token<T>(token: luau.Token<T>, trivia: VisitorState?): VisitorState
	trivia = trivia or { leading = {}, trailing = {} }
	visitor:visit_ast_token(token, trivia)
	return trivia
end

local function steal_from_expr(expr: luau.AstExpr, trivia: VisitorState?): VisitorState
	trivia = trivia or { leading = {}, trailing = {} }
	visitor:visit_ast_expression(expr, trivia)
	return trivia
end

local function steal_from_statement(stat: luau.AstStat, trivia: VisitorState?): VisitorState
	trivia = trivia or { leading = {}, trailing = {} }
	visitor:visit_ast_statement(stat, trivia)
	return trivia
end

local function push(trivia: { luau.Trivia }, onto: { luau.Trivia }): { luau.Trivia }
	table.move(onto, 1, #onto, #trivia + 1, onto)
	table.move(trivia, 1, #trivia, 1, onto)
	return onto
end

local function append(trivia: { luau.Trivia }, onto: { luau.Trivia }): { luau.Trivia }
	table.move(trivia, 1, #trivia, #onto + 1, onto)
	return onto
end

local function push_lead<T>(trivia: { luau.Trivia }, token)
	return push(trivia, token.leadingtrivia)
end

local function append_lead<T>(trivia: { luau.Trivia }, token)
	return append(trivia, token.leadingtrivia)
end

local function push_trail<T>(trivia: { luau.Trivia }, token)
	return push(trivia, token.trailingtrivia)
end

local function append_trail<T>(trivia: { luau.Trivia }, token)
	return append(trivia, token.trailingtrivia)
end

return {
	steal_from_expr = steal_from_expr,
	steal_from_stat = steal_from_statement,
	steal_from_token = steal_from_token,

	mutate = {
		push = push,
		append = append,
		push_lead = push_lead,
		append_lead = append_lead,
		push_trail = push_trail,
		append_trail = append_trail,
	},
}


local luau = require("@lute/luau")

export type Base<State = unknown> = {
	state: State,

	read visit_ast_block: (self: Base<State>, node: luau.AstStatBlock, state: State) -> State,
	read visit_ast_statement: (self: Base<State>, node: luau.AstStat, state: State) -> State,
	read visit_ast_expression: (self: Base<State>, node: luau.AstExpr, state: State) -> State,
	read visit_ast_type: (self: Base<State>, node: luau.AstType, state: State) -> State,
	read visit_ast_type_pack: (self: Base<State>, node: luau.AstTypePack, state: State) -> State,
	read visit_ast_token: <Kind>(self: Base<State>, node: luau.Token<Kind>, state: State) -> State,

	vtoken: (luau.Token) -> (),
	vlocal: (luau.AstLocal) -> boolean,
	vfunction_body: (luau.AstFunctionBody) -> boolean,

	vexpression: (luau.AstExpr) -> boolean,
	eexpression: (luau.AstExpr) -> (),
	vstatement: (luau.AstStat) -> boolean,
	estatement: (luau.AstStat) -> (),
	vtype: (luau.AstType) -> boolean,
	etype: (luau.AstType) -> (),
	vtype_pack: (luau.AstTypePack) -> boolean,
	etype_pack: (luau.AstTypePack) -> (),

	vblock: (luau.AstStatBlock) -> boolean,
	vassign: (luau.AstStatAssign) -> boolean,
	vcompound_assign: (luau.AstStatCompoundAssign) -> boolean,
	vfor_in: (luau.AstStatForIn) -> boolean,
	vfor: (luau.AstStatFor) -> boolean,
	vfunction: (luau.AstStatFunction) -> boolean,
	vif: (luau.AstStatIf) -> boolean,
	vlocal_declaration: (luau.AstStatLocal) -> boolean,
	vlocal_function: (luau.AstStatLocalFunction) -> boolean,
	vrepeat: (luau.AstStatRepeat) -> boolean,
	vreturn: (luau.AstStatReturn) -> boolean,
	vtype_alias: (luau.AstStatTypeAlias) -> boolean,
	vtype_function: (luau.AstStatTypeFunction) -> boolean,
	vwhile: (luau.AstStatWhile) -> boolean,

	vboolean: (luau.AstExprConstantBool) -> boolean,
	vnil: (luau.AstExprConstantNil) -> boolean,
	vnumber: (luau.AstExprConstantNumber) -> boolean,
	vstring: (luau.AstExprConstantString) -> boolean,
	vvarargs: (luau.AstExprVarargs) -> boolean,
	vanonymous_function: (luau.AstExprAnonymousFunction) -> boolean,
	vbinary: (luau.AstExprBinary) -> boolean,
	vcall: (luau.AstExprCall) -> boolean,
	vglobal: (luau.AstExprGlobal) -> boolean,
	vgroup: (luau.AstExprGroup) -> boolean,
	vif_expression: (luau.AstExprIfElse) -> boolean,
	vindex_expr: (luau.AstExprIndexExpr) -> boolean,
	vindex_name: (luau.AstExprIndexName) -> boolean,
	vinterpolated_string: (luau.AstExprInterpString) -> boolean,
	vlocal_reference: (luau.AstExprLocal) -> boolean,
	vtable_item: (luau.AstExprTableItem) -> boolean,
	vtable: (luau.AstExprTable) -> boolean,
	vtype_assertion: (luau.AstExprTypeAssertion) -> boolean,
	vunary: (luau.AstExprUnary) -> boolean,
}

export type AddType = {
	read contains_type: true,
	type_varray: (luau.AstTypeArray) -> boolean,
	type_vboolean: (luau.AstTypeSingletonBool) -> boolean,
	type_vfunction: (luau.AstTypeFunction) -> boolean,
	type_vgroup: (luau.AstTypeGroup) -> boolean,
	type_vintersection: (luau.AstTypeIntersection) -> boolean,
	type_vreference: (luau.AstTypeReference) -> boolean,
	type_vstring: (luau.AstTypeSingletonString) -> boolean,
	type_vtable: (luau.AstTypeTable) -> boolean,
	type_vtypeof: (luau.AstTypeTypeof) -> boolean,
	type_vunion: (luau.AstTypeUnion) -> boolean,
	type_vexplicitpack: (luau.AstTypePackExplicit) -> boolean,
	type_vgenericpack: (luau.AstTypePackGeneric) -> boolean,
	type_vvariadicpack: (luau.AstTypePackVariadic) -> boolean,
}

--[=[
at the moment this API can lead to certain inconsistencies w.r.t types, not going to deal with it for now.
]=]
export type AddScope = {
	read contains_scope: true,
	-- called when a scope is opened
	scope_open: () -> (),
	-- called when a scope is closed
	scope_close: () -> (),

	-- insert a local into the current scope
	scope_insert_local: (identifier: luau.AstLocal, assignment: luau.AstExpr?) -> (),
	-- insert a local function into the current scope
	scope_insert_local_function: (fun: luau.AstStatLocalFunction) -> (),
}

export type AddProcess = {
	read contains_process: true,
	rexpression: (luau.AstExpr) -> boolean,
	rstatement: (luau.AstStat) -> boolean,
	rtype: (luau.AstType) -> boolean,
	rtype_pack: (luau.AstTypePack) -> boolean,
	replace_value: (luau.AstExpr) -> luau.AstExpr?,
}

type IdentityAny<State = unknown> =
	& Base<State>
	& (AddType | { contains_type: nil })
	& (AddScope | { contains_scope: nil })
	& (AddProcess | { contains_process: nil })

local function EXHAUSTIVE_MATCH<Msg>(value: never, msg: Msg?): never
	error(msg or `Unknown value in exhaustive match: {value}`)
end

local v: IdentityAny

local function visit_token<T>(token: luau.Token<T>)
	v.vtoken(token)
end

local function visit_punctuated<T, Separator>(list: luau.Punctuated<T, Separator>, apply: (T) -> ())
	for index, item in list do
		apply(item.node)
		if item.separator then
			visit_token(item.separator)
		end
	end
end

local function visit_local(node: luau.AstLocal)
	if not v.vlocal(node) then
		return
	end
	visit_token(node.name)
	if node.colon then
		visit_token(node.colon)
	end
	if node.annotation then
		visit_type(node.annotation)
	end
end

local vscope: AddScope?

local function visit_block(node: luau.AstStatBlock): ()
	if not v.vblock(node) then
		return
	end
	if vscope then
		vscope.scope_open()
		for _, statement in node.statements do
			visit_statement(statement)
		end
		vscope.scope_close()
	else
		for _, statement in node.statements do
			visit_statement(statement)
		end
		return
	end
end

local function visit_if(node: luau.AstStatIf)
	if not v.vif(node) then
		return
	end
	visit_token(node.ifkeyword)
	visit_expression(node.condition)
	visit_token(node.thenkeyword)
	visit_block(node.consequent)
	for _, elseif_node in node.elseifs do
		visit_token(elseif_node.elseifkeyword)
		visit_expression(elseif_node.condition)
		visit_token(elseif_node.thenkeyword)
		visit_block(elseif_node.consequent)
	end
	if node.elsekeyword then
		visit_token(node.elsekeyword)
	end
	if node.antecedent then
		visit_block(node.antecedent)
	end
	visit_token(node.endkeyword)
end

local function visit_while(node: luau.AstStatWhile)
	if not v.vwhile(node) then
		return
	end
	visit_token(node.whilekeyword)
	visit_expression(node.condition)
	visit_token(node.dokeyword)
	visit_block(node.body)
	visit_token(node.endkeyword)
end

local function visit_repeat(node: luau.AstStatRepeat)
	if not v.vrepeat(node) then
		return
	end
	visit_token(node.repeatKeyword)
	visit_block(node.body)
	visit_token(node.untilKeyword)
	visit_expression(node.condition)
end

local function visit_return(node: luau.AstStatReturn)
	if not v.vreturn(node) then
		return
	end
	visit_token(node.returnkeyword)
	visit_punctuated(node.expressions, visit_expression)
end

local function visit_local_statement(node: luau.AstStatLocal)
	if not v.vlocal_declaration(node) then
		return
	end
	visit_token(node.localkeyword)
	visit_punctuated(node.variables, visit_local)
	if node.equals then
		visit_token(node.equals)
	end
	visit_punctuated(node.values, visit_expression)
	if vscope then
		for index, variable in node.variables do
			local value_pair = node.values[index]
			local value = value_pair and value_pair.node
			vscope.scope_insert_local(variable.node, value)
		end
	end
end

local function visit_for(node: luau.AstStatFor)
	if not v.vfor(node) then
		return
	end
	visit_token(node.forkeyword)
	visit_local(node.variable)
	visit_token(node.equals)
	visit_expression(node.from)
	visit_token(node.tocomma)
	visit_expression(node.to)
	if node.stepcomma then
		visit_token(node.stepcomma)
	end
	if node.step then
		visit_expression(node.step)
	end
	visit_token(node.dokeyword)
	if not vscope then
		visit_block(node.body)
	else
		if v.vblock(node.body) then
			vscope.scope_open()
			vscope.scope_insert_local(node.variable)
			for _, statement in node.body.statements do
				visit_statement(statement)
			end
			vscope.scope_close()
		end
	end
	visit_token(node.endkeyword)
end

local function visit_for_in(node: luau.AstStatForIn)
	if not v.vfor_in(node) then
		return
	end
	visit_token(node.forkeyword)
	visit_punctuated(node.variables, visit_local)
	visit_token(node.inkeyword)
	visit_punctuated(node.values, visit_expression)
	visit_token(node.dokeyword)
	if not vscope then
		visit_block(node.body)
	else
		if v.vblock(node.body) then
			vscope.scope_open()
			for _, variable in node.variables do
				vscope.scope_insert_local(variable.node)
			end
			for _, statement in node.body.statements do
				visit_statement(statement)
			end
			vscope.scope_close()
		end
	end
	visit_token(node.endkeyword)
end

local function visit_assign(node: luau.AstStatAssign)
	if not v.vassign(node) then
		return
	end
	visit_punctuated(node.variables, visit_expression)
	visit_token(node.equals)
	visit_punctuated(node.values, visit_expression)
end

local function visit_compound_assign(node: luau.AstStatCompoundAssign)
	if not v.vcompound_assign(node) then
		return
	end
	visit_expression(node.variable)
	visit_token(node.operand)
	visit_expression(node.value)
end

local vtype: AddType?

local function visit_generic(node: luau.AstGenericType)
	visit_token(node.name)
	if node.equals then
		visit_token(node.equals)
	end
	if node.default then
		visit_type(node.default)
	end
end

local function visit_generic_pack(node: luau.AstGenericTypePack)
	visit_token(node.name)
	visit_token(node.ellipsis)
	if node.equals then
		visit_token(node.equals)
	end
	if node.default then
		visit_type_pack(node.default)
	end
end

local function visit_type_alias(node: luau.AstStatTypeAlias)
	if not v.vtype_alias(node) then
		return
	end
	if node.export then
		visit_token(node.export)
	end
	visit_token(node.typeToken)
	visit_token(node.name)
	if node.opengenerics then
		visit_token(node.opengenerics)
	end
	if node.generics then
		visit_punctuated(node.generics, visit_generic)
	end
	if node.genericpacks then
		visit_punctuated(node.genericpacks, visit_generic_pack)
	end
	if node.closegenerics then
		visit_token(node.closegenerics)
	end
	visit_token(node.equals)
	visit_type(node.type)
end

local function visit_string(node: luau.AstExprConstantString)
	if not v.vstring(node) then
		return
	end
	visit_token(node)
end

local function visit_nil(node: luau.AstExprConstantNil)
	if not v.vnil(node) then
		return
	end
	visit_token(node)
end

local function visit_boolean(node: luau.AstExprConstantBool)
	if not v.vboolean(node) then
		return
	end
	visit_token(node)
end

local function visit_number(node: luau.AstExprConstantNumber)
	if not v.vnumber(node) then
		return
	end
	visit_token(node)
end

local function visit_varargs(node: luau.AstExprVarargs)
	if not v.vvarargs(node) then
		return
	end
	visit_token(node)
end

local function visit_local_reference(node: luau.AstExprLocal)
	if not v.vlocal_reference(node) then
		return
	end
	visit_token(node.token)
end

local function visit_global(node: luau.AstExprGlobal)
	if not v.vglobal(node) then
		return
	end
	visit_token(node.name)
end

local function visit_call(node: luau.AstExprCall)
	if not v.vcall(node) then
		return
	end
	visit_expression(node.func)
	if node.openparens then
		visit_token(node.openparens)
	end
	visit_punctuated(node.arguments, visit_expression)
	if node.closeparens then
		visit_token(node.closeparens)
	end
end

local function visit_unary(node: luau.AstExprUnary)
	if not v.vunary(node) then
		return
	end
	visit_token(node.operator)
	visit_expression(node.operand)
end

local function visit_binary(node: luau.AstExprBinary)
	if not v.vbinary(node) then
		return
	end
	visit_expression(node.lhsoperand)
	visit_token(node.operator)
	visit_expression(node.rhsoperand)
end

local function visit_function_body(node: luau.AstFunctionBody, local_fun: luau.AstStatLocalFunction?)
	if not v.vfunction_body(node) then
		-- need to make sure the local fun gets inserted regardless of if the body is visited
		if vscope and local_fun then
			vscope.scope_insert_local_function(local_fun)
		end
		return
	end
	if node.opengenerics then
		visit_token(node.opengenerics)
	end
	if node.generics then
		visit_punctuated(node.generics, visit_generic)
	end
	if node.genericpacks then
		visit_punctuated(node.genericpacks, visit_generic_pack)
	end
	if node.closegenerics then
		visit_token(node.closegenerics)
	end
	visit_token(node.openparens)
	local parameters = node.parameters
	visit_punctuated(parameters, visit_local)
	if node.vararg then
		visit_token(node.vararg)
	end
	if node.varargcolon then
		visit_token(node.varargcolon)
	end
	if node.varargannotation then
		visit_type_pack(node.varargannotation)
	end
	visit_token(node.closeparens)
	if node.returnspecifier then
		visit_token(node.returnspecifier)
	end
	if node.returnannotation then
		visit_type_pack(node.returnannotation)
	end

	if not vscope then
		visit_block(node.body)
	else
		if local_fun then
			vscope.scope_insert_local_function(local_fun)
		end
		if v.vblock(node.body) then
			vscope.scope_open()
			local self = node.self
			if self then
				vscope.scope_insert_local(self)
			end
			for _, parameter in node.parameters do
				vscope.scope_insert_local(parameter.node)
			end
			for _, statement in node.body.statements do
				visit_statement(statement)
			end
			vscope.scope_close()
		end
	end
	visit_token(node.endkeyword)
end

local function visit_attribute(node: luau.AstAttribute)
	visit_token(node)
end

local function visit_anonymous_function(node: luau.AstExprAnonymousFunction)
	if not v.vanonymous_function(node) then
		return
	end
	for _, attribute in node.attributes do
		visit_attribute(attribute)
	end
	visit_token(node.functionkeyword)
	visit_function_body(node.body)
end

local function visit_function(node: luau.AstStatFunction)
	if not v.vfunction(node) then
		return
	end
	for _, attribute in node.attributes do
		visit_attribute(attribute)
	end
	visit_token(node.functionkeyword)
	visit_expression(node.name)
	visit_function_body(node.body)
end

local function visit_local_function(node: luau.AstStatLocalFunction)
	if not v.vlocal_function(node) then
		return
	end
	for _, attribute in node.attributes do
		visit_attribute(attribute)
	end
	visit_token(node.localkeyword)
	visit_token(node.functionkeyword)
	visit_local(node.name)
	if vscope then
		vscope.scope_insert_local_function(node)
	end
	visit_function_body(node.body)
end

local function visit_stat_type_function(node: luau.AstStatTypeFunction)
	if not v.vtype_function(node) then
		return
	end
	if node.export then
		visit_token(node.export)
	end
	visit_token(node.type)
	visit_token(node.functionkeyword)
	visit_token(node.name)
	visit_function_body(node.body)
end

local function visit_table_item(node: luau.AstExprTableItem)
	if not v.vtable_item(node) then
		return
	end
	if node.kind == "list" then
		visit_expression(node.value)
	elseif node.kind == "record" then
		visit_token(node.key)
		visit_token(node.equals)
		visit_expression(node.value)
	elseif node.kind == "general" then
		visit_token(node.indexeropen)
		visit_expression(node.key)
		visit_token(node.indexerclose)
		visit_token(node.equals)
		visit_expression(node.value)
	else
		EXHAUSTIVE_MATCH(node.kind)
	end

	if node.separator then
		visit_token(node.separator)
	end
end

local function visit_table(node: luau.AstExprTable)
	if not v.vtable(node) then
		return
	end
	visit_token(node.openbrace)
	for _, item in node.entries do
		visit_table_item(item)
	end
	visit_token(node.closebrace)
end

local function visit_index_name(node: luau.AstExprIndexName)
	if not v.vindex_name(node) then
		return
	end
	visit_expression(node.expression)
	visit_token(node.accessor)
	visit_token(node.index)
end

local function visit_index_expr(node: luau.AstExprIndexExpr)
	if not v.vindex_expr(node) then
		return
	end
	visit_expression(node.expression)
	visit_token(node.openbrackets)
	visit_expression(node.index)
	visit_token(node.closebrackets)
end

local function visit_group(node: luau.AstExprGroup)
	if not v.vgroup(node) then
		return
	end
	visit_token(node.openparens)
	visit_expression(node.expression)
	visit_token(node.closeparens)
end

local function visit_interpolated_string(node: luau.AstExprInterpString)
	if not v.vinterpolated_string(node) then
		return
	end
	for i = 1, #node.strings do
		visit_token(node.strings[i])
		if i <= #node.expressions then
			visit_expression(node.expressions[i])
		end
	end
end

local function visit_type_assertion(node: luau.AstExprTypeAssertion)
	if not v.vtype_assertion(node) then
		return
	end
	visit_expression(node.operand)
	visit_token(node.operator)
	visit_type(node.annotation)
end

local function visit_if_expression(node: luau.AstExprIfElse)
	if not v.vif_expression(node) then
		return
	end
	visit_token(node.ifkeyword)
	visit_expression(node.condition)
	visit_token(node.thenkeyword)
	visit_expression(node.consequent)
	for _, elseifs in node.elseifs do
		visit_token(elseifs.elseifkeyword)
		visit_expression(elseifs.condition)
		visit_token(elseifs.thenkeyword)
		visit_expression(elseifs.consequent)
	end
	visit_token(node.elsekeyword)
	visit_expression(node.antecedent)
end

local function visit_type_or_pack(node: luau.AstType | luau.AstTypePack)
	if node.tag == "explicit" or node.tag == "generic" or node.tag == "variadic" then
		visit_type_pack(node)
	else
		visit_type(node)
	end
end

local function visit_type_reference(node: luau.AstTypeReference)
	if vtype and vtype.type_vreference(node) then
		if node.prefix then
			visit_token(node.prefix)
		end
		if node.prefixpoint then
			visit_token(node.prefixpoint)
		end
		visit_token(node.name)
		if node.openparameters then
			visit_token(node.openparameters)
		end
		if node.parameters then
			visit_punctuated(node.parameters, visit_type_or_pack)
		end
		if node.closeparameters then
			visit_token(node.closeparameters)
		end
	end
end

local function visit_type_boolean(node: luau.AstTypeSingletonBool)
	if vtype and vtype.type_vboolean(node) then
		visit_token(node)
	end
end

local function visit_type_string(node: luau.AstTypeSingletonString)
	if vtype and vtype.type_vstring(node) then
		visit_token(node)
	end
end

local function visit_type_typeof(node: luau.AstTypeTypeof)
	if vtype and vtype.type_vtypeof(node) then
		visit_token(node.typeof)
		visit_token(node.openparens)
		visit_expression(node.expression)
		visit_token(node.closeparens)
	end
end

local function visit_type_group(node: luau.AstTypeGroup)
	if vtype and vtype.type_vgroup(node) then
		visit_token(node.openparens)
		visit_type(node.type)
		visit_token(node.closeparens)
	end
end

local function visit_type_union(node: luau.AstTypeUnion)
	if vtype and vtype.type_vunion(node) then
		if node.leading then
			visit_token(node.leading)
		end
		visit_punctuated(node.types, visit_type)
	end
end

local function visit_type_intersection(node: luau.AstTypeIntersection)
	if vtype and vtype.type_vintersection(node) then
		if node.leading then
			visit_token(node.leading)
		end
		visit_punctuated(node.types, visit_type)
	end
end

local function visit_type_array(node: luau.AstTypeArray)
	if vtype and vtype.type_varray(node) then
		visit_token(node.openbrace)
		if node.access then
			visit_token(node.access)
		end
		visit_type(node.type)
		visit_token(node.closebrace)
	end
end

local function visit_type_table(node: luau.AstTypeTable)
	if vtype and vtype.type_vtable(node) then
		visit_token(node.openbrace)
		for _, entry in node.entries do
			if entry.access then
				visit_token(entry.access)
			end
			if entry.kind == "indexer" then
				visit_token(entry.indexeropen)
				visit_type(entry.key)
				visit_token(entry.indexerclose)
			elseif entry.kind == "stringproperty" then
				visit_token(entry.indexeropen)
				visit_type_string(entry.key)
				visit_token(entry.indexerclose)
			else
				visit_token(entry.key)
			end
			visit_token(entry.colon)
			visit_type(entry.value)
			if entry.separator then
				visit_token(entry.separator)
			end
		end
		visit_token(node.closebrace)
	end
end

local function visit_type_function_parameter(node: luau.AstTypeFunctionParameter)
	if node.name then
		visit_token(node.name)
	end
	if node.colon then
		visit_token(node.colon)
	end
	visit_type(node.type)
end

local function visit_type_function(node: luau.AstTypeFunction)
	if vtype and vtype.type_vfunction(node) then
		if node.opengenerics then
			visit_token(node.opengenerics)
		end
		if node.generics then
			visit_punctuated(node.generics, visit_generic)
		end
		if node.genericpacks then
			visit_punctuated(node.genericpacks, visit_generic_pack)
		end
		if node.closegenerics then
			visit_token(node.closegenerics)
		end
		visit_token(node.openparens)
		visit_punctuated(node.parameters, visit_type_function_parameter)
		if node.vararg then
			visit_type_pack(node.vararg)
		end
		visit_token(node.closeparens)
		visit_token(node.returnarrow)
		visit_type_pack(node.returntypes)
	end
end

local function visit_type_pack_explicit(node: luau.AstTypePackExplicit)
	if vtype and vtype.type_vexplicitpack(node) then
		if node.openparens then
			visit_token(node.openparens)
		end
		visit_punctuated(node.types, visit_type)
		if node.tailtype then
			visit_type_pack(node.tailtype)
		end
		if node.closeparens then
			visit_token(node.closeparens)
		end
	end
end

local function visit_type_pack_generic(node: luau.AstTypePackGeneric)
	if vtype and vtype.type_vgenericpack(node) then
		visit_token(node.name)
		visit_token(node.ellipsis)
	end
end

local function visit_type_pack_variadic(node: luau.AstTypePackVariadic)
	if vtype and vtype.type_vvariadicpack(node) then
		if node.ellipsis then
			visit_token(node.ellipsis)
		end
		visit_type(node.type)
	end
end

function visit_expression(node: luau.AstExpr)
	if v.vexpression(node) then
		if node.tag == "nil" then
			visit_nil(node)
		elseif node.tag == "boolean" then
			visit_boolean(node)
		elseif node.tag == "number" then
			visit_number(node)
		elseif node.tag == "string" then
			visit_string(node)
		elseif node.tag == "local" then
			visit_local_reference(node)
		elseif node.tag == "global" then
			visit_global(node)
		elseif node.tag == "vararg" then
			visit_varargs(node)
		elseif node.tag == "call" then
			visit_call(node)
		elseif node.tag == "unary" then
			visit_unary(node)
		elseif node.tag == "binary" then
			visit_binary(node)
		elseif node.tag == "function" then
			visit_anonymous_function(node)
		elseif node.tag == "table" then
			visit_table(node)
		elseif node.tag == "indexname" then
			visit_index_name(node)
		elseif node.tag == "index" then
			visit_index_expr(node)
		elseif node.tag == "group" then
			visit_group(node)
		elseif node.tag == "interpolatedstring" then
			visit_interpolated_string(node)
		elseif node.tag == "cast" then
			visit_type_assertion(node)
		elseif node.tag == "conditional" then
			visit_if_expression(node)
		else
			EXHAUSTIVE_MATCH(node.tag)
		end
		v.eexpression(node)
	end
end

function visit_statement(node: luau.AstStat)
	if v.vstatement(node) then
		if node.tag == "block" then
			visit_block(node)
		elseif node.tag == "conditional" then
			visit_if(node)
		elseif node.tag == "expression" then
			visit_expression(node.expression)
		elseif node.tag == "local" then
			visit_local_statement(node)
		elseif node.tag == "return" then
			visit_return(node)
		elseif node.tag == "while" then
			visit_while(node)
		elseif node.tag == "break" then
			visit_token(node)
		elseif node.tag == "continue" then
			visit_token(node)
		elseif node.tag == "repeat" then
			visit_repeat(node)
		elseif node.tag == "for" then
			visit_for(node)
		elseif node.tag == "forin" then
			visit_for_in(node)
		elseif node.tag == "assign" then
			visit_assign(node)
		elseif node.tag == "compoundassign" then
			visit_compound_assign(node)
		elseif node.tag == "function" then
			visit_function(node)
		elseif node.tag == "localfunction" then
			visit_local_function(node)
		elseif node.tag == "typealias" then
			visit_type_alias(node)
		elseif node.tag == "typefunction" then
			visit_stat_type_function(node)
		else
			EXHAUSTIVE_MATCH(node.tag)
		end
		v.estatement(node)
	end
end

function visit_type(node: luau.AstType)
	if v.vtype(node) then
		if node.tag == "reference" then
			visit_type_reference(node)
		elseif node.tag == "boolean" then
			visit_type_boolean(node)
		elseif node.tag == "string" then
			visit_type_string(node)
		elseif node.tag == "typeof" then
			visit_type_typeof(node)
		elseif node.tag == "group" then
			visit_type_group(node)
		elseif node.tag == "union" then
			visit_type_union(node)
		elseif node.tag == "intersection" then
			visit_type_intersection(node)
		elseif node.tag == "optional" then
			visit_token(node)
		elseif node.tag == "array" then
			visit_type_array(node)
		elseif node.tag == "table" then
			visit_type_table(node)
		elseif node.tag == "function" then
			visit_type_function(node)
		else
			EXHAUSTIVE_MATCH(node.tag)
		end
		v.etype(node)
	end
end

function visit_type_pack(node: luau.AstTypePack)
	if v.vtype_pack(node) then
		if node.tag == "explicit" then
			visit_type_pack_explicit(node)
		elseif node.tag == "generic" then
			visit_type_pack_generic(node)
		elseif node.tag == "variadic" then
			visit_type_pack_variadic(node)
		else
			EXHAUSTIVE_MATCH(node.tag)
		end
		v.etype_pack(node)
	end
end

local function always_explore(any: any): boolean
	return true
end
local function noop(...: any): () end

-- this is so cursed
local function api_wrapper<Node, State>(fun: (node: Node) -> ()): (self: Base<State>, node: Node, state: State) -> State
	local function visit_ast(self: IdentityAny<State>, node: Node, state: State): State
		local oldv = v
		v = self
		local oldtype = vtype
		vtype = if self.contains_type then self else nil
		local oldscope = vscope
		vscope = if self.contains_scope then self else nil
		local began = self.state
		self.state = state
		fun(node)
		local stopped = self.state
		self.state = began
		v, vtype, vscope = oldv, oldtype, oldscope
		return stopped
	end
	-- :(
	return visit_ast :: any
end

local visit_ast_block = api_wrapper(visit_block)
local visit_ast_statement = api_wrapper(visit_statement)
local visit_ast_expression = api_wrapper(visit_expression)
local visit_ast_type = api_wrapper(visit_type)
local visit_ast_type_pack = api_wrapper(visit_type_pack)
local visit_ast_token = api_wrapper(visit_token)

local function visitor_base<State>(): Base<State>
	local created: Base<State> = {
		state = nil,

		visit_ast_block = visit_ast_block,
		visit_ast_statement = visit_ast_statement,
		visit_ast_expression = visit_ast_expression,
		visit_ast_type = visit_ast_type,
		visit_ast_type_pack = visit_ast_type_pack,
		visit_ast_token = visit_ast_token :: <Kind>(self: Base<State>, node: luau.Token<Kind>, state: State) -> State,

		vtoken = noop,

		vlocal = always_explore,
		vfunction_body = always_explore,

		vexpression = always_explore,
		eexpression = noop,
		vstatement = always_explore,
		estatement = noop,
		vtype = always_explore,
		etype = noop,
		vtype_pack = always_explore,
		etype_pack = noop,

		vblock = always_explore,
		vassign = always_explore,
		vcompound_assign = always_explore,
		vfor_in = always_explore,
		vfor = always_explore,
		vfunction = always_explore,
		vif = always_explore,
		vlocal_declaration = always_explore,
		vlocal_function = always_explore,
		vrepeat = always_explore,
		vreturn = always_explore,
		vtype_alias = always_explore,
		vtype_function = always_explore,
		vwhile = always_explore,

		vboolean = always_explore,
		vnil = always_explore,
		vnumber = always_explore,
		vstring = always_explore,
		vvarargs = always_explore,
		vanonymous_function = always_explore,
		vbinary = always_explore,
		vcall = always_explore,
		vglobal = always_explore,
		vgroup = always_explore,
		vif_expression = always_explore,
		vindex_expr = always_explore,
		vindex_name = always_explore,
		vinterpolated_string = always_explore,
		vlocal_reference = always_explore,
		vtable_item = always_explore,
		vtable = always_explore,
		vtype_assertion = always_explore,
		vunary = always_explore,

		scope = nil,
		type = nil,
		typepack = nil,
	}

	return created
end

local default_type: AddType = {
	contains_type = true,
	type_varray = always_explore,
	type_vboolean = always_explore,
	type_vfunction = always_explore,
	type_vgroup = always_explore,
	type_vintersection = always_explore,
	type_vreference = always_explore,
	type_vstring = always_explore,
	type_vtable = always_explore,
	type_vtypeof = always_explore,
	type_vunion = always_explore,

	type_vexplicitpack = always_explore,
	type_vgenericpack = always_explore,
	type_vvariadicpack = always_explore,
}
local default_scope: AddScope = {
	contains_scope = true,
	scope_open = noop,
	scope_close = noop,

	scope_insert_local = noop,
	scope_insert_local_function = noop,
}
local function add_type<Visitor>(visitor: Visitor): Visitor & AddType
	-- :(
	visitor = visitor :: any
	for key, value in default_type :: any do
		visitor[key] = value
	end
	return visitor
end
local function add_scope<Visitor>(visitor: Visitor): Visitor & AddScope
	-- :(
	visitor = visitor :: any
	for key, value in default_scope :: any do
		visitor[key] = value
	end
	return visitor
end

export type Create<State = unknown> = Base<State> & AddType

--[=[
Creates a visitor with type and typepack visiting.
]=]
local function visitor_create<State>(): Create<State>
	return add_type(visitor_base())
end

return {
	create = visitor_create,
	base = visitor_base,

	add_type = add_type,
	add_scope = add_scope,

	inner_visit_block = visit_block,
	inner_visit_statement = visit_statement,
	inner_visit_expression = visit_expression,
	inner_visit_type = visit_type,
	inner_visit_type_pack = visit_type_pack,
	inner_visit_token = visit_token,
	inner_visit_local = visit_local,
	inner_visit_punctuated = visit_punctuated,
}


--[[
  This script will:
  - Migrate the latest version of the DS2 data store to the migrated data store,
  - Delete all of the entries in the DS2 data store except for the latest version.

  This script migrates the latest version of the DS2 data store to the migrated data store,
  and keeps the latest version in the DS2 data store in case a creator wants to rollback.
]]
local DataStoreService = game:GetService("DataStoreService")
local MIGRATED_DS_SCOPE = "global" -- Change to a different scope as needed

-- Wait for sufficient data store budget
local function waitForBudgetAsync(requestType: Enum.DataStoreRequestType, minBudget: number): ()
    while true do
        local budget = DataStoreService:GetRequestBudgetForRequestType(requestType)
        if budget >= minBudget then
            return
        end
        task.wait(0.5) -- Wait 0.5 second before checking again
    end
end

return function(dataStoreName)
    -- Parse the datastore name
    local ds2Key, playerIdString = string.match(dataStoreName, "^(.*)/([^/]*)$")

    if not ds2Key or not playerIdString then
        error("Invalid Data Store Name: " .. dataStoreName)
    end

    local playerId = tonumber(playerIdString)

    -- Initialize data stores
    local ds2SDS = DataStoreService:GetDataStore(dataStoreName)
    local ds2ODS = DataStoreService:GetOrderedDataStore(dataStoreName)
    local migratedDS = DataStoreService:GetDataStore(ds2Key, MIGRATED_DS_SCOPE)

    -- Store the first sorted page
    waitForBudgetAsync(Enum.DataStoreRequestType.GetSortedAsync, 1)
    local pages = ds2ODS:GetSortedAsync(false, 100)

    -------------
    -- MIGRATE --
    -------------

    -- Modified from BerezaaMethodDataMigrationTool
    -- See https://create.roblox.com/store/asset/100741523916104/BETA-Berezaa-Method-Data-Migration-Tool
    local migratedData = migratedDS:GetAsync(playerId)
    if migratedData == nil then
        local latestVersion = pages:GetCurrentPage()[1]
        if latestVersion then
            -- Get the data value from the SDS
            local version = latestVersion.value
            waitForBudgetAsync(Enum.DataStoreRequestType.GetAsync, 2)
            local data = ds2SDS:GetAsync(version)

            -- Move this data into the new SDS
            waitForBudgetAsync(Enum.DataStoreRequestType.SetIncrementAsync, 1)
            migratedDS:UpdateAsync(playerId, function(oldData)
                if oldData then
                    return oldData
                end
                return data
            end)
        end
        print("Info: Migrated latest version from " .. dataStoreName .. " to " .. ds2Key .. ".")
    else
        print("Info: Data already exists for player " .. playerId .. " in " .. ds2Key .. ".")
    end

    ------------
    -- DELETE --
    ------------

    -- Loop through all of the entries
    local firstPage = true
    while true do
        for i, key in ipairs(pages:GetCurrentPage()) do
            if i ~= 1 or not firstPage then
                -- Attempt deletion
                waitForBudgetAsync(Enum.DataStoreRequestType.SetIncrementAsync, 1)
                waitForBudgetAsync(Enum.DataStoreRequestType.SetIncrementSortedAsync, 1)
                ds2ODS:RemoveAsync(key.value)
                ds2SDS:RemoveAsync(key.value)
            end
        end

        if pages.IsFinished then
            break
        end

        waitForBudgetAsync(Enum.DataStoreRequestType.GetSortedAsync, 1)
        pages:AdvanceToNextPageAsync()
        firstPage = false
    end

    print("Info: Deleted all data from " .. dataStoreName .. " except for the latest version.")
end


--[[
  This script will:
  - Migrate the latest version of the DS2 data store to the migrated data store,
  - Delete all of the entries in the DS2 data store except for the latest version.

  This script migrates the latest version of the DS2 data store to the migrated data store,
  and keeps the latest version in the DS2 data store in case a creator wants to rollback.
]]
local DataStoreService = game:GetService("DataStoreService")
local MIGRATED_DS_SCOPE = "global" -- Change to a different scope as needed

-- Wait for sufficient data store budget
local function waitForBudgetAsync(requestType: Enum.DataStoreRequestType, minBudget: number): ()
    while true do
        local budget = DataStoreService:GetRequestBudgetForRequestType(requestType)
        if budget >= minBudget then
            return
        end
        task.wait(0.5) -- Wait 0.5 second before checking again
    end
end

return function(dataStoreName)
    -- Parse the datastore name
    local ds2Key, playerIdString = string.match(dataStoreName, "^(.*)/([^/]*)$")

    if not ds2Key or not playerIdString then
        error("Invalid Data Store Name: " .. dataStoreName)
    end

    local playerId = tonumber(playerIdString)

    -- Initialize data stores
    local ds2SDS = DataStoreService:GetDataStore(dataStoreName)
    local ds2ODS = DataStoreService:GetOrderedDataStore(dataStoreName)
    local migratedDS = DataStoreService:GetDataStore(ds2Key, MIGRATED_DS_SCOPE)

    -------------
    -- MIGRATE --
    -------------

    -- Modified from BerezaaMethodDataMigrationTool
    local migratedData = migratedDS:GetAsync(playerId)
    if migratedData == nil then
        waitForBudgetAsync(Enum.DataStoreRequestType.GetSortedAsync, 1)
        local pages = ds2ODS:GetSortedAsync(false, 1)

        -- Get the latest version
        local latestVersion = pages:GetCurrentPage()[1]
        if latestVersion then
            -- Get the data value from the SDS
            local version = latestVersion.value
            waitForBudgetAsync(Enum.DataStoreRequestType.GetAsync, 2)
            local data = ds2SDS:GetAsync(version)

            -- Move this data into the new SDS
            waitForBudgetAsync(Enum.DataStoreRequestType.SetIncrementAsync, 1)
            migratedDS:UpdateAsync(playerId, function(oldData)
                if oldData then
                    return oldData
                end
                return data
            end)
        end
        print("Info: Migrated latest version from " .. dataStoreName .. " to " .. ds2Key .. ".")
    else
        print("Info: Data already exists for player " .. playerId .. " in " .. ds2Key .. ".")
    end
end


--!nolint LocalShadow

local process = require("@lune/process")

local M = {}

type ArgKind = "POSITIONAL" | "FLAG" | "OPTION"
type ArgOptions = {
    help: string?,
    aliases: { string }?,
    default: string?, -- if this is not nil, the arg will be optional
}

-- this is what is stored, we guarante non nullity when args are added so this types makes Luau feel better
type ArgOptionsSafe = {
    help: string,
    aliases: { string },
    default: string?,
}

type ArgData = {
    name: string,
    kind: ArgKind,
    options: ArgOptionsSafe,
}

type ParseResult = {
    values: { [string]: string },
    flags: { [string]: boolean },
    fwd_args: { string }, -- all args after `--`
}

type SubcommandParseResult = {
    command: string,
    result: ParseResult,
}

local DEFAULT_OPTIONS: ArgOptionsSafe = {
    help = "",
    default = nil,
    aliases = {},
}

local DEFAULT_helpset = {
    ["-h"] = true,
    ["--help"] = true,
}

local function validate_key_or_error(key: string)
    if (key:sub(1, 2) == "--" or key:sub(1, 1) == "-") and not key:find(" ") then
        return
    end
    error(`arg key {key} is invalid. Keys must start with either '-' or '--' and may not contain spaces`)
end

local function validate_subcommand_name_or_error(name: string)
    if name:sub(1, 2) == "--" or name:sub(1, 1) == "-" or name:find(" ") then
        error(`subcommand name '{name}' is invalid. Name must not start with '-' or '--' and may not contain spaces`)
    end
end

type HelpSection = { title: string, lines: { { string } | string }? }
local function make_help(sections: { HelpSection }, indent: number?): string
    local function align_cols(rows: { { string } }, sep: string?): { { string } }
        local sep = if sep == nil then " " else sep
        local max_col_lengths: { number } = {}
        for _, row in rows do
            for i, s in row do
                if max_col_lengths[i] == nil or s:len() > max_col_lengths[i] then
                    max_col_lengths[i] = s:len()
                end
            end
        end
        local aligned_rows: { { string } } = {}
        for _, row in rows do
            local line: { string } = {}
            for i, col in row do
                table.insert(line, col)
                if i < #row then
                    local spacing = sep
                    local diff = max_col_lengths[i] - col:len()
                    if diff > 0 then
                        spacing = string.rep(" ", diff) .. sep
                    end
                    table.insert(line, spacing)
                end
            end
            table.insert(aligned_rows, line)
        end
        return aligned_rows
    end

    local function append_list<T>(dest: { T }, src: { T })
        for _, v in src do
            table.insert(dest, v)
        end
    end

    local ind = "  "
    if indent ~= nil then
        ind = string.rep(" ", indent)
    end

    local help_lines: { { string } } = {}

    for _, s in sections do
        table.insert(help_lines, { s.title })
        if s.lines ~= nil then
            local section_lines: { { string } } = {}
            for _, l in s.lines do
                if typeof(l) == "string" then
                    table.insert(section_lines, { ind, l })
                else
                    table.insert(section_lines, { ind, table.unpack(l) })
                end
            end
            append_list(help_lines, align_cols(section_lines))
        end
        table.insert(help_lines, {})
    end

    -- remove last empty line
    table.remove(help_lines)

    local help_text = ""
    for _, line in help_lines do
        help_text ..= table.concat(line, " ") .. "\n"
    end
    return help_text
end
-- subcommand cli, must be first positional arg
function M.new_subcommands(name, description: string?, helpkeys: { string }?)
    -- args that trigger print help + abort
    local helpset: { [string]: boolean } = DEFAULT_helpset
    if helpkeys ~= nil then
        helpset = {}
        for _, k in helpkeys do
            helpset[k] = true
        end
    end

    local cli = {}
    cli.name = name
    cli.description = description
    cli._helpkeys = helpkeys
    cli._helpset = helpset
    cli._subcommands = {}

    function cli:add_subcommand(name: string, description: string?)
        validate_subcommand_name_or_error(name)
        local sc = M.new(name, description, cli._helpkeys)
        cli._subcommands[name] = sc
        return sc
    end

    function cli:parse(args_in: { string }): (SubcommandParseResult?, string?)
        local args = {}
        for _, v in args_in do
            table.insert(args, v)
        end
        if #args == 0 then
            return nil, "insufficient arguments, required at least 1, got 0"
        end
        local command = args[1]
        assert(command ~= nil)

        -- asking for help?
        if cli._helpset[command] ~= nil then
            print(cli:help())
            process.exit(0)
        end

        local subcommand = cli._subcommands[command]
        if subcommand == nil then
            return nil, `'{command} is not a valid subcommand`
        end

        local sub_args = { table.unpack(args, 2) }
        local res, err = subcommand:parse(sub_args)
        if err ~= nil then
            return nil, err
        end
        assert(res ~= nil)

        return {
            command = command,
            result = res,
        }, nil
    end

    function cli:help(indent: number?): string
        local help_sections: { HelpSection } = {}

        table.insert(help_sections, { title = `usage: {cli.name} <command>` })

        if cli.description ~= nil then
            table.insert(help_sections, {
                title = "description:",
                lines = { { cli.description } },
            })
        end

        local subcommand_lines = {}
        for _, c in cli._subcommands do
            local desc = if c.description then c.description else ""
            table.insert(subcommand_lines, { c.name, desc })
        end

        table.insert(help_sections, {
            title = "commands:",
            lines = subcommand_lines,
        })

        return make_help(help_sections, indent)
    end

    return cli
end

function M.new(name: string, description: string?, helpkeys: { string }?)
    -- args that trigger print help + abort
    local helpset: { [string]: boolean } = DEFAULT_helpset
    if helpkeys ~= nil then
        helpset = {}
        for _, k in helpkeys do
            helpset[k] = true
        end
    end

    -- I'm fairly certain this is not the right pattern for making objects, but it results in single definition + great LSP so I'm happy
    local cli = {}
    cli.name = name
    cli.description = description
    cli._helpset = helpset
    cli._positionals = {} :: { ArgData }
    cli._flags = {} :: { ArgData }
    cli._options = {} :: { ArgData }
    cli._argdata_set = {} :: { [string]: ArgData }
    cli._required_list = {} :: { string }
    cli._lookup = {} :: { [string]: ArgData }
    cli._required_positional_count = 0
    cli._default_result = { values = {}, flags = {} } :: ParseResult

    local function add_arg_lookups(keys: { string }, arg_data: ArgData)
        for _, k in keys do
            if cli._helpset[k] ~= nil then
                error(`key '{k}' is already used as a help key.`)
            end
            validate_key_or_error(k)
            if cli._lookup[k] ~= nil then
                error(`key '{k}' already exists.`)
            end
            cli._lookup[k] = arg_data
        end
    end

    local function add_name_error_on_duplicate(name: string, data: ArgData)
        if cli._argdata_set[name] ~= nil then
            error(`arg with name {name} already exists`)
        end
        cli._argdata_set[name] = data
    end

    local function make_safe_options(options: ArgOptions?): ArgOptionsSafe
        if options == nil then
            return DEFAULT_OPTIONS
        end
        assert(options)

        options.aliases = if options.aliases == nil then {} else options.aliases

        return options :: ArgOptionsSafe
    end

    -- return is guaranteed to start with at least one '-', if none, adds '--'
    local function to_lookup_name(name: string): string
        if name:sub(1, 1) == "-" then
            return name
        end
        return "--" .. name
    end

    -- return is guaranted to not start with '-'
    local function to_result_name(name: string): string
        while name:sub(1, 1) == "-" do
            name = name:sub(2, -1)
        end
        return name
    end

    function cli:add_positional(name: string, options: ArgOptions?)
        if name:sub(1, 1) == "-" then
            error(`invalid arg name '{name}', positional cannot start with dashes`)
        end

        local options = make_safe_options(options)

        local arg_data: ArgData = {
            name = name,
            kind = "POSITIONAL",
            options = options,
        }
        add_name_error_on_duplicate(name, arg_data)
        if options.default == nil then
            -- check positional required ordering, can't have a required pos after an optional one
            local last = cli._positionals[#cli._positionals]
            if last ~= nil and last.options.default ~= nil then
                error(
                    `{name} is required, but {last.name} is optional. Cannot have required positional after optional positional`
                )
            end

            table.insert(cli._required_list, name)
        else
            cli._default_result.values[name] = options.default
        end
        table.insert(cli._positionals, arg_data)
    end

    function cli:add_flag(name: string, options: ArgOptions?)
        if options and options.default ~= nil then
            -- these are called constants :)
            error(`flag {name} has non nil default, default value is not supported for flags`)
        end

        local options = make_safe_options(options)

        local lookup_name = to_lookup_name(name)
        name = to_result_name(name)
        for i, v in options.aliases do
            options.aliases[i] = to_lookup_name(v)
        end

        local arg_data: ArgData = {
            name = name,
            kind = "FLAG",
            options = options,
        }
        add_name_error_on_duplicate(name, arg_data)
        add_arg_lookups({ lookup_name, table.unpack(options.aliases) }, arg_data)

        if options.default ~= nil then
            error(`flag {name} has non nil default value, defaults are not supported for flag args`)
        end

        cli._default_result.flags[name] = false
        table.insert(cli._flags, arg_data)
    end

    function cli:add_option(name: string, options: ArgOptions?)
        local options = make_safe_options(options)

        local lookup_name = to_lookup_name(name)
        name = to_result_name(name)
        for i, v in options.aliases do
            options.aliases[i] = to_lookup_name(v)
        end

        local arg_data: ArgData = {
            name = name,
            kind = "OPTION",
            options = options,
        }
        add_name_error_on_duplicate(name, arg_data)
        add_arg_lookups({ lookup_name, table.unpack(options.aliases) }, arg_data)
        if options.default == nil then
            table.insert(cli._required_list, name)
        else
            if options.default == nil then
                error(`optional arg {name} must have a default value`)
            end
            assert(options.default ~= nil)
            cli._default_result.values[name] = options.default
        end
        table.insert(cli._options, arg_data)
    end

    -- return: data, err - where data is a table, and err is a string
    -- errors early, first encountered error ends the parse
    -- if a help key is found will print help text and exit the process
    function cli:parse(args_in: { string }): (ParseResult?, string?)
        local args = {}
        for _, v in args_in do
            table.insert(args, v)
        end
        local parsed: ParseResult = { values = {}, flags = {}, fwd_args = {} } :: ParseResult

        local positional_idx = 1
        local skip_next = false -- used for options
        for i, arg in args do
            -- asking for help?
            if cli._helpset[arg] ~= nil then
                print(cli:help())
                process.exit(0)
            end

            if skip_next then
                skip_next = false
                continue
            end

            -- rest of args should be fwd args
            if arg == "--" then
                parsed.fwd_args = { table.unpack(args, i + 1) }
                break
            end

            local has_dash = arg:sub(1, 1) == "-"
            local arg_data: ArgData? = nil
            local value: string? = nil

            if not has_dash then -- positional
                if positional_idx > #cli._positionals then
                    return nil, `too many positional arguments, expected {#cli._positionals}.`
                end
                arg_data = cli._positionals[positional_idx]
                positional_idx += 1
                value = arg
            else -- flag or option
                arg_data = cli._lookup[arg]
            end

            if arg_data == nil then
                return nil, `unknown flag or option: '{arg}'`
            end
            assert(arg_data ~= nil)

            -- handle flags early
            if arg_data.kind == "FLAG" then
                parsed.flags[arg_data.name] = true
                continue
            end

            -- get option value
            if arg_data.kind == "OPTION" then
                if i < #args then
                    value = args[i + 1]
                    skip_next = true
                end
            end

            -- resolve positionals and options
            if value == nil then
                return nil, `no value provided for option '{arg}'`
            end
            assert(value ~= nil)

            parsed.values[arg_data.name] = value
        end

        -- check all required options are provided
        for _, name in cli._required_list do
            if parsed.values[name] == nil and parsed.flags[name] == nil then
                local kind = cli._argdata_set[name].kind:lower()
                return nil, `required {kind} arg '{name}' was not found`
            end
        end

        -- apply defaults for missing optional
        for k, v in cli._default_result.flags do
            if parsed.flags[k] == nil then
                parsed.flags[k] = v
            end
        end
        for k, v in cli._default_result.values do
            if parsed.values[k] == nil then
                parsed.values[k] = v
            end
        end

        return parsed, nil
    end

    function cli:help(indent: number?): string
        local function align_cols(rows: { { string } }, sep: string?): { { string } }
            local sep = if sep == nil then " " else sep
            local max_col_lengths: { number } = {}
            for _, row in rows do
                for i, s in row do
                    if max_col_lengths[i] == nil or s:len() > max_col_lengths[i] then
                        max_col_lengths[i] = s:len()
                    end
                end
            end
            local aligned_rows: { { string } } = {}
            for _, row in rows do
                local line: { string } = {}
                for i, col in row do
                    table.insert(line, col)
                    if i < #row then
                        local spacing = sep
                        local diff = max_col_lengths[i] - col:len()
                        if diff > 0 then
                            spacing = string.rep(" ", diff) .. sep
                        end
                        table.insert(line, spacing)
                    end
                end
                table.insert(aligned_rows, line)
            end
            return aligned_rows
        end

        local help_sections: { HelpSection } = {}

        -- usage
        local usage = `usage: {cli.name}`
        if #cli._flags > 0 or #cli._options > 0 then
            usage ..= " [options]"
        end
        for _, arg in cli._positionals do
            if arg.options.default == nil then
                usage ..= ` <{arg.name}>`
            else
                usage ..= ` [{arg.name}]`
            end
        end

        table.insert(help_sections, { title = usage })

        -- description
        if cli.description ~= nil then
            table.insert(help_sections, { title = "description:", lines = { cli.description } })
        end

        local function make_arg_line(arg: ArgData): { string }
            local keys = arg.name
            if arg.kind ~= "POSITIONAL" then
                keys = `--{arg.name}`
                for _, a in arg.options.aliases do
                    keys ..= `, {a}`
                end
            end

            local reqdef = ""
            -- optional & default
            --[[
            if arg.kind ~= "FLAG" then
                if arg.options.default == nil then
                    reqdef ..= "[required]"
                else
                    assert(arg.options.default)
                    reqdef ..= `[default:  '{arg.options.default}']`
                end
            end
            ]]

            local help = ""
            if arg.options.help then
                if help ~= "" then
                    help ..= " "
                end
                help ..= arg.options.help
            end

            return { keys, reqdef, help }
        end

        -- positionals
        if #cli._positionals > 0 then
            local positional_lines: { { string } } = {}
            for _, arg in cli._positionals do
                table.insert(positional_lines, make_arg_line(arg))
            end
            table.insert(help_sections, { title = "positional arguments:", lines = positional_lines })
        end

        -- flags
        if #cli._flags > 0 then
            local flag_lines: { { string } } = {}
            for _, arg in cli._flags do
                table.insert(flag_lines, make_arg_line(arg))
            end
            table.insert(help_sections, { title = "flags:", lines = flag_lines })
        end

        -- options
        if #cli._options > 0 then
            local option_lines: { { string } } = {}
            for _, arg in cli._options do
                table.insert(option_lines, make_arg_line(arg))
            end
            table.insert(help_sections, { title = "options:", lines = option_lines })
        end

        return make_help(help_sections, indent)
    end

    return cli
end

export type Cli = typeof(M.new(""))
export type CliSubcommands = typeof(M.new_subcommands(""))

return M


--!nolint LocalShadow

local fs = require("@lune/fs")
local LuaEncode = require("./LuaEncode")
local tasks_client = require("./tasks_client")
local utils = require("./utils")

local m_tasks = {}

export type Error = tasks_client.TasksError
export type CreateTaskParams = tasks_client.CreateTaskParams
export type PollTaskOpts = tasks_client.PollTaskOpts

local function make_task_result(output, path: string, poll_res: any, req_params: CreateTaskParams)
    local res = {
        output = output,
        path = path,
        poll_res = poll_res,
        req_params = req_params,
    }

    function res:fetch_logs(): ({ string }?, Error?)
        self = self :: typeof(res)
        return tasks_client.fetch_task_logs(self.req_params.api_key, self.path)
    end

    return res
end

local function make_task_handle<Rtn>(
    path: string,
    task_create_res: any,
    req_params: CreateTaskParams,
    poll_opts: PollTaskOpts,
    _rtn_sig: Rtn
)
    local handle = {
        path = path,
        task_create_res = task_create_res,
        req_params = req_params,
        poll_opts = poll_opts,
    }

    type TaskResult = typeof(make_task_result(({} :: any) :: Rtn, "", nil, nil :: any))

    function handle:wait(): (TaskResult?, Error?)
        self = self :: typeof(handle)
        local res, err = tasks_client.poll_task_completion(self.req_params.api_key, self.path, self.poll_opts)
        if err ~= nil then
            return nil, err
        end

        local output = utils.safe_access(res)["output"]["results"][1]()

        return make_task_result(output :: Rtn, path, res, req_params)
    end

    return handle
end

function m_tasks.encode_task_from_file(file_path: string, ...): (string?, Error?)
    local task_fn_str = fs.readFile(file_path)
    return m_tasks.encode_task(task_fn_str, ...)
end

function m_tasks.encode_task(task_fn_str: string, ...): (string?, Error?)
    local args = { ... }
    local args_encode = ""
    local ok, err = pcall(function()
        args_encode = LuaEncode(args, { OutputWarnings = false })
    end)
    if not ok then
        return nil, { message = `error encoding args: {err}` }
    end

    local task_encode =
        `function t() {task_fn_str}\nend;\nfunction a() return {args_encode}\nend\nreturn t()(table.unpack(a()))`
    return task_encode
end

function m_tasks.create_from_file(file_path: string, req_params: CreateTaskParams, _fn_sig)
    local task_str = fs.readFile(file_path)
    return m_tasks.create(task_str, req_params, _fn_sig)
end

function m_tasks.create<Rtn, Arg...>(task_str: string, req_params: CreateTaskParams, _fn_sig: ((Arg...) -> Rtn)?)
    local task = {
        task_str = task_str,
        req_params = req_params,
        poll_opts = {} :: PollTaskOpts,
    }

    type TaskHandle = typeof(make_task_handle(
        "",
        nil,
        { api_key = "", universe_id = "", place_id = "" },
        {},
        (nil :: any) :: Rtn
    ))

    function task:start(...: Arg...): (TaskHandle?, Error?)
        self = self :: typeof(task)
        local task_encode, err = m_tasks.encode_task(self.task_str, ...)
        if err ~= nil then
            return nil, err
        end
        assert(task_encode ~= nil)

        local res, err = tasks_client.create_task(task_encode, self.req_params)
        if err ~= nil then
            return nil, err
        end
        assert(res)
        -- create_task would have errored if res had no memeber 'path'
        return make_task_handle(res.path, res, self.req_params, self.poll_opts, (nil :: any) :: Rtn)
    end

    type TaskResult = typeof(make_task_result(({} :: any) :: Rtn, "", nil, nil :: any))
    --- wraps task creation and polling
    function task:call(...: Arg...): (TaskResult?, Error?)
        self = self :: typeof(task)
        local handle, err = self:start(...)
        if err ~= nil then
            return nil, err
        end
        assert(handle ~= nil)
        handle = handle :: TaskHandle
        local res, err = handle:wait()
        if err ~= nil then
            return nil, err
        end
        assert(res ~= nil)
        return res
    end

    return task
end

return m_tasks


local m_utils = {}

function m_utils.notnil_or_default<T>(v: T?, d: T): T
    if v == nil then
        return d
    end
    return v
end

-- allows chaining member accesses forwarding nil if any intermediate step was nil
-- []/index to access members
-- () to get the value, optionally pass in a default
function m_utils.safe_access(t: { [any]: any }?)
    local v = {
        ___value = t,
    }

    setmetatable(v, {
        __index = function(_, k: any)
            if v.___value == nil then
                return v
            end
            v.___value = v.___value[k]
            return v
        end,

        __call = function(d: any?)
            if v.___value ~= nil then
                return v.___value :: any
            end
            return d
        end,
    })

    return v
end

return m_utils


--!strict

--[[
    Core CLI functionality for parsing and handling command line arguments.
]]

-- Native Libraries
local process = require('@lune/process')

-- External Libraries
local frkcli = require('../../libs/frkcli/frkcli')

-- Local Libraries
local Constants = require('../shared/constants')

--[[
    Generates the CLI configuration and mappings.
    @return (any, { [string]: string }, { [string]: string }) - Returns the CLI instance and configuration mappings
]]
local function GenerateCLI(): (any, { [string]: string }, { [string]: string })
    local cli = frkcli.new_subcommands('lune run batch-process', 'Data Stores Batch Processing CLI')
    local configToOptionFlagMap: { [string]: string } = {}
    local inverseFlagMap: { [string]: string } = {}

    --[[
        Helper function to add configuration options to a command.
        @param command any - The command to add options to
        @param configs { [string]: Types.ConfigOption } - The configuration options to add
    ]]
    local function addConfigOptions(command, configs: { string }): ()
        for i: number, configName: string in ipairs(configs) do
            -- Get the config option
            local configOption = Constants.ConfigurationOptions[configName]
            if not configOption then
                continue
            end

            -- Add the option to the command
            command:add_option(configOption.option, {
                help = configOption.description,
                aliases = { configOption.shortOption },
                default = Constants.PlaceHolder
            })
            configToOptionFlagMap[configName] = configOption.option
        end
    end

    --[[
        Helper function to add configuration flags to a command.
        @param command any - The command to add flags to
        @param configs { [string]: Types.ConfigFlag } - The configuration flags to add
    ]]
    local function addConfigFlags(command, configs: { string }): ()
        for i: number, configName: string in ipairs(configs) do
            -- Get the config flag
            local configFlag = Constants.ConfigurationFlags[configName]
            if not configFlag then
                continue
            end

            -- Add the flag to the command
            command:add_flag(configFlag.flag, {
                help = configFlag.description
            })
            configToOptionFlagMap[configName] = configFlag.flag

            -- Add the inverse flag to the command if it exists
            if configFlag.flagInverse then
                command:add_flag(configFlag.flagInverse, {
                    help = configFlag.descriptionInverse or ('not (' .. configFlag.description .. ')')
                })
                inverseFlagMap[configFlag.flag] = configFlag.flagInverse
            end
        end
    end

    -- Add each subcommand to the CLI
    for command, commandDetails in Constants.CommandDetails do
        -- Create the subcommand
        local subcommand = cli:add_subcommand(command, commandDetails.description)
        if commandDetails.positionalArgs then
            for _, positionalArg in ipairs(commandDetails.positionalArgs) do
                subcommand:add_positional(positionalArg.name, {
                    help = positionalArg.description,
                })
            end
        end

        -- Add a config option to the subcommand
        subcommand:add_option('config', {
            help = 'The config file to use for the process',
            aliases = { '-c' },
            default = ''
        })

        -- Add subcommand-specific config options and flags
        addConfigOptions(subcommand, commandDetails.configurations)
        addConfigFlags(subcommand, commandDetails.configurations)
    end

    return cli, configToOptionFlagMap, inverseFlagMap
end

-- CLI
local CLI = {}

--[[
    Parses command line arguments into a structured format.
    @return (Types.CommandType, Types.ConfigValues, Types.ConfigValues, { [string]: string }, { [string]: string }) - Returns the parsed command, options, flags, and mappings
]]
CLI.Parse = function(): { [string]: any }
    local cli, configToOptionFlagMap, inverseFlagMap = GenerateCLI()
    -- Parse the arguments
    local res, err = cli:parse(process.args)

    -- If there is an error, print it and exit
    if err then
        error(err)
    end

    -- Return the command, options, and flags
    return {
        command = res.command,
        values = res.result.values,
        flags = res.result.flags,
        configToOptionFlagMap = configToOptionFlagMap,
        inverseFlagMap = inverseFlagMap
    }
end

return CLI


--!strict

--[[
    Command for stopping a batch processing operation.
]]

-- Native Libraries
local process = require('@lune/process')

-- Local Libraries
local CLIO = require('../util/clio')
local ConfigValidators = require('../config/validators')
local Constants = require('../shared/constants')
local OpenCloudClient = require('../util/open-cloud-client')
local Types = require('../shared/types')

-- List
local List = {}

--[[
    Lists all batch processes.
    @param configs Types.ListConfigs - The configurations for the process
]]
List.Execute = function(configs: Types.ListConfigs): ()
     local openCloudClient = OpenCloudClient.CreateOpenCloudClient(process.env['API_KEY'], configs.universeId)

    -- Validate the API Key
    local validateAPIKeySuccess: boolean, validateAPIKeyErrorMessage: string? = ConfigValidators.ValidateAPIKey(openCloudClient, process.env['API_KEY'])
    if not validateAPIKeySuccess then
        CLIO.Error(validateAPIKeyErrorMessage :: string)
        return
    end

     -- loop over all of the batch processes in data stores
    local cursor = ""
    local uniqueBatchProcesses = {}
    repeat
        local page = openCloudClient.listDataStoreEntries(Constants.BatchProcessMapName, "global", 1, cursor)
        if page.status == 404 then
            break
        elseif page.status ~= 200 then
            CLIO.Error("There was an error listing the batch processes. Please try again later. Reason: " .. OpenCloudClient.GetBodyAsString(page))
            return
        end

        if not page.body or not page.body.dataStoreEntries then
            break
        end

        for _, batchProcess in ipairs(page.body.dataStoreEntries) do
            uniqueBatchProcesses[batchProcess.id] = true
        end
        cursor = page.body.nextPageToken
    until cursor == nil or cursor == ""

    -- loop over all of the batch processes in memory stores
    cursor = ""
    repeat
        local page = openCloudClient.getSortedMapRange(Constants.BatchProcessMapName, true, 200, cursor)
        if page.status ~= 200 then
            CLIO.Error("There was an error listing the batch processes. Please try again later. Reason: " .. OpenCloudClient.GetBodyAsString(page))
            return
        end
        for _, batchProcess in ipairs(page.body.items) do
            local processInfo: Types.BatchProcess = batchProcess.value
            if not uniqueBatchProcesses[processInfo.name] then
                -- If the process is not in data stores, try to backup the process to data stores
                local backupProcessToDataStores = openCloudClient.createDataStoreEntry(Constants.BatchProcessMapName, "global", processInfo.name, processInfo)
                if backupProcessToDataStores.status ~= 200 then
                    CLIO.Error("There was an error backing up the process to data stores. Please try again later. Reason: " .. OpenCloudClient.GetBodyAsString(backupProcessToDataStores))
                end
                uniqueBatchProcesses[processInfo.name] = true
            end
        end
        cursor = page.body.nextPageToken
    until cursor == nil or cursor == ""

    -- Extract the batch process names
    local batchProcessNames = {}
    for name, _ in pairs(uniqueBatchProcesses) do
        table.insert(batchProcessNames, name)
    end

    if #batchProcessNames == 0 then
        CLIO.Print("No batch processes found.")
    else
        CLIO.PrintBatchProcessList(batchProcessNames)
    end
end

return List