local AST = require"ast"

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

local Logger = require"logger"

local NEXT_LABEL_ID = 1

local Scope = {}
Scope.__index = Scope
function Scope.new(parent)
	return setmetatable({parent = parent, items = {}}, Scope)
end
function Scope:find(name)
	if self.items[name] then
		return self.items[name]
	end
	if self.parent then
		return self.parent:find(name)
	end
	return nil
end
function Scope:add(name)
	self.items[name] = AST.VReg(name)
	return self.items[name]
end

local Parser = {}
Parser.__index = Parser

function Parser:parse_root()
	local root = AST.root()
	
	while true do
		if self:peek(0).type == "eof" then
			break
		end
		
		local declaration, err = self:parse_declaration()
		
		if not declaration then
			self:log("err", err or "can't parse")
			break
		end
		
		table.insert(root.children, declaration)
	end
	
	return root
end

function Parser:parse_declaration()
	local old_i = self.i
	
	local export = self:maybe"export"
	
	if not self:maybe"ident" then
		self.i = old_i
		return nil, "expected identifier"
	end
	
	local name = self:last().data
	
	if not self:maybe":" then
		return nil, "expected :"
	end
	
	local expr, err = self:parse_expr(0)
	
	if not expr then
		return nil, err
	end
	
	return AST.decl(name, expr, export)
end

function Parser:parse_expr(precedence)
	local old_idx = self.i
	
	if precedence == 0 then
		local ret, err = self:parse_expr(precedence + 1)
		
		if not ret then
			self.i = old_idx
			return nil, err
		end
		
		while self:maybe">" or self:maybe"<" or self:maybe"==" or self:maybe"!=" or self:maybe">=" or self:maybe"<=" do
			local op = self:last().type
			
			local a = ret
			local b, err = self:parse_expr(precedence + 1)
			
			if not b then
				self.i = old_idx
				return nil, err
			end
			
			assert(a.etype == b.etype)
			
			ret = AST.binop(a.etype, a, op, b)
		end
		
		return ret
	elseif precedence == 1 then
		local ret, err = self:parse_expr(precedence + 1)
		
		if not ret then
			self.i = old_idx
			return nil, err
		end
		
		while self:maybe"+" or self:maybe"-" do
			local op = self:last().type
			
			local a = ret
			local b, err = self:parse_expr(precedence + 1)
			
			if not b then
				self.i = old_idx
				return nil, err
			end
			
			assert(a.etype == b.etype)
			
			ret = AST.binop(a.etype, a, op, b)
		end
		
		return ret
	elseif precedence == 2 then
		local ret, err = self:parse_expr(precedence + 1)
		
		if not ret then
			self.i = old_idx
			return nil, err
		end
		
		while self:maybe"*" or self:maybe"/" do
			local op = self:last().type
			
			local a = ret
			local b, err = self:parse_expr(precedence + 1)
			
			if not b then
				self.i = old_idx
				return nil, err
			end
				
			assert(a.etype == b.etype)
			
			ret = AST.binop(a.etype, a, op, b)
		end
		
		return ret
	elseif precedence == 3 then
		local ret, err
		while self:maybe"&" or self:maybe"*" do
			local op = self:last().type
			
			local a, err = self:parse_expr(precedence)
			if not a then
				self.i = old_idx
				return nil, err
			end
			
			ret = AST.unop(op == "&" and ETypes.ref(a.etype) or ETypes.deref(a.etype), op, a)
		end
		
		if not ret then
			ret, err = self:parse_expr(precedence + 1)
		end
		
		if not ret then
			self.i = old_idx
			return nil, err
		end
		
		return ret
	elseif precedence == 4 then
		local a, err = self:parse_expr(precedence + 1)
		
		if not a then
			self.i = old_idx
			return nil, err
		end
		
		while self:maybe"[" do
			if a.etype.kind == "pointer" then
				a = AST.unop(a.etype.to, "*", a)
			end
			
			assert(a.etype.kind == "array", "Indexing a non-array")
			
			local b, err = self:parse_expr(0)
			
			if not b then
				self.i = old_idx
				return nil, err
			end
			
			a = AST.unop(a.etype.element_etype, "*", AST.binop(nil, a, "+", b))
			
			assert(self:maybe"]")
		end
		
		return a
	elseif precedence == 5 then
		local asdf = self.i
		
		-- Okay for etype to be nil
		local etype = self:parse_etype(0)
		
		if self:maybe"?" then
			return AST.unknown(etype)
		elseif self:maybe'string' then
			return AST.cast(AST.string(self:last().data), etype)
		elseif self:maybe"num" then
			local tok = self:last().data
			local base = tonumber(tok:match"^([0-9]+)r") or 10
			local val = tok:match"r([0-9a-zA-Z]+)$" or tok
			return AST.int(etype, tonumber(val, base))
		elseif self:maybe"ident" then
			assert(not etype)
			
			local name = self:last().data
			local vreg = self.scope:find(name)
			if not vreg then
				self.i = old_idx
				return nil, "Undeclared variable " .. name
			end
			
			return AST.var(vreg.etype, vreg)
		elseif self:maybe"[" then
			assert(etype)
			
			local children = {}
			
			if not self:maybe"]" then
				while true do
					local ex = self:parse_expr(0)
					assert(ex)
					
					if etype then
						ex = AST.cast(ex, etype.element_etype)
					end
					
					table.insert(children, ex)
					
					if self:maybe"]" then
						break
					end
					
					if not self:maybe"," then
						local last = self:last()
						self.i = old_idx
						return nil, "expected comma", last
					end
				end
			end
			
			local n = AST.array(etype, 0)
			for _, child in ipairs(children) do
				table.insert(n.children, child)
			end
			return n
		elseif self:maybe"{" then
			local n = AST.func(etype)
			while not self:maybe"}" do
				local status, err = self:parse_stmt(n)
				
				if not status then
					self.i = old_idx
					return nil, err
				end
			end
			return n
		end
	end
	
	self.i = old_idx
	return nil
end

function Parser:parse_stmt(chunk)
	local old_idx = self.i
	
	if self:maybe";" then
		return true
	elseif self:maybe"return" then
		local ex = self:parse_expr(0)
		
		-- ex can be null
		
		table.insert(chunk.children, AST.ret(ex))
		
		return true
	elseif self:maybe"if" then
		local condition = self:parse_expr(0)
		
		local lbl_id = NEXT_LABEL_ID
		NEXT_LABEL_ID = NEXT_LABEL_ID + 1
		
		table.insert(chunk.children, AST.jump(condition, lbl_id))
		
		if not self:maybe"{" then
			self.i = old_idx
			return false, "expected {"
		end
		
		while not self:maybe"}" do
			local status, err = self:parse_stmt(chunk)
			if not status then
				self.i = old_idx
				return nil, err
			end
		end
		
		table.insert(chunk.children, AST.label(lbl_id))
		
		return true
	elseif self:peek(0).type == "ident" and self:peek(1).type == "=" then
		local name = self:next().data
		self:next()
		
		local expr = self:parse_expr(0)
		
		local vreg = self.scope:find(name)
		if not vreg then
			assert(expr.etype)
			
			vreg = self.scope:add(name)
			vreg.etype = expr.etype
		else
			assert(vreg.etype == expr.etype, "Type mismatch in assignment")
		end
		
		table.insert(chunk.children, AST.assign(AST.var(vreg.etype, vreg), expr))
		
		return true
	end
	
	-- Try parsing assignment
	
	local ex = self:parse_expr(0)
	if not ex then
		self.i = old_idx
		return nil, "expected expression"
	end
	
	if not self:maybe"=" then
		self.i = old_idx
		return nil, "expected ="
	end
	
	local ex2 = self:parse_expr(0)
	if not ex2 then
		self.i = old_idx
		return nil, "expected expression"
	end
	
	table.insert(chunk.children, AST.assign(ex, AST.cast(ex2, ex.etype)))
	
	return true
end

function Parser:parse_etype(precedence)
	local old_idx = self.i
	
	if precedence == 0 then
		local ret = self:parse_etype(precedence + 1)
		
		if not ret then
			return nil
		end
		
		while self:maybe"->" do
			local a = ret
			local b = self:parse_etype(precedence + 1)
			
			ret = ETypes.func(a, b)
			
			while self:maybe("mod") do
				local mod_type = self:last().data
				
				if mod_type == "@save" then
					if not self:maybe"(" then
						local last = self:last()
						self.i = old_idx
						return nil, "expected (", last
					end
					
					local ex = self:parse_expr(0)
					if not ex or ex.kind ~= "expr-array" then
						local last = self:last()
						self.i = old_idx
						return nil, "expected array", last
					end
					
					if not self:maybe")" then
						local last = self:last()
						self.i = old_idx
						return nil, "expected )", last
					end
					
					local items = {}
					for _, child in ipairs(ex.children) do
						assert(child.kind == "expr-string")
						
						if not require"target".REGS[child.value] then
							self:log("warn", "skipping unknown register " .. child.value)
						end
						
						items[child.value] = true
					end
					
					ret.modifiers["save"] = items
				else
					self:log("warn", "skipping unknown modifier " .. mod_type)
					
					if self:maybe"(" then
						local depth = 1
						while true do
							local t = self:next()
							if t.type == "(" then
								depth = depth + 1
							elseif t.type == ")" then
								depth = depth - 1
								if depth == 0 then
									break
								end
							elseif t.type == "eof" then
								break
							end
						end
					end
				end
			end
		end
		
		return ret
	elseif precedence == 1 then
		local a = self:parse_etype(precedence + 1)
		
		if not a then
			return nil
		end
		
		while self:maybe"*" or self:maybe"[" do
			local old_idx = self.i
			
			if self:last().type == "*" then
				a = ETypes.ref(a)
			else
				local length_expr = self:parse_expr(0)
				
				if not length_expr or not self:maybe"]" or (length_expr.kind ~= "expr-unknown" and length_expr.kind ~= "expr-int") then
					-- This cannot be an array type, most likely a case like (string[?] [...])
					self.i = old_idx - 1
					break
				end
				
				local length
				if length_expr.kind == "expr-int" then
					length = length_expr.value
				end
				
				a = ETypes.array(a, length)
			end
		end
		return a
	elseif precedence == 2 then
		if self:maybe"(" then
			if self:maybe")" then
				return ETypes.struct({})
			else
				self:go_back()
			end
		end
		
		if self:maybe"stringkw" then
			return ETypes.string()
		end
		
		if not self:maybe"ident" then
			self.i = old_idx
			return nil
		end
		
		local str = self:last().data
		
		local ret
		if str:match"[ui][0-9]+" then
			ret = ETypes.scalar(str:sub(1, 1) == "u", tonumber(str:sub(2)))
		elseif str:match"[ui]gpr" then
			ret = ETypes.scalar(str:sub(1, 1) == "u", Target.GPR_SIZE)
		else
			self.i = old_idx
			return nil
		end
		
		return ret
	end
	
	self.i = old_idx
	return nil
end

function Parser:go_back()
	if self.i <= 1 then
		error("Already at the beginning")
	end
	self.i = self.i - 1
end

function Parser:last()
	return self.tokens[self.i - 1]
end

function Parser:peek(idx)
	if self.i + idx > #self.tokens then
		return {type = "eof"}
	end
	return self.tokens[self.i + idx]
end

function Parser:maybe(token_type)
	if self.i > #self.tokens then
		return false
	end
	
	if self.tokens[self.i].type == token_type then
		self.i = self.i + 1
		return true
	end
	
	return false
end

function Parser:next()
	local ret = self:peek(0)
	if ret.type ~= "eof" then
		self.i = self.i + 1
	end
	return ret
end

function Parser:log(msg_type, fmt, ...)
	local tok = self:last()
	Logger.log(msg_type, tok.y .. ":" .. tok.x .. ", " .. fmt, ...)
end

return function(tokens)
	return setmetatable({i = 1, tokens = tokens, scope = Scope.new(nil)}, Parser)
end