local CG = {}
CG.__index = CG

local AST = require"ast"
local Target = require"target"
local ETypes = require"etype"

local SIMPLE_BINOPS = {["+"] = "add", ["-"] = "sub", ["^"] = "xor", ["|"] = "or", ["&"] = "and"}
local SIMPLE_UNOPS = {["~"] = "not", ["-"] = "neg"}

local COMP_SIGNED = {["=="] = "e", ["!="] = "ne", ["<"] = "l", [">"] = "g", [">="] = "ge", ["<="] = "le"}
local COMP_UNSIGNED = {["=="] = "e", ["!="] = "ne", ["<"] = "b", [">"] = "a", [">="] = "be", ["<="] = "ae"}

function CG:xop(ex)
	if ex.kind == "expr-var" then
		return ex.vreg.cgi.register
	elseif ex.kind == "expr-int" then
		return tostring(ex.value)
	elseif ex.kind == "expr-unop" and ex.op == "*" then
		-- Memory ops
		if ex.a.kind == "expr-binop" and ex.a.op == "+" and ex.a.a.kind == "expr-var" and ex.a.b.kind == "expr-int" then
			-- *(var1 + int)
			return "[" .. ex.a.a.vreg.cgi.register .. " + " .. ex.a.b.value .. "]"
		elseif ex.a.kind == "expr-binop" and ex.a.op == "+" and ex.a.a.kind == "expr-var" and ex.a.b.kind == "expr-var" then
			-- *(var1 + var2)
			return "[" .. ex.a.a.vreg.cgi.register .. " + " .. ex.a.b.vreg.cgi.register .. "]"
		end
	end
	
	error("Unimplemented " .. tostring(ex))
end

function CG:emit(chunk)
	if chunk.cgi.stack_reservation > 0 then
		print("sub esp, " .. chunk.cgi.stack_reservation)
	end
	
	for _, stmt in ipairs(chunk.children) do
	
		if stmt.kind == "stmt-assign" then
			if stmt.src.kind == "expr-binop" then
				assert(SIMPLE_BINOPS[stmt.src.op])
				assert(self:xop(stmt.dest) == self:xop(stmt.src.a))
				
				print(SIMPLE_BINOPS[stmt.src.op] .. " " .. self:xop(stmt.dest) .. ", " .. self:xop(stmt.src.b))
			else
				print("mov " .. self:xop(stmt.dest) .. ", " .. self:xop(stmt.src))
			end
		elseif stmt.kind == "stmt-jump" then
			assert(stmt.condition.kind == "expr-binop")
			
			local cond = stmt.condition
			local op = cond.op
			
			assert(cond.a.etype == cond.b.etype)
			assert(cond.a.etype.kind == "scalar")
			assert(op == ">" or op == "<" or op == "==" or op == "!=" or op == ">=" or op == "<=")
			
			print("cmp " .. self:xop(cond.a) .. ", " .. self:xop(cond.b))
			
			local unsigned = cond.a.etype.unsigned
			
			print("j" .. (unsigned and COMP_UNSIGNED or COMP_SIGNED)[op] .. " .L" .. stmt.target)
		elseif stmt.kind == "stmt-label" then
			print(".L" .. stmt.name .. ":")
		elseif stmt.kind == "stmt-return" then
			if chunk.cgi.stack_reservation > 0 then
				print("add esp, " .. chunk.cgi.stack_reservation)
			end
			
			print("ret")
		else
			error("Unimplemented " .. tostring(stmt))
		end
		
	end
end

function CG:reg_alloc(chunk, vreguses)
	local vreg_live_start = {}
	
	for stmt_idx, stmt in ipairs(chunk.children) do
		if stmt.kind == "stmt-assign" and stmt.dest.kind == "expr-var" then
			local vreg = stmt.dest.vreg
			
			if vreg_live_start[vreg] then
				vreg_live_start[vreg] = math.min(vreg_live_start[vreg], stmt_idx)
			else
				vreg_live_start[vreg] = stmt_idx
			end
		end
	end
	
	local vreg_live_fin = {}
	
	for vreg, uses in pairs(vreguses) do
		vreg_live_fin[vreg] = require"set".max(uses)
	end
	
	-- Determine register classes as defined in target.lua
	for vreg in pairs(vreg_live_start) do
		vreg.cgi = {}
		if vreg.etype:byte_size() == 1 then
			vreg.cgi.register_class = "reg8"
		else
			vreg.cgi.register_class = "regn8"
		end
	end
	
	local edges = {}
	for vreg1, start1 in pairs(vreg_live_start) do
		edges[vreg1] = {}
		
		for vreg2, start2 in pairs(vreg_live_start) do
			if vreg1 ~= vreg2 then
				local fin1 = vreg_live_fin[vreg1]
				local fin2 = vreg_live_fin[vreg2]
				
				local live_range_intersection = ((start1 <= start2 and start2 <= fin1) or (start1 <= fin2 and fin2 <= fin1))
				local resource_intersection = (Target.REG_CLASSES[vreg1.cgi.register_class].mask & Target.REG_CLASSES[vreg2.cgi.register_class].mask) ~= 0
				
				if live_range_intersection and resource_intersection then
					edges[vreg1][vreg2] = true
				end
			end
		end
	end
	
	for vreg in pairs(vreg_live_start) do
		local found_reg
		for _, reg in ipairs(Target.REG_CLASSES[vreg.cgi.register_class].items) do
			local available = true
			for vreg2 in pairs(vreg_live_start) do
				if vreg2.cgi.register == reg then
					available = false
					break
				end
			end
			if available then
				found_reg = reg
				break
			end
		end
		
		if not found_reg then
			error("Spilling!")
		end
		
		vreg.cgi.register = found_reg
	end
end

function CG:compute_uses(chunk)
	local stmt_idx = 0
	
	local ret = {}
	
	chunk:generic_visitor(function(n)
		if n.kind:sub(1, 5) == "stmt-" then
			stmt_idx = stmt_idx + 1
		elseif n.kind == "expr-var" then
			if not ret[n.vreg] then
				ret[n.vreg] = {}
			end
			
			ret[n.vreg][stmt_idx] = true
		end
	end)
	
	return ret
end

function CG:compute_defs(chunk)
	local totaldefs = {}
	for stmt_idx, stmt in ipairs(chunk.children) do
		if stmt.kind == "stmt-assign" and stmt.dest.kind == "expr-var" then
			if not totaldefs[stmt.dest.vreg] then
				totaldefs[stmt.dest.vreg] = {}
			end
			table.insert(totaldefs[stmt.dest.vreg], stmt_idx)
		end
	end
	
	local outdefs = {}
	local indefs = {}
	
	local changed = {}
	for stmt_idx, stmt in ipairs(chunk.children) do
		outdefs[stmt_idx] = {}
		changed[stmt_idx] = true
	end
	
	while #changed > 0 do
		local stmt_idx
		for s in pairs(changed) do stmt_idx = s break end
		changed[stmt_idx] = nil
		local stmt = chunk.children[stmt_idx]
		
		local predecessors = {}
		if stmt_idx > 1 then
			predecessors[stmt_idx - 1] = true
		end
		if stmt.kind == "stmt-label" then
			for stmt2_idx, stmt2 in pairs(chunk.children) do
				if stmt2.kind == "stmt-jump" and stmt2.target == stmt.name then
					predecessors[stmt2_idx] = true
				end
			end
		end
		
		indefs[stmt_idx] = {}
		for pred in pairs(predecessors) do
			require"set".merge(indefs[stmt_idx], outdefs[pred])
		end
		
		local newout = {}
		local kill = nil
		if stmt.kind == "stmt-assign" and stmt.dest.kind == "expr-var" then
			kill = stmt.dest.vreg
			newout[stmt_idx] = true
		end
		for indef in pairs(indefs[stmt_idx]) do
			if chunk.children[indef].kind ~= "stmt-assign" or chunk.children[indef].dest.kind ~= "expr-var" or chunk.children[indef].dest.vreg ~= kill then
				newout[indef] = true
			end
		end
		
		if require"set".equal(outdefs[stmt_idx], newout) then
			break
		end
		
		local successors = {}
		if stmt_idx < #chunk.children then
			successors[stmt_idx + 1] = true
		end
		if stmt.kind == "jump" then
			for stmt2_idx, stmt2 in pairs(chunk.children) do
				if stmt2.kind == "label" and stmt2.name == stmt.target then
					successors[stmt2_idx] = true
				end
			end
		end
		
		for succ in pairs(successors) do
			changed[succ] = true
		end
		
		outdefs[stmt_idx] = newout
	end
	
	return indefs, outdefs
end

function CG:get_stack_vreg(chunk)
	for vreg in pairs(self:compute_uses(chunk)) do
		if vreg.cgi.register == "esp" then
			return vreg
		end
	end
	
	local vreg = AST.VReg("@stack", ETypes.scalar(true, Target.GPR_SIZE))
	vreg.cgi = {}
	vreg.cgi.register = "esp"
	return vreg
end

function CG:pass_save(chunk)
	local saves = chunk.etype.modifiers["save"]
	
	if not saves then
		return
	end
	
	local used_saves = {}
	for vreg in pairs(self:compute_uses(chunk)) do
		if saves[vreg.cgi.register] then
			table.insert(used_saves, vreg)
		end
	end
	
	if #used_saves == 0 then
		return
	end
	
	local stack_vreg = self:get_stack_vreg(chunk)
	
	for i = 1, #used_saves do
		table.insert(chunk.children, i,
			AST.assign(
				AST.unop(stack_vreg.etype.to, "*",
					AST.binop(stack_vreg.etype, AST.var(nil, stack_vreg), "+", AST.int(stack_vreg.etype, (i - 1) * Target.GPR_SIZE // 8))),
				AST.var(nil, used_saves[i])))
		
		chunk.cgi.stack_reservation = chunk.cgi.stack_reservation + 4
	end
	
	local handled_rets = {}
	for stmt_idx, stmt in ipairs(chunk.children) do
		if stmt.kind == "stmt-return" then
			if not handled_rets[stmt] then
				for i = 1, #used_saves do
					table.insert(chunk.children, stmt_idx,
						AST.assign(
							AST.var(nil, used_saves[i]),
							AST.unop(stack_vreg.etype.to, "*",
								AST.binop(stack_vreg.etype, AST.var(nil, stack_vreg), "+", AST.int(stack_vreg.etype, (i - 1) * Target.GPR_SIZE // 8)))))
				end
				handled_rets[stmt] = true
			end
		end
	end
end

function CG:process(chunk)
	assert(not chunk.cgi)
	chunk.cgi = {}
	chunk.cgi.stack_reservation = 0
	
	--local indefs, outdefs = self:compute_defs(chunk)
	local uses = self:compute_uses(chunk)
	self:reg_alloc(chunk, uses)
	
	self:pass_save(chunk)
end

local function apply(modules)
	local cg = setmetatable({}, CG)
	
	for module, root in pairs(modules) do
		-- Go through all declarations in module
		for _, decl in pairs(root.children) do
			if decl.export then
				print("global " .. decl.name)
			end
			print(decl.name .. ":")
			
			local ex = decl.expr
			
			if ex.etype.kind == "func" then
				cg:process(ex)
				cg:emit(ex)
			else
				error("Unimplemented " .. decl.expr.etype)
			end
		end
	end
end

return {apply = apply}