#!/usr/local/bin/lua51

-- 6l

-----------------------------------------------------------------------------
-- INITIALIZING DEFAULTS
-----------------------------------------------------------------------------
progname = "6l"
local oargs = require("oargs")
local log = require("log")
local stdio = require("stdio")
local bit = require("bit")
local tcmod = require("tcmod")

local args = oargs.compile(arg, { "-o" })

if #args.none == 0 then
	print("usage: 6l file ...")
	os.exit(1)
end

local output = (args["-o"] == nil) and "6.out" or args["-o"]

local info = { 
	symbols = { },
	offset = 0,
	org = 0,
}

local dst = io.open(output, "wb")
if dst == nil then
	log.error("failed to open " .. output)
end

local function search_ignore(src, ttype, ignore)
	local stream = { }
	while true do
		local tk = tcmod.readhdr(src)
		local imatch = false

		if tk.type ~= tcmod.tokens.nwl then
			table.insert(stream, tk)
		end

		for i = 1, #ignore do
			if tk.type == tcmod.tokens[ignore[i]] then
				imatch = true
				break
			end
		end

		if tk.type == tcmod.tokens[ttype] then
			return stream
		elseif imatch == false then
			return nil
		end
	end

	return nil
end

local function asmlink(hdr, src, status)
	if hdr.type == tcmod.tokens["}"] then
		status.inasm = false
	elseif hdr.type == tcmod.tokens.nwl then
		return
	elseif hdr.type == tcmod.tokens[";"] then
		tcmod.readasmline(src)
	elseif tcmod.isasmdu(hdr) == true then
		local stream = tcmod.readasmline(src)
		table.insert(stream, 1, hdr)

		local bytes = tcmod.fmtasmdu(info, stream)
		info.offset = info.offset + #bytes

		for i = 1, #bytes do stdio.writeu8(dst, bytes[i]) end

	-- Either a label or instruction or throw an error
	else
		local cur = src:seek("cur")
		local stream = tcmod.readasmline(src)
		table.insert(stream, 1, hdr) -- Insert header at start

		if #stream >= 2 and stream[2].type == tcmod.tokens[":"] then
			src:seek("set", cur)
			tcmod.readhdr(src) -- Skip colon

			if stream[1].type ~= tcmod.tokens.id then
				log.error("attempt to set label with other type than id at line " .. stream[1].line)
			end

			local name = stream[1].raw

			if info.symbols[name] ~= nil then
				log.error("label " .. name .. " already exists. Line " .. stream[1].line)
			end

			info.symbols[name] = info.offset

		else
			local bytes = tcmod.fmtasminst(info, stream)
			info.offset = info.offset + #bytes

			for i = 1, #bytes do stdio.writeu8(dst, bytes[i]) end
		end
	end
end

local function sclink(hdr, src, status)
	if hdr.type == tcmod.tokens["asm"] then
		if search_ignore(src, "{", { "nwl" }) == nil then
			log.error("illegal token after asm at line " .. hdr.line)
		end

		status.inasm = true

	elseif hdr.type == tcmod.tokens["}"] then
		if status.infunc == true then
			status.infunc = false

			-- MOV RSP, RBP
			-- POP RBP
			-- RET
			stdio.writeu8(dst, tcmod.hexval("0h48")) -- MOV RSP, RBP (3 bytes)
			stdio.writeu8(dst, tcmod.hexval("0h89"))
			stdio.writeu8(dst, tcmod.hexval("0hEC"))

			stdio.writeu8(dst, tcmod.hexval("0h5D")) -- POP RBP (1 byte)
			stdio.writeu8(dst, tcmod.hexval("0hC3")) -- RET (1 byte)

			-- Here we decrement to build space for local variables
			local cur = dst:seek("cur")
			dst:seek("set", status.funcaddr - info.org + 7)

			status.rsp_sub = bit.align16bytes(status.rsp_sub)
			local bytes = bit.nbytes(status.rsp_sub, 2)

			for i = 1, #bytes do stdio.writeu8(dst, bytes[i]) end

			dst:seek("set", cur)

			info.offset = info.offset + 5
			status.rsp_sub = 0
		else
			log.error("trailing close bracket at line " .. hdr.line)
		end
	elseif tcmod.isctype(hdr) then
		local cur = src:seek("cur")
		local stream = search_ignore(src, "(", { "id", "nwl" })

		if stream ~= nil then
			if status.infunc then log.error("nested functions are illegal line " .. hdr.line) end
			status.funcaddr = info.offset

			local fargs = search_ignore(src, "{", { "id", "void", "u8", "u16", "u32", "u64", "(", ")", "nwl" })
			
			if fargs == nil then
				log.error("illegal function format line " .. stream[1].line)
			end

			status.infunc = true

			local fname = stream[1].raw
			if info.symbols[fname] ~= nil then
				log.error("double definition of symbol line " .. hdr.line)
			end

			info.symbols[fname] = info.offset

			-- PUSH RBP
			-- MOV RBP, RSP
			-- SUB RSP, 0 ( for now )
			stdio.writeu8(dst, tcmod.hexval("0h55")) -- PUSH RBP (1 byte)

			stdio.writeu8(dst, tcmod.hexval("0h48")) -- MOV RBP, RSP (3 bytes)
			stdio.writeu8(dst, tcmod.hexval("0h89"))
			stdio.writeu8(dst, tcmod.hexval("0hE5"))

			stdio.writeu8(dst, tcmod.hexval("0h48")) -- SUB RSP, 0 (7 bytes)
			stdio.writeu8(dst, tcmod.hexval("0h81"))
			stdio.writeu8(dst, tcmod.hexval("0hEC"))
			stdio.writeu8(dst, tcmod.hexval("0h00"))
			stdio.writeu8(dst, tcmod.hexval("0h00"))
			stdio.writeu8(dst, tcmod.hexval("0h00"))
			stdio.writeu8(dst, tcmod.hexval("0h00"))

			info.offset = info.offset + 11


			return
		end

		--[[
		src:seek("set", cur)
		stream = search_ignore(src, "=", { "id", "nwl" })

		if stream ~= nil then
			log.error("bad instruction format line " .. hdr.line)
			return
		end
		]]--

		src:seek("set", cur)
		stream = search_ignore(src, ";", { "id", "nwl", "*" })

		if stream ~= nil then
			if status.infunc == true then
				if stream[1].type == tcmod.tokens["*"] then
					status.rsp_sub = status.rsp_sub + 8

					-- mov u64 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0h48")) -- REX.W prefix
					stdio.writeu8(dst, tcmod.hexval("0hC7"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writeu8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 9
				elseif hdr.type == tcmod.tokens["u8"] or hdr.type == tcmod.tokens["s8"] then
					-- Write 0 to address
					status.rsp_sub = status.rsp_sub + 1

					-- mov u8 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0hC6"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writes8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 4

				elseif hdr.type == tcmod.tokens["u16"] or hdr.type == tcmod.tokens["s16"] then
					status.rsp_sub = status.rsp_sub + 2

					-- mov u16 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0h66"))
					stdio.writeu8(dst, tcmod.hexval("0hC7"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writeu8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 6

				elseif hdr.type == tcmod.tokens["u32"] or hdr.type == tcmod.tokens["s32"] then
					status.rsp_sub = status.rsp_sub + 4

					-- mov u32 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0hC7"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writeu8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 8

				elseif hdr.type == tcmod.tokens["u64"] or hdr.type == tcmod.tokens["s64"] then
					status.rsp_sub = status.rsp_sub + 8

					-- mov u64 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0h48")) -- REX.W prefix
					stdio.writeu8(dst, tcmod.hexval("0hC7"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writeu8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 9

				elseif hdr.type == tcmod.tokens["f32"] then
					status.rsp_sub = status.rsp_sub + 8

					-- mov u64 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0h48")) -- REX.W prefix
					stdio.writeu8(dst, tcmod.hexval("0hC7"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writeu8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 9

				elseif hdr.type == tcmod.tokens["f64"] then
					status.rsp_sub = status.rsp_sub + 8

					-- mov u64 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0h48")) -- REX.W prefix
					stdio.writeu8(dst, tcmod.hexval("0hC7"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writeu8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 9

					status.rsp_sub = status.rsp_sub + 8

					-- mov u64 [rbp - status.rsp_sub], 0
					stdio.writeu8(dst, tcmod.hexval("0h48")) -- REX.W prefix
					stdio.writeu8(dst, tcmod.hexval("0hC7"))
					stdio.writeu8(dst, tcmod.hexval("0h45"))
					stdio.writeu8(dst, -status.rsp_sub + MAX_U8)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)
					stdio.writeu8(dst, 0)

					info.offset = info.offset + 9

				else
					log.error("variables of type u8 are only ones supported for now line " .. hdr.line)
				end
			else
				if hdr.type == tcmod.tokens["u8"] then
				else
					log.error("variables of type u8 are only ones supported for now line " .. hdr.line)
				end
			end
			return
		end
		log.error("functions are only suportted for now illegal at line " .. hdr.line)
	end
end

local function link(src)
	local status = { 
		inasm = false,
		infunc = false,
		rsp_sub = 0,
		funcaddr = nil,
	}

	while true do
		local hdr = tcmod.readhdr(src)

		if hdr == nil then break end

		if hdr.type == tcmod.tokens.org then
			info.offset = tonumber(tcmod.hexval(hdr.raw))
			info.org = tonumber(tcmod.hexval(hdr.raw))
		else
			if status.inasm == false then
				sclink(hdr, src, status)
			else
				asmlink(hdr, src, status)
			end
		end
	end

	if status.infunc == true then
		log.error("function not closed")
	end
end


for i = 1, #args.none do
	local src = io.open(args.none[i], "rb")

	if src == nil then
		log.error("failed to open " .. args.none[i])
	end

	link(src)
end

tcmod.resolve(info, dst)