local Nodes = {}

-- Operator types
export type UnaryOpType = "Invert" | "Not" | "UAdd" | "USub"
export type BinOpType =
	"Add"
	| "Sub"
	| "Mult"
	| "MatMult"
	| "Div"
	| "Mod"
	| "Pow"
	| "LShift"
	| "RShift"
	| "BitOr"
	| "BitXor"
	| "BitAnd"
	| "FloorDiv"
export type CmpOp = "Eq" | "NotEq" | "Lt" | "LtE" | "Gt" | "GtE" | "Is" | "IsNot" | "In" | "NotIn"
export type BoolOpType = "And" | "Or"

-- Base AST node type with source location information
export type ASTNode = {
	type: string, -- Node type identifier
	lineno: number, -- Source line number
	col_offset: number, -- Column offset
	end_lineno: number?, -- End line number
	end_col_offset: number?, -- End column offset
}

-- Expression base type
export type Expr = ASTNode & {
	-- All expressions inherit from this
}

-- Statement base type
export type Stmt = ASTNode & {
	-- All statements inherit from this
}

-- MODULE
-- Module node (top-level container)
export type Module = ASTNode & {
	type: "Module",
	body: { Stmt },
	type_ignores: { any }?, -- For type: ignore comments
}

-- STATEMENTS
-- Assignment statement: x = 1, x, y = 1, 2
export type Assign = Stmt & {
	type: "Assign",
	targets: { Expr }, -- Left-hand side expressions
	value: Expr, -- Right-hand side expression
}

-- Augmented assignment: x += 1
export type AugAssign = Stmt & {
	type: "AugAssign",
	target: Expr, -- Left-hand side
	op: BinOpType, -- Operator
	value: Expr, -- Right-hand side
}

-- Expression statement: function_call()
export type ExprStmt = Stmt & {
	type: "Expr",
	value: Expr,
}

-- If statement: if condition: ... elif condition: ... else: ...
export type If = Stmt & {
	type: "If",
	test: Expr, -- Condition
	body: { Stmt }, -- If body
	orelse: { Stmt }, -- Else/elif body
}

-- While loop: while condition: ...
export type While = Stmt & {
	type: "While",
	test: Expr, -- Condition
	body: { Stmt }, -- Loop body
	orelse: { Stmt }, -- Else clause
}

-- For loop: for target in iter: ...
export type For = Stmt & {
	type: "For",
	target: Expr, -- Loop variable
	iter: Expr, -- Iterable
	body: { Stmt }, -- Loop body
	orelse: { Stmt }, -- Else clause
}

-- Function definition: def name(args): ...
export type FunctionDef = Stmt & {
	type: "FunctionDef",
	name: string, -- Function name
	args: Arguments, -- Parameters
	body: { Stmt }, -- Function body
	decorator_list: { Expr }, -- Decorators
	returns: Expr?, -- Return type annotation
}

export type ExceptHandler = {
	type: "ExceptHandler",
	lineno: number,
	col_offset: number,
	end_lineno: number?,
	end_col_offset: number?,
	body: { Stmt },
	expr: Expr?,
	name: string?,
}

export type Try = Stmt & {
	type: "Try",
	body: { Stmt },
	handlers: { ExceptHandler },
	orelse: { Stmt },
	finalbody: { Stmt },
}

-- Class definition: class Name(bases): ...
export type ClassDef = Stmt & {
	type: "ClassDef",
	name: string, -- Class name
	bases: { Expr }, -- Base classes
	keywords: { Keyword }, -- Keyword arguments
	body: { Stmt }, -- Class body
	decorator_list: { Expr }, -- Decorators
}

-- Import alias: name as asname
export type alias = {
	name: string,
	asname: string?,
	lineno: number,
	col_offset: number,
}

-- Import statement: import x, y as z
export type Import = Stmt & {
	type: "Import",
	names: { alias },
}

-- Import from statement: from .x import y
export type ImportFrom = Stmt & {
	type: "ImportFrom",
	module: string?,
	names: { alias },
	level: number?,
}

-- Return statement: return value
export type Return = Stmt & {
	type: "Return",
	value: Expr?, -- Return value (optional)
}

-- Break statement: break
export type Break = Stmt & {
	type: "Break",
}

-- Continue statement: continue
export type Continue = Stmt & {
	type: "Continue",
}

-- Pass statement: pass
export type Pass = Stmt & {
	type: "Pass",
}

-- EXPRESSIONS
-- Binary operation: left op right
export type BinOpNode = Expr & {
	type: "BinOp",
	left: Expr, -- Left operand
	op: BinOpType, -- Operator
	right: Expr, -- Right operand
}

-- Unary operation: op operand
export type UnaryOpNode = Expr & {
	type: "UnaryOp",
	op: UnaryOpType, -- Operator
	operand: Expr, -- Operand
}

-- Boolean operation: left op right (and/or)
export type BoolOpNode = Expr & {
	type: "BoolOp",
	op: BoolOpType, -- Operator (And/Or)
	values: { Expr }, -- Operands (2 or more)
}

-- Comparison: left ops comparators
export type Compare = Expr & {
	type: "Compare",
	left: Expr, -- Left operand
	ops: { CmpOp }, -- Comparison operators
	comparators: { Expr }, -- Right operands
}

-- Function call: func(args, keywords)
export type Call = Expr & {
	type: "Call",
	func: Expr, -- Function to call
	args: { Expr }, -- Positional arguments
	keywords: { Keyword }, -- Keyword arguments
}

-- Attribute access: value.attr
export type Attribute = Expr & {
	type: "Attribute",
	value: Expr, -- Object
	attr: string, -- Attribute name
	ctx: ExprContext, -- Load/Store/Del context
}

-- Subscript: value[slice]
export type Subscript = Expr & {
	type: "Subscript",
	value: Expr, -- Object
	slice: Expr, -- Index/slice
	ctx: ExprContext, -- Load/Store/Del context
}

-- Variable name: identifier
export type Name = Expr & {
	type: "Name",
	id: string, -- Identifier name
	ctx: ExprContext, -- Load/Store/Del context
}

-- Constant value: 42, "hello", True, None
export type Constant = Expr & {
	type: "Constant",
	value: any, -- The constant value
	kind: string?, -- String prefix for string literals
}

-- List literal: [1, 2, 3]
export type List = Expr & {
	type: "List",
	elts: { Expr }, -- Elements
	ctx: ExprContext, -- Load/Store/Del context
}

-- Tuple literal: (1, 2, 3)
export type Tuple = Expr & {
	type: "Tuple",
	elts: { Expr }, -- Elements
	ctx: ExprContext, -- Load/Store/Del context
}

-- Dictionary literal: {'key': 'value'}
export type Dict = Expr & {
	type: "Dict",
	keys: { Expr }, -- Keys (DictUnpack sentinel nodes for ** entries)
	values: { Expr }, -- Values
}

-- Set literal: {1, 2, 3}
export type Set = Expr & {
	type: "Set",
	elts: { Expr }, -- Elements
}

-- List comprehension: [elt for target in iter if ...]
export type Comprehension = {
	target: Expr,
	iter: Expr,
	ifs: { Expr },
	is_async: boolean,
}

export type ListComp = Expr & {
	type: "ListComp",
	elt: Expr,
	generators: { Comprehension },
}

-- Lambda expression: lambda parameters: expression
export type Lambda = Expr & {
	type: "Lambda",
	args: Arguments,
	body: Expr,
}

-- Yield expression: yield value (optional)
export type Yield = Expr & {
	type: "Yield",
	value: Expr?,
}

-- Yield from expression: yield from iterator
export type YieldFrom = Expr & {
	type: "YieldFrom",
	value: Expr,
}

-- Starred expression: *value (in lists, tuples, calls)
export type Starred = Expr & {
	type: "Starred",
	value: Expr,
	ctx: ExprContext,
}

-- Dict **unpack sentinel inside a dict literal (mirrors CPython using None in keys list)
export type DictUnpack = Expr & {
	type: "DictUnpack",
	value: Expr,
}

-- HELPER TYPES
-- Expression context (how the expression is used)
export type ExprContext = "Load" | "Store" | "Del"

-- Function arguments
export type Arguments = {
	posonlyargs: { Arg }, -- Positional-only arguments
	args: { Arg }, -- Regular arguments
	vararg: Arg?, -- *args
	kwonlyargs: { Arg }, -- Keyword-only arguments
	kw_defaults: { Expr? }, -- Keyword argument defaults
	kwarg: Arg?, -- **kwargs
	defaults: { Expr }, -- Default values
}

-- Function argument
export type Arg = {
	arg: string, -- Argument name
	annotation: Expr?, -- Type annotation
	lineno: number,
	col_offset: number,
	end_lineno: number?,
	end_col_offset: number?,
}

-- Keyword argument
export type Keyword = {
	arg: string?, -- Keyword name (None for **)
	value: Expr, -- Value
	lineno: number,
	col_offset: number,
	end_lineno: number?,
	end_col_offset: number?,
}

-- Union type for all statement types
export type AnyStmt =
	Assign
	| AugAssign
	| ExprStmt
	| If
	| While
	| For
	| FunctionDef
	| ClassDef
	| Import
	| ImportFrom
	| Return
	| Break
	| Continue
	| Pass
	| Try

-- Union type for all expression types
export type AnyExpr =
	BinOpNode
	| UnaryOpNode
	| BoolOpNode
	| Compare
	| Call
	| Attribute
	| Subscript
	| Name
	| Constant
	| List
	| Tuple
	| Dict
	| Set
	| ListComp
	| Lambda
	| Yield
	| YieldFrom
	| Starred
	| DictUnpack

-- Union type for all AST node types
export type AnyNode = Module | AnyStmt | AnyExpr

-- UTILITY FUNCTIONS
-- Create a new AST node with source location
function Nodes.new_node(node_type: string, lineno: number, col_offset: number): ASTNode
	return {
		type = node_type,
		lineno = lineno,
		col_offset = col_offset,
		end_lineno = nil,
		end_col_offset = nil,
	}
end

-- Set end location for an AST node
function Nodes.set_end_location(node: ASTNode, end_lineno: number, end_col_offset: number): ASTNode
	node.end_lineno = end_lineno
	node.end_col_offset = end_col_offset
	return node
end

-- Copy location information from one node to another
function Nodes.copy_location(from_node: ASTNode, to_node: ASTNode): ASTNode
	to_node.lineno = from_node.lineno
	to_node.col_offset = from_node.col_offset
	to_node.end_lineno = from_node.end_lineno
	to_node.end_col_offset = from_node.end_col_offset
	return to_node
end

return Nodes


local Visitor = {}

local nodes = require("./nodes")
type ASTNode = nodes.ASTNode
type AnyNode = nodes.AnyNode
type AnyStmt = nodes.AnyStmt
type AnyExpr = nodes.AnyExpr

-- Base visitor class for AST traversal
export type Visitor = {
	visit: (self: Visitor, node: ASTNode) -> any,
	generic_visit: (self: Visitor, node: ASTNode) -> any,

	-- Statement visitors
	visit_Module: ((self: Visitor, node: nodes.Module) -> any)?,
	visit_Assign: ((self: Visitor, node: nodes.Assign) -> any)?,
	visit_AugAssign: ((self: Visitor, node: nodes.AugAssign) -> any)?,
	visit_Expr: ((self: Visitor, node: nodes.ExprStmt) -> any)?,
	visit_If: ((self: Visitor, node: nodes.If) -> any)?,
	visit_While: ((self: Visitor, node: nodes.While) -> any)?,
	visit_For: ((self: Visitor, node: nodes.For) -> any)?,
	visit_FunctionDef: ((self: Visitor, node: nodes.FunctionDef) -> any)?,
	visit_ClassDef: ((self: Visitor, node: nodes.ClassDef) -> any)?,
	visit_Return: ((self: Visitor, node: nodes.Return) -> any)?,
	visit_Break: ((self: Visitor, node: nodes.Break) -> any)?,
	visit_Continue: ((self: Visitor, node: nodes.Continue) -> any)?,
	visit_Pass: ((self: Visitor, node: nodes.Pass) -> any)?,

	-- Expression visitors
	visit_BinOp: ((self: Visitor, node: nodes.BinOpNode) -> any)?,
	visit_UnaryOp: ((self: Visitor, node: nodes.UnaryOpNode) -> any)?,
	visit_BoolOp: ((self: Visitor, node: nodes.BoolOpNode) -> any)?,
	visit_Compare: ((self: Visitor, node: nodes.Compare) -> any)?,
	visit_Call: ((self: Visitor, node: nodes.Call) -> any)?,
	visit_Attribute: ((self: Visitor, node: nodes.Attribute) -> any)?,
	visit_Subscript: ((self: Visitor, node: nodes.Subscript) -> any)?,
	visit_Name: ((self: Visitor, node: nodes.Name) -> any)?,
	visit_Constant: ((self: Visitor, node: nodes.Constant) -> any)?,
	visit_List: ((self: Visitor, node: nodes.List) -> any)?,
	visit_Tuple: ((self: Visitor, node: nodes.Tuple) -> any)?,
	visit_Dict: ((self: Visitor, node: nodes.Dict) -> any)?,
	visit_Set: ((self: Visitor, node: nodes.Set) -> any)?,
}

-- Create a new visitor instance
function Visitor.new(): Visitor
	local visitor = {}
	setmetatable(visitor, { __index = Visitor })
	return visitor :: any
end

-- Main visit method - dispatches to specific visit methods
function Visitor:visit(node: ASTNode): any
	local method_name = "visit_" .. node.type
	local method = (self :: any)[method_name]

	if method then
		return method(self, node)
	else
		return self:generic_visit(node)
	end
end

-- Generic visit - recursively visits all child nodes
function Visitor:generic_visit(node: ASTNode): any
	local node_type = node.type

	-- Module
	if node_type == "Module" then
		local module_node = node :: nodes.Module
		for _, stmt in ipairs(module_node.body) do
			self:visit(stmt)
		end

	-- Statements
	elseif node_type == "Assign" then
		local assign_node = node :: nodes.Assign
		for _, target in ipairs(assign_node.targets) do
			self:visit(target)
		end
		self:visit(assign_node.value)
	elseif node_type == "AugAssign" then
		local aug_assign_node = node :: nodes.AugAssign
		self:visit(aug_assign_node.target)
		self:visit(aug_assign_node.value)
	elseif node_type == "Expr" then
		local expr_stmt_node = node :: nodes.ExprStmt
		self:visit(expr_stmt_node.value)
	elseif node_type == "If" then
		local if_node = node :: nodes.If
		self:visit(if_node.test)
		for _, stmt in ipairs(if_node.body) do
			self:visit(stmt)
		end
		for _, stmt in ipairs(if_node.orelse) do
			self:visit(stmt)
		end
	elseif node_type == "While" then
		local while_node = node :: nodes.While
		self:visit(while_node.test)
		for _, stmt in ipairs(while_node.body) do
			self:visit(stmt)
		end
		for _, stmt in ipairs(while_node.orelse) do
			self:visit(stmt)
		end
	elseif node_type == "For" then
		local for_node = node :: nodes.For
		self:visit(for_node.target)
		self:visit(for_node.iter)
		for _, stmt in ipairs(for_node.body) do
			self:visit(stmt)
		end
		for _, stmt in ipairs(for_node.orelse) do
			self:visit(stmt)
		end
	elseif node_type == "FunctionDef" then
		local func_node = node :: nodes.FunctionDef
		-- Visit arguments
		local args = func_node.args
		for _, arg in ipairs(args.posonlyargs) do
			if arg.annotation then
				self:visit(arg.annotation)
			end
		end
		for _, arg in ipairs(args.args) do
			if arg.annotation then
				self:visit(arg.annotation)
			end
		end
		if args.vararg and args.vararg.annotation then
			self:visit(args.vararg.annotation)
		end
		for _, arg in ipairs(args.kwonlyargs) do
			if arg.annotation then
				self:visit(arg.annotation)
			end
		end
		if args.kwarg and args.kwarg.annotation then
			self:visit(args.kwarg.annotation)
		end
		for _, default in ipairs(args.defaults) do
			self:visit(default)
		end
		for _, default in ipairs(args.kw_defaults) do
			if default then
				self:visit(default)
			end
		end

		-- Visit decorators
		for _, decorator in ipairs(func_node.decorator_list) do
			self:visit(decorator)
		end

		-- Visit return annotation
		if func_node.returns then
			self:visit(func_node.returns)
		end

		-- Visit body
		for _, stmt in ipairs(func_node.body) do
			self:visit(stmt)
		end
	elseif node_type == "ClassDef" then
		local class_node = node :: nodes.ClassDef
		for _, base in ipairs(class_node.bases) do
			self:visit(base)
		end
		for _, keyword in ipairs(class_node.keywords) do
			self:visit(keyword.value)
		end
		for _, decorator in ipairs(class_node.decorator_list) do
			self:visit(decorator)
		end
		for _, stmt in ipairs(class_node.body) do
			self:visit(stmt)
		end
	elseif node_type == "Return" then
		local return_node = node :: nodes.Return
		if return_node.value then
			self:visit(return_node.value)
		end

	-- Break, Continue, Pass have no child nodes

	-- Expressions
	elseif node_type == "BinOp" then
		local binop_node = node :: nodes.BinOpNode
		self:visit(binop_node.left)
		self:visit(binop_node.right)
	elseif node_type == "UnaryOp" then
		local unaryop_node = node :: nodes.UnaryOpNode
		self:visit(unaryop_node.operand)
	elseif node_type == "BoolOp" then
		local boolop_node = node :: nodes.BoolOpNode
		for _, value in ipairs(boolop_node.values) do
			self:visit(value)
		end
	elseif node_type == "Compare" then
		local compare_node = node :: nodes.Compare
		self:visit(compare_node.left)
		for _, comparator in ipairs(compare_node.comparators) do
			self:visit(comparator)
		end
	elseif node_type == "Call" then
		local call_node = node :: nodes.Call
		self:visit(call_node.func)
		for _, arg in ipairs(call_node.args) do
			self:visit(arg)
		end
		for _, keyword in ipairs(call_node.keywords) do
			self:visit(keyword.value)
		end
	elseif node_type == "Attribute" then
		local attr_node = node :: nodes.Attribute
		self:visit(attr_node.value)
	elseif node_type == "Subscript" then
		local subscript_node = node :: nodes.Subscript
		self:visit(subscript_node.value)
		self:visit(subscript_node.slice)
	elseif node_type == "List" then
		local list_node = node :: nodes.List
		for _, elt in ipairs(list_node.elts) do
			self:visit(elt)
		end
	elseif node_type == "Tuple" then
		local tuple_node = node :: nodes.Tuple
		for _, elt in ipairs(tuple_node.elts) do
			self:visit(elt)
		end
	elseif node_type == "Dict" then
		local dict_node = node :: nodes.Dict
		for i, key in ipairs(dict_node.keys) do
			if key then -- Can be nil for ** unpacking
				self:visit(key)
			end
			self:visit(dict_node.values[i])
		end
	elseif node_type == "Set" then
		local set_node = node :: nodes.Set
		for _, elt in ipairs(set_node.elts) do
			self:visit(elt)
		end

		-- Name and Constant have no child nodes
	end

	return nil
end

-- UTILITY VISITORS
-- Node transformer - creates modified copies of AST nodes
export type NodeTransformer = Visitor & {
	transform: (self: NodeTransformer, node: ASTNode) -> ASTNode,
}

function Visitor.new_transformer(): NodeTransformer
	local transformer = Visitor.new() :: any

	-- Override visit to transform nodes
	function transformer:transform(node: ASTNode): ASTNode
		return self:visit(node) or node
	end

	return transformer
end

-- Node collector - collects nodes of specific types
export type NodeCollector = Visitor & {
	collected: { ASTNode },
	node_types: { [string]: boolean },
	collect: (self: NodeCollector, node: ASTNode, node_types: { string }) -> { ASTNode },
}

function Visitor.new_collector(): NodeCollector
	local collector = Visitor.new() :: any
	collector.collected = {}
	collector.node_types = {}

	function collector:collect(node: ASTNode, node_types: { string }): { ASTNode }
		self.collected = {}
		self.node_types = {}

		for _, node_type in ipairs(node_types) do
			self.node_types[node_type] = true
		end

		self:visit(node)
		return self.collected
	end

	-- Override generic_visit to collect matching nodes
	function collector:generic_visit(node: ASTNode): any
		if self.node_types[node.type] then
			table.insert(self.collected, node)
		end

		-- Continue with normal traversal
		Visitor.generic_visit(self, node)
		return nil
	end

	return collector
end

return Visitor


local Instructions = {}

local opcodes = require("./opcodes")
type Opcode = opcodes.Opcode

-- Bytecode instruction
export type Instruction = {
	opcode: Opcode,
	arg: number?, -- Optional numeric argument (index / delta / small int)
	lineno: number?, -- Source line for debug / traceback
}

-- Create an instruction; validate arg presence vs opcode metadata
function Instructions.make(op: Opcode, arg: number?, lineno: number?): Instruction
	local needs = opcodes.requiresArg(op)
	if needs and arg == nil then
		error("Opcode " .. op .. " requires argument", 2)
	elseif (not needs) and arg ~= nil then
		error("Opcode " .. op .. " does not take argument", 2)
	end
	return { opcode = op, arg = arg, lineno = lineno }
end

-- Pretty print a single instruction
function Instructions.render(inst: Instruction): string
	if inst.arg ~= nil then
		return string.format("%s %d", inst.opcode, inst.arg)
	end
	return inst.opcode
end

-- Pretty print a list of instructions
function Instructions.renderAll(code: { Instruction }): { string }
	local out: { string } = {}
	for i, inst in ipairs(code) do
		local ln = inst.lineno and ("[" .. tostring(inst.lineno) .. "] ") or ""
		out[#out + 1] = string.format("%03d %s%s", i - 1, ln, Instructions.render(inst))
	end
	return out
end

-- Clone instruction (useful for patching / optimization passes later)
function Instructions.clone(inst: Instruction): Instruction
	return { opcode = inst.opcode, arg = inst.arg, lineno = inst.lineno }
end

return Instructions


local nodes = require("../ast/nodes")
local opcodes = require("../bytecode/opcodes")
local State = require("./state")

type CompilerState = State.CompilerState
type ExprContext = nodes.ExprContext
type ListComp = nodes.ListComp
type BinOpType = nodes.BinOpType
type CmpOp = nodes.CmpOp

local function deepCopy(value: any): any
	if type(value) ~= "table" then
		return value
	end
	local result = {}
	for k, v in pairs(value) do
		result[k] = deepCopy(v)
	end
	return result
end

local function cloneExpr(expr: any): any
	return deepCopy(expr)
end

local function setStoreContext(node: any)
	local t = node.type
	if t == "Name" or t == "Attribute" or t == "Subscript" then
		node.ctx = "Store"
	elseif t == "Tuple" or t == "List" then
		node.ctx = "Store"
		for _, elt in ipairs(node.elts) do
			setStoreContext(elt)
		end
	elseif t == "Starred" then
		node.ctx = "Store"
		setStoreContext(node.value)
	end
end

local function makeNameExpr(id: string, ctx: ExprContext, lineno: number, col: number)
	return { type = "Name", id = id, ctx = ctx, lineno = lineno, col_offset = col }
end

local function buildListCompStatements(node: ListComp, resultName: string): { any }
	local stmts: { any } = {}

	stmts[#stmts + 1] = {
		type = "Assign",
		targets = {
			{
				type = "Name",
				id = resultName,
				ctx = "Store",
				lineno = node.lineno,
				col_offset = node.col_offset,
			},
		},
		value = { type = "List", elts = {}, ctx = "Load", lineno = node.lineno, col_offset = node.col_offset },
		lineno = node.lineno,
		col_offset = node.col_offset,
	}

	local function appendCall(): any
		local valueExpr = cloneExpr(node.elt)
		local attr = {
			type = "Attribute",
			value = makeNameExpr(resultName, "Load", node.lineno, node.col_offset),
			attr = "append",
			ctx = "Load",
			lineno = node.lineno,
			col_offset = node.col_offset,
		}
		return {
			type = "Call",
			func = attr,
			args = { valueExpr },
			keywords = {},
			lineno = valueExpr.lineno or node.lineno,
			col_offset = valueExpr.col_offset or node.col_offset,
		}
	end

	local function wrapIfs(ifs: { any }, inner: { any }): { any }
		local body = inner
		for i = #ifs, 1, -1 do
			local cond = cloneExpr(ifs[i])
			body = {
				{
					type = "If",
					test = cond,
					body = body,
					orelse = {},
					lineno = cond.lineno or node.lineno,
					col_offset = cond.col_offset or node.col_offset,
				},
			}
		end
		return body
	end

	local function buildGenerator(index: number): { any }
		local gen = node.generators[index]
		local targetClone = cloneExpr(gen.target)
		setStoreContext(targetClone)
		local iterClone = cloneExpr(gen.iter)
		local innerBody: { any }
		if index == #node.generators then
			local callExpr = appendCall()
			innerBody = {
				{ type = "Expr", value = callExpr, lineno = callExpr.lineno, col_offset = callExpr.col_offset },
			}
		else
			innerBody = buildGenerator(index + 1)
		end
		if gen.ifs and #gen.ifs > 0 then
			innerBody = wrapIfs(gen.ifs, innerBody)
		end
		return {
			{
				type = "For",
				target = targetClone,
				iter = iterClone,
				body = innerBody,
				orelse = {},
				lineno = targetClone.lineno or node.lineno,
				col_offset = targetClone.col_offset or node.col_offset,
			},
		}
	end

	local loopStmts = buildGenerator(1)
	for _, stmt in ipairs(loopStmts) do
		stmts[#stmts + 1] = stmt
	end

	stmts[#stmts + 1] = {
		type = "Return",
		value = makeNameExpr(resultName, "Load", node.lineno, node.col_offset),
		lineno = node.lineno,
		col_offset = node.col_offset,
	}

	return stmts
end

return function(deps: { compileStmt: (CompilerState, any) -> () })
	local compileStmt = deps.compileStmt

	local BINOP_TO_OPCODE: { [BinOpType]: opcodes.Opcode } = {
		Add = "BINARY_ADD",
		Sub = "BINARY_SUBTRACT",
		Mult = "BINARY_MULTIPLY",
		Div = "BINARY_DIVIDE",
		Mod = "BINARY_MODULO",
		Pow = "BINARY_POWER",
		LShift = "BINARY_LSHIFT",
		RShift = "BINARY_RSHIFT",
		BitOr = "BINARY_OR",
		BitXor = "BINARY_XOR",
		BitAnd = "BINARY_AND",
		FloorDiv = "BINARY_FLOOR_DIVIDE",
		MatMult = "BINARY_MATRIX_MULTIPLY",
	}

	local COMPARE_TO_OPCODE: { [CmpOp]: number } = {
		Lt = 0,
		LtE = 1,
		Eq = 2,
		NotEq = 3,
		Gt = 4,
		GtE = 5,
		Is = 6,
		IsNot = 7,
		In = 8,
		NotIn = 9,
	}

	local compileExpr: (CompilerState, any) -> ()

	local function compileListCompExpr(st: CompilerState, node: ListComp)
		local compState = State.newState()
		compState.code.name = "<listcomp>"
		compState.code.firstlineno = node.lineno
		compState.code.argcount = 0
		compState.code.kwonlyargcount = 0
		compState.code.varnames = {}

		local resultName = State.newTempName(compState, "__listcomp_result_")
		local stmts = buildListCompStatements(node, resultName)
		for _, stmt in ipairs(stmts) do
			compileStmt(compState, stmt)
		end

		local codeIdx = State.addConst(st, compState.code)
		State.emit(st, "LOAD_CONST", codeIdx, node.lineno)
		State.emit(st, "MAKE_FUNCTION", 0, node.lineno)
		State.emit(st, "CALL_FUNCTION", 0, node.lineno)
	end

	local function compileLambdaExpr(st: CompilerState, node: nodes.Lambda)
		local lambdaState = State.newState()
		lambdaState.code.name = "<lambda>"
		lambdaState.code.firstlineno = node.lineno
		lambdaState.code.argcount = if node.args and node.args.args then #node.args.args else 0
		lambdaState.code.kwonlyargcount = 0
		lambdaState.code.varnames = {}

		if node.args and node.args.args then
			for _, arg in ipairs(node.args.args) do
				lambdaState.code.varnames[#lambdaState.code.varnames + 1] = arg.arg
			end
		end

		compileExpr(lambdaState, node.body)
		local retLine = node.body.lineno or node.lineno
		State.emit(lambdaState, "RETURN_VALUE", nil, retLine)

		local codeIdx = State.addConst(st, lambdaState.code)
		State.emit(st, "LOAD_CONST", codeIdx, node.lineno)

		local defaults = if node.args and node.args.defaults then node.args.defaults else {}
		local flags = 0
		if defaults and #defaults > 0 then
			for _, defaultExpr in ipairs(defaults) do
				compileExpr(st, defaultExpr)
			end
			State.emit(st, "BUILD_TUPLE", #defaults, node.lineno)
			flags += 0x01
		end

		State.emit(st, "MAKE_FUNCTION", flags, node.lineno)
	end

	local function withContext(node: any, ctx: ExprContext, fn: () -> ())
		local prev = node.ctx
		node.ctx = ctx
		fn()
		node.ctx = prev
	end

	local function compileName(st: CompilerState, node: any)
		if node.ctx == "Load" then
			local idx = State.addName(st, node.id)
			State.emit(st, "LOAD_NAME", idx, node.lineno)
		elseif node.ctx == "Store" then
			local idx = State.addName(st, node.id)
			State.emit(st, "STORE_NAME", idx, node.lineno)
		else
			error("Unsupported Name ctx: " .. tostring(node.ctx))
		end
	end

	local function compileAssignmentTarget(st: CompilerState, target: any)
		if target.type == "Name" then
			target.ctx = "Store"
			compileName(st, target)
		elseif target.type == "Tuple" then
			target.ctx = "Store"
			local count = #target.elts
			State.emit(st, "UNPACK_SEQUENCE", count, target.lineno)
			for _, elt in ipairs(target.elts) do
				compileAssignmentTarget(st, elt)
			end
		elseif target.type == "Attribute" then
			compileExpr(st, target.value)
			local idx = State.addName(st, target.attr)
			State.emit(st, "STORE_ATTR", idx, target.lineno)
		elseif target.type == "Subscript" then
			compileExpr(st, target.value)
			compileExpr(st, target.slice)
			State.emit(st, "STORE_SUBSCR", nil, target.lineno)
		else
			error("Unsupported assignment target type: " .. tostring(target.type))
		end
	end

	local function compileLoadTarget(st: CompilerState, target: any)
		if target.type == "Name" then
			withContext(target, "Load", function()
				compileName(st, target)
			end)
		elseif target.type == "Attribute" then
			compileExpr(st, target.value)
			local idx = State.addName(st, target.attr)
			State.emit(st, "LOAD_ATTR", idx, target.lineno)
		elseif target.type == "Subscript" then
			compileExpr(st, target.value)
			compileExpr(st, target.slice)
			State.emit(st, "LOAD_SUBSCR", nil, target.lineno)
		else
			error("Unsupported load target type: " .. tostring(target.type))
		end
	end

	local function compileStoreTarget(st: CompilerState, target: any)
		if target.type == "Name" then
			withContext(target, "Store", function()
				compileName(st, target)
			end)
		elseif target.type == "Attribute" then
			compileExpr(st, target.value)
			State.emit(st, "ROT_TWO", nil, target.lineno)
			local idx = State.addName(st, target.attr)
			State.emit(st, "STORE_ATTR", idx, target.lineno)
		elseif target.type == "Subscript" then
			compileExpr(st, target.value)
			State.emit(st, "ROT_TWO", nil, target.lineno)
			compileExpr(st, target.slice)
			State.emit(st, "ROT_TWO", nil, target.lineno)
			State.emit(st, "STORE_SUBSCR", nil, target.lineno)
		else
			error("Unsupported store target type: " .. tostring(target.type))
		end
	end

	compileExpr = function(st: CompilerState, node: any)
		local t = node.type
		if t == "Constant" then
			if node.kind == "bytes" then
				local tagged = { __bytes_literal = true, value = node.value }
				local idx = State.addConst(st, tagged)
				State.emit(st, "LOAD_CONST", idx, node.lineno)
			else
				local idx = State.addConst(st, node.value)
				State.emit(st, "LOAD_CONST", idx, node.lineno)
			end
		elseif t == "Name" then
			compileName(st, node)
		elseif t == "Attribute" then
			compileExpr(st, node.value)
			local idx = State.addName(st, node.attr)
			State.emit(st, "LOAD_ATTR", idx, node.lineno)
		elseif t == "Subscript" then
			compileExpr(st, node.value)
			compileExpr(st, node.slice)
			State.emit(st, "LOAD_SUBSCR", nil, node.lineno)
		elseif t == "BinOp" then
			compileExpr(st, node.left)
			compileExpr(st, node.right)
			local op = BINOP_TO_OPCODE[node.op]
			if not op then
				error("Unsupported binary op: " .. tostring(node.op))
			end
			State.emit(st, (op :: any) :: opcodes.Opcode, nil, node.lineno)
		elseif t == "UnaryOp" then
			compileExpr(st, node.operand)
			if node.op == "UAdd" then
				State.emit(st, "UNARY_POSITIVE", nil, node.lineno)
			elseif node.op == "USub" then
				State.emit(st, "UNARY_NEGATIVE", nil, node.lineno)
			elseif node.op == "Not" then
				State.emit(st, "UNARY_NOT", nil, node.lineno)
			elseif node.op == "Invert" then
				State.emit(st, "UNARY_INVERT", nil, node.lineno)
			else
				error("Unsupported unary op: " .. tostring(node.op))
			end
		elseif t == "Compare" then
			local opCount = #node.ops
			local comparatorCount = #node.comparators
			if opCount ~= comparatorCount or opCount == 0 then
				error("Invalid comparison operands")
			end

			compileExpr(st, node.left)
			local falseJumpFixups: { number } = {}

			for index, op in ipairs(node.ops) do
				local comparator = node.comparators[index]
				local isLast = index == opCount
				local lineno = (comparator and comparator.lineno) or node.lineno

				compileExpr(st, comparator)
				if not isLast then
					State.emit(st, "DUP_TOP", nil, lineno)
					State.emit(st, "ROT_THREE", nil, lineno)
				end

				local opcode = COMPARE_TO_OPCODE[op]
				if opcode == nil then
					error("Unsupported comparison op: " .. tostring(op))
				end
				State.emit(st, "COMPARE_OP", opcode, lineno)

				if not isLast then
					local jumpAddr = State.getCurrentAddress(st)
					State.emit(st, "JUMP_IF_FALSE_OR_POP", 0, lineno)
					falseJumpFixups[#falseJumpFixups + 1] = jumpAddr
				end
			end

			if #falseJumpFixups > 0 then
				local jumpPastCleanup = State.getCurrentAddress(st)
				State.emit(st, "JUMP_FORWARD", 0, node.lineno)

				local cleanupAddr = State.getCurrentAddress(st)
				State.emit(st, "ROT_TWO", nil, node.lineno)
				State.emit(st, "POP_TOP", nil, node.lineno)
				local doneAddr = State.getCurrentAddress(st)

				for _, addr in ipairs(falseJumpFixups) do
					State.patchJump(st, addr, cleanupAddr)
				end
				State.patchJump(st, jumpPastCleanup, doneAddr)
			end
		elseif t == "Call" then
			compileExpr(st, node.func)
			for _, arg in ipairs(node.args) do
				compileExpr(st, arg)
			end
			State.emit(st, "CALL_FUNCTION", #node.args, node.lineno)
		elseif t == "ListComp" then
			compileListCompExpr(st, node)
		elseif t == "Lambda" then
			compileLambdaExpr(st, node)
		elseif t == "Yield" then
			if not st.isGenerator then
				error("SyntaxError: 'yield' outside function", 2)
			end
			if node.value then
				compileExpr(st, node.value)
			else
				local idxNone = State.addConst(st, nil)
				State.emit(st, "LOAD_CONST", idxNone, node.lineno)
			end
			State.emit(st, "YIELD_VALUE", nil, node.lineno)
		elseif t == "YieldFrom" then
			if not st.isGenerator then
				error("SyntaxError: 'yield' outside function", 2)
			end
			compileExpr(st, node.value)
			State.emit(st, "GET_ITER", nil, node.lineno)
			local loopStart = State.getCurrentAddress(st)
			local forIterJump = State.getCurrentAddress(st)
			State.emit(st, "FOR_ITER", 0, node.lineno)
			State.emit(st, "YIELD_VALUE", nil, node.lineno)
			local currentAddr = State.getCurrentAddress(st)
			local backOffset = loopStart - currentAddr - 1
			State.emit(st, "JUMP_FORWARD", backOffset, node.lineno)
			local loopEnd = State.getCurrentAddress(st)
			State.patchJump(st, forIterJump, loopEnd)
			local idxNone = State.addConst(st, nil)
			State.emit(st, "LOAD_CONST", idxNone, node.lineno)
		elseif t == "List" then
			for _, elt in ipairs(node.elts) do
				compileExpr(st, elt)
			end
			State.emit(st, "BUILD_LIST", #node.elts, node.lineno)
		elseif t == "Tuple" then
			for _, elt in ipairs(node.elts) do
				compileExpr(st, elt)
			end
			State.emit(st, "BUILD_TUPLE", #node.elts, node.lineno)
		elseif t == "Set" then
			for _, elt in ipairs(node.elts) do
				compileExpr(st, elt)
			end
			State.emit(st, "BUILD_SET", #node.elts, node.lineno)
		elseif t == "Dict" then
			local count = #node.keys
			local hasUnpack = false
			for i = 1, count do
				local key = node.keys[i]
				if key and key.type == "DictUnpack" then
					hasUnpack = true
					break
				end
			end

			if not hasUnpack then
				for i = 1, count do
					local key = node.keys[i]
					local value = node.values[i]
					compileExpr(st, key)
					compileExpr(st, value)
				end
				State.emit(st, "BUILD_MAP", count, node.lineno)
			else
				local segments = 0
				local currentPairs: { { key: any, value: any } } = {}

				local function flushCurrentPairs()
					if #currentPairs == 0 then
						return
					end
					for _, pair in ipairs(currentPairs) do
						compileExpr(st, pair.key)
						compileExpr(st, pair.value)
					end
					State.emit(st, "BUILD_MAP", #currentPairs, node.lineno)
					segments += 1
					currentPairs = {}
				end

				for i = 1, count do
					local key = node.keys[i]
					local value = node.values[i]
					if key and key.type == "DictUnpack" then
						flushCurrentPairs()
						compileExpr(st, value)
						segments += 1
					else
						currentPairs[#currentPairs + 1] = { key = key, value = value }
					end
				end

				flushCurrentPairs()

				if segments > 0 then
					State.emit(st, "BUILD_MAP_UNPACK", segments, node.lineno)
				else
					State.emit(st, "BUILD_MAP", 0, node.lineno)
				end
			end
		else
			error("Unsupported expression node type: " .. tostring(t))
		end
	end

	local function resolveBinOpOpcode(op: BinOpType)
		return BINOP_TO_OPCODE[op]
	end

	return {
		compileExpr = compileExpr,
		compileAssignmentTarget = compileAssignmentTarget,
		compileLoadTarget = compileLoadTarget,
		compileStoreTarget = compileStoreTarget,
		resolveBinOpOpcode = resolveBinOpOpcode,
	}
end


local State = require("@self/state")
local expressionsFactory = require("@self/expressions")
local statementsFactory = require("@self/statements")

export type CodeObject = State.CodeObject

local Compiler = {}

local compileStmtImpl: ((State.CompilerState, any) -> ())? = nil

local function compileStmt(st: State.CompilerState, node: any)
	if not compileStmtImpl then
		error("Compiler pipeline not initialized", 2)
	end
	compileStmtImpl(st, node)
end

local expressions = expressionsFactory({
	compileStmt = compileStmt,
})

local statements = statementsFactory({
	compileExpr = expressions.compileExpr,
	compileAssignmentTarget = expressions.compileAssignmentTarget,
	compileLoadTarget = expressions.compileLoadTarget,
	compileStoreTarget = expressions.compileStoreTarget,
	resolveBinOpOpcode = expressions.resolveBinOpOpcode,
})

compileStmtImpl = statements.compileStmt

local function compileModule(ast: any): CodeObject
	local st = State.newState()
	st.code.firstlineno = ast.lineno
	for _, stmt in ipairs(ast.body) do
		compileStmt(st, stmt)
	end
	local idxNone = State.addConst(st, nil)
	State.emit(st, "LOAD_CONST", idxNone, ast.lineno)
	State.emit(st, "RETURN_VALUE", nil, ast.lineno)
	return st.code
end

function Compiler.compile(ast: any): CodeObject
	if ast.type == "Module" then
		return compileModule(ast)
	end

	local exprNode = ast
	if ast.type == "Expr" and ast.value then
		exprNode = ast.value
	end

	if exprNode and exprNode.type then
		local lineno = exprNode.lineno or 1
		local col = exprNode.col_offset or 0
		local fakeModule = {
			type = "Module",
			body = {
				{
					type = "Return",
					value = exprNode,
					lineno = lineno,
					col_offset = col,
				},
			},
			lineno = lineno,
			col_offset = col,
		}
		return compileModule(fakeModule)
	end

	error("Unsupported root AST for compile: " .. tostring(ast.type))
end

return Compiler


local instructions = require("../bytecode/instructions")
local opcodes = require("../bytecode/opcodes")

type Instruction = instructions.Instruction
type Opcode = opcodes.Opcode

export type CodeObject = {
	constants: { any },
	names: { string },
	varnames: { string },
	bytecode: { Instruction },
	firstlineno: number?,
	lnotab: { number }?,
	name: string?,
	argcount: number?,
	kwonlyargcount: number?,
	isGenerator: boolean?,
}

export type LoopContext = {
	startLabel: number,
	endLabel: number,
}

export type CompilerState = {
	code: CodeObject,
	constIndex: { [string]: number },
	nameIndex: { [string]: number },
	loopStack: { LoopContext },
	labelCounter: number,
	tempNameCounter: number,
	isGenerator: boolean,
}

local State = {}

local function serializeConst(value: any): string
	local t = typeof(value)
	if t == "nil" then
		return "N"
	elseif t == "number" then
		return "F" .. tostring(value)
	elseif t == "boolean" then
		return value and "B1" or "B0"
	elseif t == "string" then
		return "S" .. value
	else
		return "O" .. tostring(value)
	end
end

function State.newState(): CompilerState
	return {
		code = {
			constants = {},
			names = {},
			varnames = {},
			bytecode = {},
			firstlineno = nil,
			lnotab = nil,
			name = nil,
			argcount = nil,
			kwonlyargcount = nil,
			isGenerator = false,
		},
		constIndex = {},
		nameIndex = {},
		loopStack = {},
		labelCounter = 0,
		tempNameCounter = 0,
		isGenerator = false,
	}
end

function State.newTempName(st: CompilerState, prefix: string): string
	st.tempNameCounter += 1
	return prefix .. tostring(st.tempNameCounter)
end

function State.addConst(st: CompilerState, value: any): number
	local key = serializeConst(value)
	local idx = st.constIndex[key]
	if idx then
		return idx
	end

	local storedValue = value
	if value == nil then
		storedValue = { __nil = true }
	end

	table.insert(st.code.constants, storedValue)
	local newIdx = #st.code.constants - 1
	st.constIndex[key] = newIdx
	return newIdx
end

function State.addName(st: CompilerState, name: string): number
	local idx = st.nameIndex[name]
	if idx then
		return idx
	end
	table.insert(st.code.names, name)
	local newIdx = #st.code.names - 1
	st.nameIndex[name] = newIdx
	return newIdx
end

function State.emit(st: CompilerState, op: Opcode, arg: number?, lineno: number?)
	local inst = instructions.make(op, arg, lineno)
	st.code.bytecode[#st.code.bytecode + 1] = inst
end

function State.getCurrentAddress(st: CompilerState): number
	return #st.code.bytecode
end

function State.patchJump(st: CompilerState, jumpAddr: number, targetAddr: number?)
	local target = targetAddr or State.getCurrentAddress(st)
	local instruction = st.code.bytecode[jumpAddr + 1]
	if not instruction then
		return
	end

	if instruction.opcode == "JUMP_FORWARD" then
		local targetIndex = target + 1
		instruction.arg = targetIndex - (jumpAddr + 2)
	else
		instruction.arg = target + 1
	end
end

function State.enterLoop(st: CompilerState, startLabel: number, endLabel: number)
	table.insert(st.loopStack, { startLabel = startLabel, endLabel = endLabel })
end

function State.exitLoop(st: CompilerState)
	table.remove(st.loopStack)
end

function State.getCurrentLoop(st: CompilerState): LoopContext?
	return st.loopStack[#st.loopStack]
end

return State


local State = require("./state")

type CompilerState = State.CompilerState

type StatementDeps = {
	compileExpr: (CompilerState, any) -> (),
	compileAssignmentTarget: (CompilerState, any) -> (),
	compileLoadTarget: (CompilerState, any) -> (),
	compileStoreTarget: (CompilerState, any) -> (),
	resolveBinOpOpcode: (any) -> string?,
}

return function(deps: StatementDeps)
	local compileExpr = deps.compileExpr
	local compileAssignmentTarget = deps.compileAssignmentTarget
	local compileLoadTarget = deps.compileLoadTarget
	local compileStoreTarget = deps.compileStoreTarget
	local resolveBinOpOpcode = deps.resolveBinOpOpcode

	local compileStmt: (CompilerState, any) -> ()

	local function nodeContainsYield(node: any): boolean
		if type(node) ~= "table" then
			return false
		end
		local nodeType = node.type
		if nodeType == "Yield" or nodeType == "YieldFrom" then
			return true
		end
		if nodeType == "FunctionDef" or nodeType == "ClassDef" or nodeType == "Lambda" then
			return false
		end
		for _, value in pairs(node) do
			if type(value) == "table" then
				if value.type then
					if nodeContainsYield(value) then
						return true
					end
				else
					for _, item in ipairs(value) do
						if nodeContainsYield(item) then
							return true
						end
					end
				end
			end
		end
		return false
	end

	local function compileTryNoFinally(st: CompilerState, node: any)
		if not node.handlers or #node.handlers == 0 then
			error("try statement must have at least one handler", 2)
		end

		local setupIndex = State.getCurrentAddress(st)
		State.emit(st, "SETUP_EXCEPT", 0, node.lineno)

		for _, stmt in ipairs(node.body) do
			compileStmt(st, stmt)
		end

		State.emit(st, "POP_BLOCK", nil, node.lineno)

		if node.orelse and #node.orelse > 0 then
			for _, stmt in ipairs(node.orelse) do
				compileStmt(st, stmt)
			end
		end

		local jumpPastHandlers = State.getCurrentAddress(st)
		State.emit(st, "JUMP_FORWARD", 0, node.lineno)

		local handlerStarts = {}
		local handlerFallthroughJumps = {}
		local matchFailJumps = {}

		for index, handler in ipairs(node.handlers) do
			local handlerStart = State.getCurrentAddress(st)
			handlerStarts[index] = handlerStart

			local matchJumpIndex: number? = nil
			if handler.expr then
				State.emit(st, "DUP_TOP", nil, handler.lineno)
				local matchNameIdx = State.addName(st, "_pylua_match_exception")
				State.emit(st, "LOAD_NAME", matchNameIdx, handler.lineno)
				State.emit(st, "ROT_TWO", nil, handler.lineno)
				compileExpr(st, handler.expr)
				State.emit(st, "CALL_FUNCTION", 2, handler.lineno)
				matchJumpIndex = State.getCurrentAddress(st)
				State.emit(st, "POP_JUMP_IF_FALSE", 0, handler.lineno)
			end
			if matchJumpIndex then
				matchFailJumps[#matchFailJumps + 1] = { index = index, jump = matchJumpIndex }
			end

			if handler.name then
				local storeIdx = State.addName(st, handler.name)
				State.emit(st, "STORE_NAME", storeIdx, handler.lineno)
			else
				State.emit(st, "POP_TOP", nil, handler.lineno)
			end

			for _, stmt in ipairs(handler.body) do
				compileStmt(st, stmt)
			end

			local exitJump = State.getCurrentAddress(st)
			State.emit(st, "JUMP_FORWARD", 0, handler.lineno)
			handlerFallthroughJumps[#handlerFallthroughJumps + 1] = exitJump
		end

		local unhandledLabel = State.getCurrentAddress(st)
		if #node.handlers > 0 then
			local raiseIdx = State.addName(st, "_pylua_raise")
			State.emit(st, "LOAD_NAME", raiseIdx, node.lineno)
			State.emit(st, "ROT_TWO", nil, node.lineno)
			State.emit(st, "CALL_FUNCTION", 1, node.lineno)
			State.emit(st, "POP_TOP", nil, node.lineno)
		end

		local endLabel = State.getCurrentAddress(st)

		State.patchJump(st, setupIndex, handlerStarts[1])
		State.patchJump(st, jumpPastHandlers, endLabel)

		for _, record in ipairs(matchFailJumps) do
			local target: number? = nil
			for nextIndex = record.index + 1, #handlerStarts do
				local start = handlerStarts[nextIndex]
				if start ~= nil then
					target = start
					break
				end
			end
			State.patchJump(st, record.jump, target or unhandledLabel)
		end

		for _, jump in ipairs(handlerFallthroughJumps) do
			State.patchJump(st, jump, endLabel)
		end
	end

	local function compileTryWithFinally(st: CompilerState, node: any)
		local lineno = node.lineno
		local setupIndex = State.getCurrentAddress(st)
		State.emit(st, "SETUP_FINALLY", 0, lineno)

		if node.handlers and #node.handlers > 0 then
			local inner = {
				type = "Try",
				body = node.body,
				handlers = node.handlers,
				orelse = node.orelse,
				finalbody = {},
				lineno = node.lineno,
				col_offset = node.col_offset,
			}
			compileTryNoFinally(st, inner)
		else
			for _, stmt in ipairs(node.body) do
				compileStmt(st, stmt)
			end
			if node.orelse and #node.orelse > 0 then
				for _, stmt in ipairs(node.orelse) do
					compileStmt(st, stmt)
				end
			end
		end

		State.emit(st, "POP_BLOCK", nil, lineno)
		State.emit(st, "BEGIN_FINALLY", nil, lineno)

		local finalStart = State.getCurrentAddress(st)
		State.patchJump(st, setupIndex, finalStart)

		for _, stmt in ipairs(node.finalbody) do
			compileStmt(st, stmt)
		end

		State.emit(st, "END_FINALLY", nil, lineno)
	end

	local function compileFunctionDef(st: CompilerState, node: any)
		local funcName = node.name
		local args = node.args
		local body = node.body

		local funcState = State.newState()
		funcState.code.name = funcName
		funcState.code.firstlineno = node.lineno
		local isGenerator = false
		if body then
			for _, stmt in ipairs(body) do
				if nodeContainsYield(stmt) then
					isGenerator = true
					break
				end
			end
		end
		funcState.isGenerator = isGenerator
		funcState.code.isGenerator = isGenerator

		local argcount = 0
		local varnames: { string } = {}

		if args and args.args then
			for _, arg in ipairs(args.args) do
				argcount += 1
				varnames[#varnames + 1] = arg.arg or arg.name or ("arg" .. tostring(argcount))
			end
		end

		funcState.code.argcount = argcount
		funcState.code.varnames = varnames

		for _, stmt in ipairs(body) do
			compileStmt(funcState, stmt)
		end

		local idxNone = State.addConst(funcState, nil)
		State.emit(funcState, "LOAD_CONST", idxNone, node.lineno)
		State.emit(funcState, "RETURN_VALUE", nil, node.lineno)

		local funcCodeIdx = State.addConst(st, funcState.code)
		State.emit(st, "LOAD_CONST", funcCodeIdx, node.lineno)

		local defaultCount = 0
		local flags = 0
		if args and args.defaults and #args.defaults > 0 then
			for _, defaultExpr in ipairs(args.defaults) do
				compileExpr(st, defaultExpr)
				defaultCount += 1
			end
			State.emit(st, "BUILD_TUPLE", defaultCount, node.lineno)
			flags += 0x01
		end

		State.emit(st, "MAKE_FUNCTION", flags, node.lineno)

		local nameIdx = State.addName(st, funcName)
		State.emit(st, "STORE_NAME", nameIdx, node.lineno)
	end

	local function compileClassDef(st: CompilerState, node: any)
		local classState = State.newState()
		classState.code.name = node.name
		classState.code.firstlineno = node.lineno
		classState.code.argcount = 0
		classState.code.varnames = {}
		for _, stmt in ipairs(node.body) do
			compileStmt(classState, stmt)
		end
		local idxNone = State.addConst(classState, nil)
		State.emit(classState, "LOAD_CONST", idxNone, node.lineno)
		State.emit(classState, "RETURN_VALUE", nil, node.lineno)

		local codeIdx = State.addConst(st, classState.code)
		State.emit(st, "LOAD_CONST", codeIdx, node.lineno)
		State.emit(st, "MAKE_FUNCTION", 0, node.lineno)

		local nameIdx = State.addConst(st, node.name)
		State.emit(st, "LOAD_CONST", nameIdx, node.lineno)

		for _, base in ipairs(node.bases) do
			compileExpr(st, base)
		end

		State.emit(st, "BUILD_CLASS", #node.bases, node.lineno)

		local storeIdx = State.addName(st, node.name)
		State.emit(st, "STORE_NAME", storeIdx, node.lineno)
	end

	local function compileImport(st: CompilerState, node: any)
		for _, alias in ipairs(node.names) do
			local levelIdx = State.addConst(st, 0)
			State.emit(st, "LOAD_CONST", levelIdx, node.lineno)

			local noneIdx = State.addConst(st, nil)
			State.emit(st, "LOAD_CONST", noneIdx, node.lineno)

			local nameIdx = State.addName(st, alias.name)
			State.emit(st, "IMPORT_NAME", nameIdx, node.lineno)

			if alias.asname then
				-- import a.b.c as d
				-- IMPORT_NAME returns 'a'. We need to traverse to 'c'.
				local parts = {}
				for part in string.gmatch(alias.name, "[^%.]+") do
					table.insert(parts, part)
				end

				for i = 2, #parts do
					local partIdx = State.addName(st, parts[i])
					State.emit(st, "IMPORT_FROM", partIdx, node.lineno)
					State.emit(st, "ROT_TWO", nil, node.lineno)
					State.emit(st, "POP_TOP", nil, node.lineno)
				end

				local storeIdx = State.addName(st, alias.asname)
				State.emit(st, "STORE_NAME", storeIdx, node.lineno)
			else
				-- import a.b.c
				-- IMPORT_NAME returns 'a'. We store 'a'.
				local parts = {}
				for part in string.gmatch(alias.name, "[^%.]+") do
					table.insert(parts, part)
				end
				local topName = parts[1]
				local storeIdx = State.addName(st, topName)
				State.emit(st, "STORE_NAME", storeIdx, node.lineno)
			end
		end
	end

	local function compileImportFrom(st: CompilerState, node: any)
		local levelIdx = State.addConst(st, node.level or 0)
		State.emit(st, "LOAD_CONST", levelIdx, node.lineno)

		local names = {}
		for _, alias in ipairs(node.names) do
			table.insert(names, alias.name)
		end
		-- Create a tuple for fromlist
		local namesTuple = { __type = "tuple", __value = names }
		local namesIdx = State.addConst(st, namesTuple)
		State.emit(st, "LOAD_CONST", namesIdx, node.lineno)

		local moduleName = node.module or ""
		local modIdx = State.addName(st, moduleName)
		State.emit(st, "IMPORT_NAME", modIdx, node.lineno)

		for _, alias in ipairs(node.names) do
			if alias.name == "*" then
				State.emit(st, "IMPORT_STAR", nil, node.lineno)
			else
				local nameIdx = State.addName(st, alias.name)
				State.emit(st, "IMPORT_FROM", nameIdx, node.lineno)
				local storeName = alias.asname or alias.name
				local storeIdx = State.addName(st, storeName)
				State.emit(st, "STORE_NAME", storeIdx, node.lineno)
			end
		end

		State.emit(st, "POP_TOP", nil, node.lineno)
	end

	compileStmt = function(st: CompilerState, node: any)
		if node.type == "Expr" then
			compileExpr(st, node.value)
			State.emit(st, "POP_TOP", nil, node.lineno)
		elseif node.type == "Assign" then
			compileExpr(st, node.value)
			local target = node.targets[1]
			if not target then
				error("Assign without target")
			end
			compileStoreTarget(st, target)
		elseif node.type == "AugAssign" then
			local target = node.target
			compileLoadTarget(st, target)
			compileExpr(st, node.value)
			local op = resolveBinOpOpcode(node.op)
			if not op then
				error("Unsupported augmented assignment operator: " .. tostring(node.op))
			end
			State.emit(st, op :: any, nil, node.lineno)
			compileStoreTarget(st, target)
		elseif node.type == "Return" then
			if node.value then
				compileExpr(st, node.value)
			else
				local idx = State.addConst(st, nil)
				State.emit(st, "LOAD_CONST", idx, node.lineno)
			end
			State.emit(st, "RETURN_VALUE", nil, node.lineno)
		elseif node.type == "If" then
			compileExpr(st, node.test)
			local jumpIfFalse = State.getCurrentAddress(st)
			State.emit(st, "POP_JUMP_IF_FALSE", 0, node.lineno)

			for _, stmt in ipairs(node.body) do
				compileStmt(st, stmt)
			end

			if #node.orelse > 0 then
				local jumpToEnd = State.getCurrentAddress(st)
				State.emit(st, "JUMP_FORWARD", 0, node.lineno)
				State.patchJump(st, jumpIfFalse, State.getCurrentAddress(st))

				for _, stmt in ipairs(node.orelse) do
					compileStmt(st, stmt)
				end

				State.patchJump(st, jumpToEnd, State.getCurrentAddress(st))
			else
				State.patchJump(st, jumpIfFalse, State.getCurrentAddress(st))
			end
		elseif node.type == "While" then
			local setupLoop = State.getCurrentAddress(st)
			State.emit(st, "SETUP_LOOP", 0, node.lineno)

			local loopStart = State.getCurrentAddress(st)
			compileExpr(st, node.test)
			local jumpIfFalse = State.getCurrentAddress(st)
			State.emit(st, "POP_JUMP_IF_FALSE", 0, node.lineno)

			State.enterLoop(st, loopStart, jumpIfFalse)

			for _, stmt in ipairs(node.body) do
				compileStmt(st, stmt)
			end

			local currentAddr = State.getCurrentAddress(st)
			local relativeOffset = loopStart - currentAddr - 1
			State.emit(st, "JUMP_FORWARD", relativeOffset, node.lineno)

			local loopEnd = State.getCurrentAddress(st)
			State.patchJump(st, jumpIfFalse, loopEnd)

			if #node.orelse > 0 then
				for _, stmt in ipairs(node.orelse) do
					compileStmt(st, stmt)
				end
			end

			local blockEnd = State.getCurrentAddress(st)
			State.emit(st, "POP_BLOCK", nil, node.lineno)
			State.patchJump(st, setupLoop, blockEnd)

			State.exitLoop(st)
		elseif node.type == "For" then
			compileExpr(st, node.iter)
			State.emit(st, "GET_ITER", nil, node.lineno)

			local setupLoop = State.getCurrentAddress(st)
			State.emit(st, "SETUP_LOOP", 0, node.lineno)

			local loopStart = State.getCurrentAddress(st)

			local forIterJump = State.getCurrentAddress(st)
			State.emit(st, "FOR_ITER", 0, node.lineno)

			State.enterLoop(st, loopStart, forIterJump)

			compileAssignmentTarget(st, node.target)

			for _, stmt in ipairs(node.body) do
				compileStmt(st, stmt)
			end

			local currentAddr = State.getCurrentAddress(st)
			local relativeOffset = loopStart - currentAddr - 1
			State.emit(st, "JUMP_FORWARD", relativeOffset, node.lineno)

			local loopEnd = State.getCurrentAddress(st)
			State.patchJump(st, forIterJump, loopEnd)

			if #node.orelse > 0 then
				for _, stmt in ipairs(node.orelse) do
					compileStmt(st, stmt)
				end
			end

			local blockEnd = State.getCurrentAddress(st)
			State.emit(st, "POP_BLOCK", nil, node.lineno)
			State.patchJump(st, setupLoop, blockEnd)

			State.exitLoop(st)
		elseif node.type == "Break" then
			local loop = State.getCurrentLoop(st)
			if not loop then
				error("'break' outside loop", 2)
			end
			State.emit(st, "BREAK_LOOP", nil, node.lineno)
		elseif node.type == "Continue" then
			local loop = State.getCurrentLoop(st)
			if not loop then
				error("'continue' not properly in loop", 2)
			end
			State.emit(st, "CONTINUE_LOOP", loop.startLabel, node.lineno)
		elseif node.type == "Pass" then
			return
		elseif node.type == "Try" then
			if node.finalbody and #node.finalbody > 0 then
				compileTryWithFinally(st, node)
			else
				compileTryNoFinally(st, node)
			end
		elseif node.type == "FunctionDef" then
			compileFunctionDef(st, node)
		elseif node.type == "ClassDef" then
			compileClassDef(st, node)
		elseif node.type == "Import" then
			compileImport(st, node)
		elseif node.type == "ImportFrom" then
			compileImportFrom(st, node)
		else
			error("Unsupported statement type: " .. tostring(node.type))
		end
	end

	return {
		compileStmt = compileStmt,
	}
end


local Bridge = {}

local Base = require("../objects/base")
local Collections = require("../objects/collections")

type PyObject = Base.PyObject
type PyType = Base.PyType

type VisitedPy = { [PyObject]: any }
type VisitedLuau = { [any]: PyObject }

type HashEntry = { key: PyObject, value: PyObject }

type LuauFunctionRecord = { func: (...any) -> any }

type ClassBindingOptionsInternal = {
	pyTypeName: string,
	typeofName: string?,
	predicate: ((value: any) -> boolean)?,
	bindSelf: boolean?,
	getAttr: ((target: any, attr: string) -> any?)?,
	setAttr: ((target: any, attr: string, value: any) -> ())?,
	toLuau: ((target: any, visited: VisitedPy) -> any)?,
}

type ClassBinding = {
	pyTypeName: string,
	typeofName: string?,
	predicate: ((value: any) -> boolean)?,
	bindSelf: boolean,
	getAttr: ((target: any, attr: string) -> any?)?,
	setAttr: ((target: any, attr: string, value: any) -> ())?,
	toLuau: ((target: any, visited: VisitedPy) -> any)?,
}

type BindingPayload = {
	__pylua_binding: boolean,
	target: any,
	binding: ClassBinding,
}

export type ClassBindingOptions = ClassBindingOptionsInternal

local luauFunctionType: PyType?
local classBindings: { ClassBinding } = {}
local bindingByPyType: { [string]: ClassBinding } = {}

local bindingGetAttr: (PyObject, string) -> PyObject?
local bindingSetAttr: (PyObject, string, PyObject) -> ()

local function tryTypeof(value: any): string?
	local globalEnv = _G :: any
	local typeofFn = globalEnv and globalEnv.typeof
	if typeofFn and type(typeofFn) == "function" then
		local ok, result = pcall(typeofFn, value)
		if ok and type(result) == "string" then
			return result
		end
	end
	return nil
end

local function getTypeTag(value: any): string
	local tag = tryTypeof(value)
	if tag then
		return tag
	end
	return type(value)
end

local function isPyObject(value: any): boolean
	return type(value) == "table" and value.__type ~= nil and value.__typeobj ~= nil
end

local function isSequenceTable(tbl: { [any]: any }): boolean
	local maxIndex = 0
	local count = 0
	for key, _ in pairs(tbl) do
		if type(key) ~= "number" or key % 1 ~= 0 or key < 1 then
			return false
		end
		if key > maxIndex then
			maxIndex = key
		end
		count += 1
	end
	if count == 0 then
		return true
	end
	if maxIndex ~= count then
		return false
	end
	for i = 1, maxIndex do
		if rawget(tbl, i) == nil then
			return false
		end
	end
	return true
end

local function hasMetatableClass(tbl: any): string?
	local mt = getmetatable(tbl)
	if mt and type(mt) == "table" and mt.__pylua_class then
		return mt.__pylua_class
	end
	return nil
end

local function hashKey(obj: PyObject): string
	local t = obj.__type
	if t == "int" or t == "float" or t == "bool" or t == "str" then
		return t .. ":" .. tostring(obj.__value)
	elseif t == "NoneType" then
		return "None"
	end
	return tostring(obj)
end

local function bindingMatches(binding: ClassBinding, value: any): boolean
	if binding.typeofName then
		if getTypeTag(value) ~= binding.typeofName then
			return false
		end
	end
	if binding.predicate then
		local ok, result = pcall(binding.predicate, value)
		if not ok then
			return false
		end
		return result == true
	end
	return true
end

local function safeIndex(target: any, key: string): any?
	local ok, result = pcall(function()
		return target[key]
	end)
	if ok then
		return result
	end
	return nil
end

local function safeAssign(target: any, key: string, value: any)
	local ok, err = pcall(function()
		target[key] = value
	end)
	if not ok then
		error("AttributeError: cannot set attribute '" .. key .. "': " .. tostring(err), 2)
	end
end

local function ensureLuauFunctionType(): PyType
	if luauFunctionType then
		return luauFunctionType
	end
	local ok, existing = pcall(function(): PyType
		return Base.getTypeObject("luau_function")
	end)
	local typeObj: PyType
	if ok and existing then
		typeObj = existing
	else
		typeObj = Base.registerType("luau_function", {})
	end
	typeObj.methods = typeObj.methods or {}
	local methods = typeObj.methods :: { [string]: (PyObject, { PyObject }?) -> PyObject }
	if not methods["__call__"] then
		methods["__call__"] = function(selfObj: PyObject, args: { PyObject }?): PyObject
			local record: LuauFunctionRecord = selfObj.__value
			local converted: { any } = {}
			if args then
				for index, arg in ipairs(args) do
					converted[index] = Bridge.pyToLuau(arg)
				end
			end
			local okCall, result = pcall(record.func, table.unpack(converted))
			if not okCall then
				error(result, 0)
			end
			return Bridge.luauToPy(result)
		end
	end
	luauFunctionType = typeObj
	return typeObj
end

function Bridge.wrapLuauFunction(fn: (...any) -> any): PyObject
	if type(fn) ~= "function" then
		error("wrapLuauFunction requires a Luau function", 2)
	end
	local typeObj = ensureLuauFunctionType()
	local wrapper: PyObject = {
		__type = "luau_function",
		__value = { func = fn },
		__dict = {},
		__typeobj = typeObj,
	}
	return wrapper
end

bindingGetAttr = function(obj: PyObject, attr: string): PyObject?
	local payloadRaw = obj.__value
	if type(payloadRaw) ~= "table" then
		return nil
	end
	local payload = payloadRaw :: BindingPayload
	local binding = payload.binding
	if not binding then
		return nil
	end
	local rawValue: any?
	if binding.getAttr then
		local ok, result = pcall(binding.getAttr, payload.target, attr)
		if ok then
			rawValue = result
		else
			rawValue = nil
		end
	end
	if rawValue == nil then
		rawValue = safeIndex(payload.target, attr)
	end
	if rawValue == nil then
		return nil
	end
	if isPyObject(rawValue) then
		return rawValue :: PyObject
	end
	if type(rawValue) == "function" then
		local method = rawValue
		local callFn: (...any) -> any
		if binding.bindSelf then
			callFn = function(...: any): any
				return method(payload.target, ...)
			end
		else
			callFn = method
		end
		return Bridge.wrapLuauFunction(callFn)
	end
	return Bridge.luauToPy(rawValue)
end

bindingSetAttr = function(obj: PyObject, attr: string, value: PyObject)
	local payloadRaw = obj.__value
	if type(payloadRaw) ~= "table" then
		error("AttributeError: cannot set attribute on unbound object", 2)
	end
	local payload = payloadRaw :: BindingPayload
	local binding = payload.binding
	if not binding then
		error("AttributeError: cannot set attribute on unbound object", 2)
	end
	local luauValue = Bridge.pyToLuau(value)
	if binding.setAttr then
		local ok, errMsg = pcall(function()
			binding.setAttr(payload.target, attr, luauValue)
		end)
		if not ok then
			error(errMsg, 2)
		end
	else
		safeAssign(payload.target, attr, luauValue)
	end
end

local function ensureBindingType(binding: ClassBinding): PyType
	local ok, existing = pcall(function(): PyType
		return Base.getTypeObject(binding.pyTypeName)
	end)
	local typeObj: PyType
	if ok and existing then
		typeObj = existing
	else
		typeObj = Base.registerType(binding.pyTypeName, {})
	end
	typeObj.__getattr = bindingGetAttr
	typeObj.__setattr = bindingSetAttr
	typeObj.methods = typeObj.methods or {}
	return typeObj
end

local function createBindingInstance(value: any, visited: VisitedLuau, binding: ClassBinding): PyObject
	local existing = visited[value]
	if existing then
		return existing
	end
	ensureBindingType(binding)
	local payload: BindingPayload = {
		__pylua_binding = true,
		target = value,
		binding = binding,
	}
	local wrapper = Base.newPyObject(binding.pyTypeName, payload)
	visited[value] = wrapper
	return wrapper
end

local function tryClassBinding(value: any, visited: VisitedLuau): PyObject?
	if isPyObject(value) then
		return nil
	end
	for _, binding in ipairs(classBindings) do
		if bindingMatches(binding, value) then
			return createBindingInstance(value, visited, binding)
		end
	end
	return nil
end

function Bridge.registerClassBinding(options: ClassBindingOptionsInternal)
	if type(options) ~= "table" then
		error("registerClassBinding requires options table", 2)
	end
	local pyTypeName = options.pyTypeName
	if type(pyTypeName) ~= "string" or pyTypeName == "" then
		error("registerClassBinding requires pyTypeName", 2)
	end
	if options.typeofName == nil and options.predicate == nil then
		error("registerClassBinding requires typeofName or predicate", 2)
	end
	local bindSelf = if options.bindSelf ~= nil then options.bindSelf else true
	local existing = bindingByPyType[pyTypeName]
	local binding: ClassBinding
	if existing then
		existing.typeofName = options.typeofName
		existing.predicate = options.predicate
		existing.bindSelf = bindSelf
		existing.getAttr = options.getAttr
		existing.setAttr = options.setAttr
		existing.toLuau = options.toLuau
		binding = existing
	else
		binding = {
			pyTypeName = pyTypeName,
			typeofName = options.typeofName,
			predicate = options.predicate,
			bindSelf = bindSelf,
			getAttr = options.getAttr,
			setAttr = options.setAttr,
			toLuau = options.toLuau,
		}
		bindingByPyType[pyTypeName] = binding
		classBindings[#classBindings + 1] = binding
	end
	ensureBindingType(binding)
	return binding
end

function Bridge.clearClassBindings()
	for index = #classBindings, 1, -1 do
		classBindings[index] = nil
	end
	for key in pairs(bindingByPyType) do
		bindingByPyType[key] = nil
	end
end

local function luauToPyInternal(value: any, visited: VisitedLuau): PyObject
	local valueType = type(value)
	if valueType == "nil" then
		return Base.newNone()
	elseif valueType == "boolean" or valueType == "number" or valueType == "string" then
		return Base.ensurePyObject(value)
	elseif valueType == "function" then
		return Bridge.wrapLuauFunction(value)
	elseif valueType ~= "table" then
		local bindingWrapper = tryClassBinding(value, visited)
		if bindingWrapper then
			return bindingWrapper
		end
		return Base.newPyObject("object", value)
	end

	if isPyObject(value) then
		return value
	end

	if value.__nil then
		return Base.newNone()
	end

	local bindingWrapper = tryClassBinding(value, visited)
	if bindingWrapper then
		return bindingWrapper
	end

	if visited[value] then
		return visited[value]
	end

	-- Check for Luau class with __pylua_class marker in metatable
	local className = hasMetatableClass(value)
	if className then
		-- Create an auto-binding for this metatable-based class
		local autoBinding: ClassBinding = {
			pyTypeName = className,
			bindSelf = true,
		}
		ensureBindingType(autoBinding)
		local payload: BindingPayload = {
			__pylua_binding = true,
			target = value,
			binding = autoBinding,
		}
		local wrapper = Base.newPyObject(className, payload)
		visited[value] = wrapper
		return wrapper
	end

	if isSequenceTable(value) then
		local list = Collections.newList({})
		visited[value] = list
		local storage = list.__value :: { PyObject }
		local length = #value
		for index = 1, length do
			storage[index] = luauToPyInternal(value[index], visited)
		end
		return list
	end

	local dict = Collections.newDict({})
	visited[value] = dict
	local storage = dict.__value :: { [string]: HashEntry }
	for key, entryValue in pairs(value) do
		local pyKey = luauToPyInternal(key, visited)
		local pyValue = luauToPyInternal(entryValue, visited)
		storage[hashKey(pyKey)] = { key = pyKey, value = pyValue }
	end
	return dict
end

function Bridge.luauToPy(value: any): PyObject
	return luauToPyInternal(value, {})
end

local function pyToLuauInternal(obj: PyObject, visited: VisitedPy): any
	if visited[obj] ~= nil then
		return visited[obj]
	end

	local typeName = obj.__type
	if typeName == "NoneType" then
		return nil
	elseif typeName == "bool" then
		return obj.__value == true
	elseif typeName == "int" or typeName == "float" then
		return obj.__value
	elseif typeName == "str" or typeName == "bytes" then
		return obj.__value
	elseif typeName == "list" or typeName == "tuple" then
		local arr: { any } = {}
		visited[obj] = arr
		local source = obj.__value :: { PyObject }
		for index, item in ipairs(source) do
			arr[index] = pyToLuauInternal(item, visited)
		end
		return arr
	elseif typeName == "dict" then
		local tbl: { [any]: any } = {}
		visited[obj] = tbl
		local storage = obj.__value :: { [string]: HashEntry }
		for _, entry in pairs(storage) do
			local keyValue = pyToLuauInternal(entry.key, visited)
			local keyType = type(keyValue)
			if keyType ~= "string" and keyType ~= "number" and keyType ~= "boolean" then
				error("TypeError: unhashable type: '" .. entry.key.__type .. "'", 2)
			end
			tbl[keyValue] = pyToLuauInternal(entry.value, visited)
		end
		return tbl
	elseif typeName == "set" then
		local arr: { any } = {}
		visited[obj] = arr
		local storage = obj.__value :: { [string]: { value: PyObject } }
		local index = 1
		for _, entry in pairs(storage) do
			arr[index] = pyToLuauInternal(entry.value, visited)
			index += 1
		end
		return arr
	elseif type(obj.__value) == "table" and (obj.__value :: any).__pylua_binding then
		local payload = obj.__value :: BindingPayload
		local binding = payload.binding
		if binding.toLuau then
			visited[obj] = payload.target
			local converted = binding.toLuau(payload.target, visited)
			visited[obj] = converted
			return converted
		end
		visited[obj] = payload.target
		return payload.target
	elseif typeName == "luau_function" then
		local record: LuauFunctionRecord = obj.__value
		return record.func
	end

	return obj.__value
end

function Bridge.pyToLuau(value: any): any
	local obj: PyObject
	if isPyObject(value) then
		obj = value
	else
		obj = Base.ensurePyObject(value)
	end
	return pyToLuauInternal(obj, {})
end

return Bridge


local Builtins = {}

local base = require("./base")
type PyObject = base.PyObject

-- Forward references to type registration already performed in base bootstrap

-- None
export type PyNone = PyObject & { __type: "NoneType", __value: nil }
function Builtins.None(): PyNone
	return base.newNone() :: PyNone
end

-- Booleans
export type PyBool = PyObject & { __type: "bool", __value: boolean }

function Builtins.True(): PyBool
	return base.newBool(true) :: PyBool
end

function Builtins.False(): PyBool
	return base.newBool(false) :: PyBool
end

-- Numbers
export type PyInt = PyObject & { __type: "int", __value: number }
export type PyFloat = PyObject & { __type: "float", __value: number }

function Builtins.Int(n: number): PyInt
	if n % 1 ~= 0 then
		error("Int() requires integral value", 2)
	end
	return base.newInt(n) :: PyInt
end

function Builtins.Float(n: number): PyFloat
	return base.newFloat(n) :: PyFloat
end

export type PyStr = PyObject & { __type: "str", __value: string }

function Builtins.Str(s: string): PyStr
	return base.newPyObject("str", s) :: PyStr
end

-- Bytes
export type PyBytes = PyObject & { __type: "bytes", __value: string }

function Builtins.Bytes(val: any): PyBytes
	return base.newBytes(val) :: PyBytes
end

-- Minimal bytes operations (repr, eq, len, getitem)
local getTypeObject = base.getTypeObject
local bytesType = getTypeObject("bytes")
bytesType.methods = bytesType.methods or {}
local bytesMethods = bytesType.methods :: { [string]: (PyObject, { PyObject }?) -> PyObject }

-- __repr__
bytesMethods["__repr__"] = function(selfObj: PyObject)
	local s = selfObj.__value :: string
	local out = { "b'" }
	for i = 1, #s do
		local b = string.byte(s, i)
		if b >= 32 and b <= 126 and b ~= 39 and b ~= 92 then -- printable, not ' or \
			table.insert(out, string.char(b))
		elseif b == 39 then
			table.insert(out, "\\'")
		elseif b == 92 then
			table.insert(out, "\\\\")
		else
			table.insert(out, string.format("\\x%02x", b))
		end
	end
	table.insert(out, "'")
	return base.newPyObject("str", table.concat(out, ""))
end

-- __eq__
bytesMethods["__eq__"] = function(selfObj: PyObject, args: { PyObject }?)
	local other = args and args[1]
	if not other or other.__type ~= "bytes" then
		return base.newBool(false)
	end
	return base.newBool(selfObj.__value == other.__value)
end

-- __len__
bytesMethods["__len__"] = function(selfObj: PyObject)
	return base.newPyObject("int", #(selfObj.__value :: string))
end

-- __getitem__ (indexing)
bytesMethods["__getitem__"] = function(selfObj: PyObject, args: { PyObject }?)
	local idx = args and args[1]
	if not idx or (idx.__type ~= "int" and idx.__type ~= "float") then
		error("TypeError: bytes indices must be integers", 2)
	end
	local i = idx.__value
	if i < 0 then
		i = #selfObj.__value + i + 1
	else
		i = i + 1 -- Python is 0-based, Luau is 1-based
	end
	if i < 1 or i > #selfObj.__value then
		error("IndexError: bytes index out of range", 2)
	end
	local b = string.byte(selfObj.__value, i)
	return base.newPyObject("int", b)
end

-- __add__ (concatenate)
bytesMethods["__add__"] = function(selfObj: PyObject, args: { PyObject }?)
	local other = args and args[1]
	if not other or other.__type ~= "bytes" then
		error("TypeError: can't concat bytes to non-bytes", 2)
	end
	return base.newPyObject("bytes", (selfObj.__value :: string) .. (other.__value :: string))
end

-- __mul__ (repeat)
bytesMethods["__mul__"] = function(selfObj: PyObject, args: { PyObject }?)
	local nObj = args and args[1]
	if not nObj or nObj.__type ~= "int" then
		error("TypeError: can't multiply sequence by non-int", 2)
	end
	local n: number = nObj.__value
	if n <= 0 then
		return base.newPyObject("bytes", "")
	end
	local pieces = table.create(n)
	for i = 1, n do
		pieces[i] = selfObj.__value
	end
	return base.newPyObject("bytes", table.concat(pieces, ""))
end

-- Type objects (very small placeholder)
export type PyTypeObj = PyObject & { __type: "type", __value: { name: string } }

local function _makeTypeObject(name: string): PyTypeObj
	return base.newPyObject("type", { name = name }) :: PyTypeObj
end

-- Arithmetic method attachments (monkey patch into registered types)
local function attachPrimitiveMethods()
	local function wrapNum(result: number, left: PyObject, right: PyObject?): PyObject
		-- Preserve float if any operand is float
		if left.__type == "float" or (right and right.__type == "float") then
			return base.newFloat(result)
		end
		-- If non-integer result from int op -> float
		if result % 1 ~= 0 then
			return base.newFloat(result)
		end
		return base.newInt(result)
	end

	local registry = base.typeRegistrySnapshot()
	if registry.int then
		-- We can't mutate type tables directly from snapshot; fetch actual via internal API.
		-- Using protected access: base.getTypeObject (not exported earlier, but present).
		local getTypeObject = base.getTypeObject
		local intType = getTypeObject("int")
		local floatType = getTypeObject("float")
		local strType = getTypeObject("str")

		intType.methods = intType.methods or {}
		floatType.methods = floatType.methods or {}
		strType.methods = strType.methods or {}

		local intMethods = intType.methods :: { [string]: (PyObject, { PyObject }?) -> PyObject }
		local floatMethods = floatType.methods :: { [string]: (PyObject, { PyObject }?) -> PyObject }
		local strMethods = strType.methods :: { [string]: (PyObject, { PyObject }?) -> PyObject }

		strMethods["__getitem__"] = function(selfObj: PyObject, args: { PyObject }?)
			if not args or not args[1] then
				error("TypeError: string indices must be integers", 2)
			end
			local indexObj = base.ensurePyObject(args[1])
			if indexObj.__type ~= "int" then
				error("TypeError: string indices must be integers", 2)
			end
			local s = selfObj.__value :: string
			local length = #s
			local idx = indexObj.__value :: number
			if idx < 0 then
				idx += length
			end
			if idx < 0 or idx >= length then
				error("IndexError: string index out of range", 2)
			end
			local pos = idx + 1
			local char = string.sub(s, pos, pos)
			return base.newPyObject("str", char)
		end

		local numAdd = function(selfObj: PyObject, args: { PyObject }?)
			local other = args and args[1] or base.ensurePyObject(0)
			if other.__type ~= "int" and other.__type ~= "float" then
				error("TypeError: unsupported operand type(s) for +", 2)
			end
			return wrapNum(selfObj.__value + other.__value, selfObj, other)
		end

		local numSub = function(selfObj: PyObject, args: { PyObject }?)
			local other = args and args[1] or base.ensurePyObject(0)
			if other.__type ~= "int" and other.__type ~= "float" then
				error("TypeError: unsupported operand type(s) for -", 2)
			end
			return wrapNum(selfObj.__value - other.__value, selfObj, other)
		end

		local numMul = function(selfObj: PyObject, args: { PyObject }?)
			local other = args and args[1] or base.ensurePyObject(0)
			if other.__type ~= "int" and other.__type ~= "float" then
				error("TypeError: unsupported operand type(s) for *", 2)
			end
			return wrapNum(selfObj.__value * other.__value, selfObj, other)
		end

		local numTrueDiv = function(selfObj: PyObject, args: { PyObject }?)
			local other = args and args[1] or base.ensurePyObject(1)
			if other.__value == 0 then
				error("ZeroDivisionError: division by zero", 2)
			end
			return base.newFloat(selfObj.__value / other.__value)
		end

		local numFloorDiv = function(selfObj: PyObject, args: { PyObject }?)
			local other = args and args[1] or base.ensurePyObject(1)
			if other.__type ~= "int" and other.__type ~= "float" then
				error("TypeError: unsupported operand type(s) for //", 2)
			end
			if other.__value == 0 then
				error("ZeroDivisionError: integer division or modulo by zero", 2)
			end
			local quotient = selfObj.__value / other.__value
			local floored = math.floor(quotient)
			if selfObj.__type == "float" or other.__type == "float" then
				return base.newFloat(floored)
			end
			return base.newInt(floored)
		end

		intMethods["__add__"] = numAdd
		intMethods["__sub__"] = numSub
		intMethods["__mul__"] = numMul
		intMethods["__truediv__"] = numTrueDiv
		intMethods["__floordiv__"] = numFloorDiv
		floatMethods["__add__"] = numAdd
		floatMethods["__sub__"] = numSub
		floatMethods["__mul__"] = numMul
		floatMethods["__truediv__"] = numTrueDiv
		floatMethods["__floordiv__"] = numFloorDiv

		-- String concatenation & repetition minimal subset
		strMethods["__add__"] = function(selfObj: PyObject, args: { PyObject }?)
			local other = args and args[1] or base.ensurePyObject("")
			if other.__type ~= "str" then
				error("TypeError: can only concatenate str (not '" .. other.__type .. "') to str", 2)
			end
			return base.newPyObject("str", (selfObj.__value :: any) .. (other.__value :: any))
		end
		strMethods["__mul__"] = function(selfObj: PyObject, args: { PyObject }?)
			local other = args and args[1] or base.ensurePyObject(0)
			if other.__type ~= "int" then
				error("TypeError: can't multiply sequence by non-int", 2)
			end
			local n: number = other.__value
			if n < 0 then
				n = 0
			end
			local pieces = table.create(n)
			for i = 1, n do
				pieces[i] = selfObj.__value
			end
			return base.newPyObject("str", table.concat(pieces, ""))
		end
	end
end

attachPrimitiveMethods()

return Builtins


local Collections = {}

local base = require("./base")
local builtins = require("./builtins")

type PyObject = base.PyObject
type PyType = base.PyType
export type PyList = PyObject & { __type: "list", __value: { PyObject } }
export type PyTuple = PyObject & { __type: "tuple", __value: { PyObject } }
export type PyDict = PyObject & { __type: "dict", __value: { [string]: { key: PyObject, value: PyObject } } }
export type PySet = PyObject & { __type: "set", __value: { [string]: { value: PyObject } } }
type MethodMap = { [string]: (PyObject, { PyObject }?) -> PyObject }

-- Helper utilities

-- Cache for hash keys to avoid repeated computation
-- Uses weak keys (__mode = "k") so cached entries are automatically cleaned up by GC
-- when the PyObject keys are no longer referenced elsewhere
local hashCache = setmetatable({}, { __mode = "k" }) :: { [any]: string }

local function hashKey(obj: PyObject): string
	-- Quick cache lookup for frequently accessed objects
	local cached = hashCache[obj]
	if cached then
		return cached
	end

	local t = obj.__type
	local hash: string
	if t == "int" or t == "float" or t == "bool" or t == "str" then
		-- Use string.format for slightly better performance than concatenation
		hash = string.format("%s:%s", t, tostring(obj.__value))
	elseif t == "NoneType" then
		hash = "None"
	else
		hash = tostring(obj)
	end

	-- Cache the result (weak keys allow GC to clean up unused entries)
	hashCache[obj] = hash

	return hash
end

local function ensureIntIndex(value: PyObject, caller: string): number
	if value.__type ~= "int" then
		error("TypeError: " .. caller .. " indices must be integers", 2)
	end
	return value.__value :: number
end

local function resolveSequenceIndex(length: number, index: number, caller: string): number
	local idx = index
	if idx < 0 then
		idx += length
	end
	if idx < 0 or idx >= length then
		error("IndexError: " .. caller .. " index out of range", 2)
	end
	return idx + 1 -- convert to 1-based for Luau arrays
end

local function pyEquals(lhs: PyObject, rhs: PyObject): boolean
	if lhs == rhs then
		return true
	end
	local ok, eqMethod = pcall(function()
		return base.getattr(lhs, "__eq__")
	end)
	if ok and eqMethod then
		local result = base.call(eqMethod, { rhs })
		if type(result) == "table" and result.__type ~= nil then
			return base.truthy(result)
		elseif type(result) == "boolean" then
			return result
		end
		return base.truthy(base.ensurePyObject(result))
	end
	if lhs.__type == rhs.__type then
		local lv = lhs.__value
		local rv = rhs.__value
		if lv ~= nil and rv ~= nil then
			return lv == rv
		end
	end
	return false
end

local function sequenceContains(seq: { PyObject }, needle: PyObject): boolean
	for _, item in ipairs(seq) do
		if pyEquals(item, needle) then
			return true
		end
	end
	return false
end

local function boolResult(value: boolean): PyObject
	return base.newBool(value)
end

local function extractNumber(cell: PyObject, context: string): (number, boolean)
	if cell.__type == "int" then
		return cell.__value, false
	elseif cell.__type == "float" then
		return cell.__value, true
	end
	error("TypeError: " .. context .. " must contain numeric values", 2)
end

local function validateMatrix(listObj: PyObject, label: string): ({ PyObject }, number, number)
	if listObj.__type ~= "list" then
		error("TypeError: matrix multiplication requires list operands", 2)
	end
	local rows = listObj.__value :: { PyObject }
	if #rows == 0 then
		return rows, 0, 0
	end
	local firstRow = rows[1]
	if not firstRow or firstRow.__type ~= "list" then
		error("TypeError: " .. label .. " operand for @ must be a list of lists", 2)
	end
	local expectedCols = #(firstRow.__value :: { PyObject })
	for _, rowObj in ipairs(rows) do
		if rowObj.__type ~= "list" then
			error("TypeError: " .. label .. " operand for @ must be a list of lists", 2)
		end
		local rowValues = rowObj.__value :: { PyObject }
		if #rowValues ~= expectedCols then
			error("ValueError: inconsistent row lengths in " .. label .. " matrix", 2)
		end
	end
	return rows, #rows, expectedCols
end

local function matrixMultiply(left: PyObject, right: PyObject): PyObject
	local leftRows, leftRowCount, leftColCount = validateMatrix(left, "left")
	local rightRows, rightRowCount, rightColCount = validateMatrix(right, "right")
	if leftColCount ~= rightRowCount then
		error("ValueError: matrix dimension mismatch for @", 2)
	end
	local resultRows: { PyObject } = {}
	if leftRowCount == 0 or rightColCount == 0 then
		return Collections.newList(resultRows)
	end
	for i = 1, leftRowCount do
		local leftRowObj = leftRows[i]
		local leftValues = leftRowObj.__value :: { PyObject }
		local rowValues: { PyObject } = {}
		for j = 1, rightColCount do
			local sum = 0
			local sumHasFloat = false
			for k = 1, leftColCount do
				local rightRowObj = rightRows[k]
				local rightValues = rightRowObj.__value :: { PyObject }
				local leftCell = leftValues[k]
				local rightCell = rightValues[j]
				if not leftCell or not rightCell then
					error("ValueError: matrix dimension mismatch for @", 2)
				end
				local leftNumber, leftIsFloat = extractNumber(leftCell, "matrix")
				local rightNumber, rightIsFloat = extractNumber(rightCell, "matrix")
				local product = leftNumber * rightNumber
				sum += product
				sumHasFloat = sumHasFloat or leftIsFloat or rightIsFloat or (product % 1 ~= 0)
			end
			local cell: PyObject
			if sumHasFloat or (sum % 1 ~= 0) then
				cell = base.newFloat(sum)
			else
				cell = base.newInt(sum)
			end
			rowValues[#rowValues + 1] = cell
		end
		resultRows[#resultRows + 1] = Collections.newList(rowValues)
	end
	return Collections.newList(resultRows)
end

local getTypeObject = base.getTypeObject
local function ensureType(name: string)
	local ok = pcall(function()
		getTypeObject(name)
	end)
	if not ok then
		base.registerType(name, {})
	end
end

local function ensureMethods(pyType: PyType): MethodMap
	if not pyType.methods then
		pyType.methods = {} :: MethodMap
	end
	return pyType.methods :: MethodMap
end

ensureType("list")
ensureType("tuple")
ensureType("dict")
ensureType("set")
ensureType("list_iterator")

-- List implementation

function Collections.newList(values: { PyObject }?): PyList
	local arr: { PyObject } = {}
	if values then
		for i, v in ipairs(values) do
			arr[i] = base.ensurePyObject(v)
		end
	end
	return base.newPyObject("list", arr) :: PyList
end

do
	local listType = getTypeObject("list")
	local listMethods = ensureMethods(listType)

	listMethods["append"] = function(selfObj: PyObject, args: { PyObject }?)
		local value = args and args[1]
		if not value then
			error("TypeError: append() missing value", 2)
		end
		table.insert(selfObj.__value, base.ensurePyObject(value))
		return builtins.None()
	end

	listMethods["pop"] = function(selfObj: PyObject, args: { PyObject }?)
		local arr = selfObj.__value :: { PyObject }
		if #arr == 0 then
			error("IndexError: pop from empty list", 2)
		end
		local index: number
		if args and args[1] then
			local idxNum = ensureIntIndex(base.ensurePyObject(args[1]), "list")
			index = resolveSequenceIndex(#arr, idxNum, "list")
		else
			index = #arr
		end
		local removed = table.remove(arr, index)
		if not removed then
			error("IndexError: pop index invalid", 2)
		end
		return removed
	end

	listMethods["__len__"] = function(selfObj: PyObject)
		return base.newInt(#(selfObj.__value :: { PyObject }))
	end

	listMethods["__iter__"] = function(selfObj: PyObject)
		local iterState = { index = 0, list = selfObj }
		local iterator = base.newPyObject("list_iterator", iterState)
		local iterType = getTypeObject("list_iterator")
		local iterMethods = ensureMethods(iterType)
		if not iterMethods["__next__"] then
			iterMethods["__next__"] = function(it: PyObject)
				local state = it.__value
				state.index += 1
				local arr = state.list.__value :: { PyObject }
				if state.index > #arr then
					error("StopIteration", 2)
				end
				return arr[state.index]
			end
		end
		return iterator
	end

	listMethods["__add__"] = function(selfObj: PyObject, args: { PyObject }?)
		local other = args and args[1]
		if not other or other.__type ~= "list" then
			error("TypeError: can only concatenate list (not '" .. (other and other.__type or "nil") .. "') to list", 2)
		end
		local selfArr = selfObj.__value :: { PyObject }
		local otherArr = other.__value :: { PyObject }
		-- Pre-allocate combined array with known size
		local combined = table.create(#selfArr + #otherArr)
		local idx = 1
		for _, item in ipairs(selfArr) do
			combined[idx] = item
			idx += 1
		end
		for _, item in ipairs(otherArr) do
			combined[idx] = item
			idx += 1
		end
		return Collections.newList(combined)
	end

	listMethods["__getitem__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			error("TypeError: list indices must be integers", 2)
		end
		local idxNum = ensureIntIndex(base.ensurePyObject(args[1]), "list")
		local arr = selfObj.__value :: { PyObject }
		local resolved = resolveSequenceIndex(#arr, idxNum, "list")
		return arr[resolved]
	end

	listMethods["__setitem__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or #args < 2 then
			error("TypeError: list assignment requires index and value", 2)
		end
		local idxNum = ensureIntIndex(base.ensurePyObject(args[1]), "list")
		local value = base.ensurePyObject(args[2])
		local arr = selfObj.__value :: { PyObject }
		local resolved = resolveSequenceIndex(#arr, idxNum, "list")
		arr[resolved] = value
		return builtins.None()
	end

	listMethods["__contains__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			return boolResult(false)
		end
		local needle = base.ensurePyObject(args[1])
		local found = sequenceContains(selfObj.__value :: { PyObject }, needle)
		return boolResult(found)
	end

	listMethods["__matmul__"] = function(selfObj: PyObject, args: { PyObject }?)
		local other = args and args[1]
		if not other then
			error("TypeError: matrix multiplication requires right operand", 2)
		end
		return matrixMultiply(selfObj, base.ensurePyObject(other))
	end

	listMethods["__rmatmul__"] = function(selfObj: PyObject, args: { PyObject }?)
		local other = args and args[1]
		if not other then
			error("TypeError: matrix multiplication requires left operand", 2)
		end
		return matrixMultiply(base.ensurePyObject(other), selfObj)
	end
end

-- Tuple implementation (immutable)

function Collections.newTuple(values: { PyObject }?): PyTuple
	local arr: { PyObject } = {}
	if values then
		for i, v in ipairs(values) do
			arr[i] = base.ensurePyObject(v)
		end
	end
	return base.newPyObject("tuple", arr) :: PyTuple
end

do
	local tupleType = getTypeObject("tuple")
	local tupleMethods = ensureMethods(tupleType)

	tupleMethods["__len__"] = function(selfObj: PyObject)
		return base.newInt(#(selfObj.__value :: { PyObject }))
	end

	tupleMethods["__iter__"] = function(selfObj: PyObject)
		local listVersion: PyObject = Collections.newList(selfObj.__value)
		local iterMethod = base.getattr(listVersion, "__iter__")
		return base.call(iterMethod, {})
	end

	tupleMethods["__getitem__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			error("TypeError: tuple indices must be integers", 2)
		end
		local idxNum = ensureIntIndex(base.ensurePyObject(args[1]), "tuple")
		local arr = selfObj.__value :: { PyObject }
		local resolved = resolveSequenceIndex(#arr, idxNum, "tuple")
		return arr[resolved]
	end

	tupleMethods["__contains__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			return boolResult(false)
		end
		local needle = base.ensurePyObject(args[1])
		local found = sequenceContains(selfObj.__value :: { PyObject }, needle)
		return boolResult(found)
	end
end

-- Dict implementation

function Collections.newDict(entries: { { PyObject } }?): PyDict
	local storage: { [string]: { key: PyObject, value: PyObject } } = {}
	if entries then
		for _, pair in ipairs(entries) do
			local keyObj = base.ensurePyObject(pair[1])
			local valueObj = base.ensurePyObject(pair[2])
			storage[hashKey(keyObj)] = { key = keyObj, value = valueObj }
		end
	end
	return base.newPyObject("dict", storage) :: PyDict
end

do
	local dictType = getTypeObject("dict")
	local dictMethods = ensureMethods(dictType)

	dictMethods["set"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or #args < 2 then
			error("TypeError: set requires key and value", 2)
		end
		local key = base.ensurePyObject((args :: any)[1])
		local value = base.ensurePyObject((args :: any)[2])
		local store = selfObj.__value :: { [string]: { key: PyObject, value: PyObject } }
		store[hashKey(key)] = { key = key, value = value }
		return builtins.None()
	end

	dictMethods["get"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or #args < 1 then
			error("TypeError: get requires key", 2)
		end
		local key = base.ensurePyObject((args :: any)[1])
		local entry = (selfObj.__value :: any)[hashKey(key)]
		if entry then
			return entry.value
		end
		if args and (args :: any)[2] then
			return (args :: any)[2]
		end
		return builtins.None()
	end

	dictMethods["__len__"] = function(selfObj: PyObject)
		local count = 0
		for _ in pairs(selfObj.__value) do
			count += 1
		end
		return base.newInt(count)
	end

	dictMethods["__getitem__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			error("TypeError: dict indices must be provided", 2)
		end
		local key = base.ensurePyObject(args[1])
		local entry = (selfObj.__value :: any)[hashKey(key)]
		if not entry then
			error("KeyError: " .. tostring(key.__value), 2)
		end
		return entry.value
	end

	dictMethods["__setitem__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or #args < 2 then
			error("TypeError: dict assignment requires key and value", 2)
		end
		local key = base.ensurePyObject(args[1])
		local value = base.ensurePyObject(args[2])
		local store = selfObj.__value :: { [string]: { key: PyObject, value: PyObject } }
		store[hashKey(key)] = { key = key, value = value }
		return builtins.None()
	end

	dictMethods["__contains__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			return boolResult(false)
		end
		local key = base.ensurePyObject(args[1])
		local entry = (selfObj.__value :: any)[hashKey(key)]
		return boolResult(entry ~= nil)
	end

	dictMethods["keys"] = function(selfObj: PyObject)
		local values: { PyObject } = {}
		for _, entry in pairs(selfObj.__value) do
			values[#values + 1] = entry.key
		end
		return Collections.newList(values)
	end

	dictMethods["values"] = function(selfObj: PyObject)
		local values: { PyObject } = {}
		for _, entry in pairs(selfObj.__value) do
			values[#values + 1] = entry.value
		end
		return Collections.newList(values)
	end

	dictMethods["items"] = function(selfObj: PyObject)
		local pairsList: { PyObject } = {}
		for _, entry in pairs(selfObj.__value) do
			pairsList[#pairsList + 1] = Collections.newTuple({ entry.key, entry.value })
		end
		return Collections.newList(pairsList)
	end
end

-- Set implementation

function Collections.newSet(values: { PyObject }?): PySet
	local storage: { [string]: { value: PyObject } } = {}
	if values then
		for _, v in ipairs(values) do
			local valueObj = base.ensurePyObject(v)
			storage[hashKey(valueObj)] = { value = valueObj }
		end
	end
	return base.newPyObject("set", storage) :: PySet
end

do
	local setType = getTypeObject("set")
	local setMethods = ensureMethods(setType)

	setMethods["add"] = function(selfObj: PyObject, args: { PyObject }?)
		local value = args and args[1]
		if not value then
			error("TypeError: add requires value", 2)
		end
		local ensured = base.ensurePyObject(value)
		local store = selfObj.__value :: { [string]: { value: PyObject } }
		store[hashKey(ensured)] = { value = ensured }
		return builtins.None()
	end

	setMethods["__len__"] = function(selfObj: PyObject)
		local count = 0
		for _ in pairs(selfObj.__value) do
			count += 1
		end
		return base.newInt(count)
	end

	setMethods["__contains__"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			return boolResult(false)
		end
		local value = base.ensurePyObject(args[1])
		return boolResult((selfObj.__value :: any)[hashKey(value)] ~= nil)
	end

	setMethods["remove"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			error("TypeError: remove() requires a value", 2)
		end
		local value = base.ensurePyObject(args[1])
		local key = hashKey(value)
		local store: any = selfObj.__value
		if not store[key] then
			error("KeyError: value not found in set", 2)
		end
		store[key] = nil
		return builtins.None()
	end

	setMethods["discard"] = function(selfObj: PyObject, args: { PyObject }?)
		if not args or not args[1] then
			error("TypeError: discard() requires a value", 2)
		end
		local value = base.ensurePyObject(args[1])
		local store = selfObj.__value :: { [string]: { value: PyObject } }
		store[hashKey(value)] = nil
		return builtins.None()
	end
end

-- Public API shortcuts
Collections.list = Collections.newList
Collections.tuple = Collections.newTuple
Collections.dict = Collections.newDict
Collections.set = Collections.newSet

return Collections


local Base = require("./base")

local Module = {}

-- Register the module type
local moduleType = Base.registerType("module", {
	bases = { Base.getTypeObject("object") },
})

function Module.new(name: string, doc: string?): Base.PyObject
	local dict = {
		["__name__"] = Base.newPyObject("str", name),
		["__doc__"] = if doc then Base.newPyObject("str", doc) else Base.newNone(),
		["__package__"] = Base.newNone(),
		["__loader__"] = Base.newNone(),
		["__spec__"] = Base.newNone(),
	}

	local moduleObj = {
		__type = "module",
		__value = nil,
		__dict = dict,
		__typeobj = moduleType,
	}

	return moduleObj
end

return Module


local errors = {}

export type Location = { line: number, column: number }

local function fmtLocation(loc: Location?): string
	if not loc then
		return "<unknown>"
	end
	return string.format("line %d, column %d", loc.line, loc.column)
end

function errors.expected(expected: string, got: string?, loc: Location?)
	return string.format("ParseError: expected %s but got %s at %s", expected, got or "<eof>", fmtLocation(loc))
end

-- Simple raise helper (throws Lua error)
function errors.raise(msg: string, loc: Location?)
	if loc then
		error(string.format("%s at %s", msg, fmtLocation(loc)), 0)
	else
		error(msg, 0)
	end
end

-- Placeholder synchronization (future): skip until one of NEWLINE/DEDENT/ENDMARKER
function errors.synchronize(tokens: { any }, startPos: number): number
	-- Not implemented yet; returns original position
	return startPos
end

return errors