local arch8608 = require("rom-8608-defs")

local aliases = {
	["jpz imm8"] = {"jeq imm8"},
	["jnz imm8"] = {"jne imm8"},
	["jmp q"   ] = {"ret"     },
}

local function trim(s) return s:gsub("^ +", ""):gsub(" +$", "").."" end
local function validWordsFromInstrs(instrs)
	local words = {}
	for mnem, _ in pairs(instrs) do
		for word in mnem:gmatch("[^ ]+") do
			words[word] = true
		end
	end
	return words
end
local function decodeNumber(n)
	n = trim(n)
	local sign = 1; if n:sub(1, 1)=="-" then sign = -1; n = n:sub(2, #n); end;
	if     n:sub(1, 1)=="$"   then return sign*(tonumber(n:sub(2, #n  ), 16) or error("invalid number "..n)), math.ceil((#n-1)/2)
	elseif n:sub(1, 2)=="0x"  then return sign*(tonumber(n:sub(3, #n  ), 16) or error("invalid number "..n)), math.ceil((#n-2)/2)
	elseif n:sub(#n, #n)=="h" then return sign*(tonumber(n:sub(1, #n-1), 16) or error("invalid number "..n)), math.ceil((#n-1)/2)
	else
		local v = sign*(tonumber(n) or error("invalid number "..n))
		if v>=-128 and v<=255 then return v, 1
		elseif v>=-32768 and v<=65535 then return v, 2
		else error("out-of-range number "..v) end
	end
end
local function mnemFromLine(line, instrs, validWords)
	local firstWord = line:match("^[^ ]+")
	local imms = {}
	local function addNum(n)
		n = trim(n)
		local val, len = decodeNumber(n)
		local linei8  = line:gsub(n, "imm8", 1, true)
		local linei16 = line:gsub(n, "imm16", 1, true)
		if len==1 and (not instrs[linei8]) and instrs[linei16] then len = 2 end
		table.insert(imms, { val = val, len = len } )
		return " imm"..(len*8).." "
	end
	local function addLabel(n)
		n = trim(n)
		local len = 2
		local linei8 = line:gsub(n, "imm8", 1, true)
		if instrs[linei8] then len = 1 end
		table.insert(imms, { label = n, len = len } )
		return " imm"..(len*8).." "
	end
	
	local mnem = " "..line:gsub(" ", "  ").." "
	mnem = mnem:gsub("%- *", " %+%-")
	mnem = mnem:gsub("([%*%+])", " %1 ")
	mnem = mnem:gsub(" %-?%$[0-9a-fA-F]+ "       , function(n)                                 return addNum  (n)     end)
	mnem = mnem:gsub(" %-?0x[0-9a-fA-F]+ "       , function(n)                                 return addNum  (n)     end)
	mnem = mnem:gsub(" %-?[0-9a-fA-F]+h "        , function(n) if not validWords[trim(n)] then return addNum  (n) end end)
	mnem = mnem:gsub(" %-?[0-9]+ "               , function(n) if not validWords[trim(n)] then return addNum  (n) end end)
	mnem = mnem:gsub(" [a-zA-Z_][a-zA-Z0-9_%.]* ", function(n) if not validWords[trim(n)] then return addLabel(n) end end)
	mnem = trim(mnem):gsub(" +", " ")
	
	if not instrs[mnem] then mnem = mnem:gsub("%+ imm", "imm") end
	
	return mnem, imms
end
local function addByte(state, val)
	assert(val>=-128 and val<=255, "invalid byte "..val)
	assert(state.memory[state.curAddr]==nil, "overwriting memory at "..state.curAddr)
	state.memory[state.curAddr] = val%256
	state.curAddr = state.curAddr + 1
end
local function addWord(state, val)
	assert(val>=0 and val<=65535, "invalid word "..val)
	addByte(state, math.floor(val/256))
	addByte(state, val%256)
end
local function addSpace(state, len)
	for i = 1, len do
		assert(state.memory[state.curAddr]==nil, "overwriting memory at "..state.curAddr)
		state.memory[state.curAddr] = false
		state.curAddr = state.curAddr + 1
	end
end
local function assembleInstruction(line, state, instrs, validWords)
	local mnem, imms = mnemFromLine(line, instrs, validWords)
	local opcode = instrs[mnem] or error("invalid instruction \""..line.."\" (mnem \""..mnem.."\")")
	local writeimms = true
	local padlen = 0
	if type(opcode)=="function" then
		padlen, writeimms = opcode(imms)
		addSpace(state, padlen)
	elseif opcode>=0 then
		addByte(state, opcode)
	end
	if writeimms then
		for _, imm in ipairs(imms) do
			if imm.val then
				if     imm.len==1 then addByte(state, imm.val)
				elseif imm.len==2 then addWord(state, imm.val)
				else error("invalid imm len") end
			elseif imm.label then
				table.insert(state.labelReplacements, {
					name = imm.label,
					addr = state.curAddr,
					len = imm.len,
					rel = imm.len==1,
				})
				state.curAddr = state.curAddr + imm.len
			else error("invalid imm") end
		end
	end
end
local directiveFunctions = {
	fn    = function(state, fn) state.fileName = fn end,
	ln    = function(state, ln) state.lineNum = tonumber(ln) end,
	org   = function(state, addr) state.curAddr = decodeNumber(addr) end,
	align = function(state, alns) local aln = decodeNumber(alns); if state.curAddr % aln ~= 0 then state.curAddr = state.curAddr + (aln - state.curAddr%aln) end end,
}
local function assembleCode(code, instrs)
	local validWords = validWordsFromInstrs(instrs)
	
	local state = {
		lineNum = 0,
		fileName = "",
		curAddr = 0,
		memory = {},
		labelReplacements = {},
		labelAddrs = {},
	}
	
	for line in code:gmatch("[^\n]+") do
		line = trim(line)
		if line:sub(1, 1)=="." then -- directive
			local dir, rest = line:match("^%.([^ ]+) *(.*)$")
			assert(dir and rest, "no directive on line "..line)
			local dirf = directiveFunctions[dir] or error("invalid directive "..dir)
			dirf(state, rest)
		elseif line:sub(#line, #line)==":" then -- label
			local name = line:sub(1, #line-1)
			assert(not state.labelAddrs[name], "redefinition of label "..name)
			state.labelAddrs[name] = state.curAddr
		elseif line:find("[^ ]") then
			assembleInstruction(line, state, instrs, validWords)
		end
	end
	
	for _, rep in ipairs(state.labelReplacements) do
		local labelAddr = state.labelAddrs[rep.name] or error("no label named "..rep.name)
		state.curAddr = rep.addr
		if     rep.len==1 then addByte(state, labelAddr-(rep.addr+1))
		elseif rep.len==2 then addWord(state, labelAddr)
		else error("invalid labelreplace len") end
	end
	
	return state.memory
end

local function readFile(fn)
	local fi = io.open(fn, "r") or error("could not open file "..fn)
	local text = fi:read("*a")
	fi:close()
	return text
end

local function separateCommas(l)
	local c = {}; for a in l:gmatch("[^,]+") do table.insert(c, trim(a)) end; return c;
end
local function evaluateExpression(expr)
	expr = expr:gsub("[^%+%-%*%/]+", function(word)
		local val = decodeNumber(word) or error("invalid number in expression: "..word)
		return val
	end)
	assert(not expr:find("[^a-zA-Z0-9_%(%)%+%-%*%/ \t\r\n]"), "invalid char in expression: "..expr)
	local exprf = loadstring("return "..expr)
	local eval = exprf() or error("invalid expr: "..expr)
	return eval
end
local function preprocessCode(code)
	code = "\n"..code.."\n"
	
	local funcmacros = {}
	code = code:gsub(".define ([%.a-zA-Z0-9_]+)%(([^%)]+)%) ([^\n]+)", function(name, args, repl)
		local argt = separateCommas(args)
		for argidx, arg in ipairs(argt) do assert(not arg:find("[^a-zA-Z0-9_]"), "invalid character in macro arg name: "..name.." "..arg) end
		repl = " "..repl.." "
		local invoc = 0
		funcmacros[name] = function(b, callargs)
			invoc = invoc + 1
			local callargt = separateCommas(callargs)
			local callrepl = repl
			for argidx, arg in ipairs(argt) do
				local callarg = callargt[argidx]
				callrepl = callrepl:gsub("([^a-zA-Z0-9_])"..arg.."([^a-zA-Z0-9_])", "%1"..callarg.."%2")
			end
			callrepl = callrepl:gsub("(_BRACE_[0-9]+_)", "%1"..invoc.."_")
			return b..callrepl
		end
		return ""
	end)
	for name, replf in pairs(funcmacros) do code = code:gsub("([^a-zA-Z0-9_])"..name.." *%(([^%)]+)%)", replf) end
	
	local simplemacros = {}
	code = code:gsub("%.define +([%.a-zA-Z0-9_]+) +([^\n]+)", function(name, repl)
		assert(not simplemacros[name], "Redefinition of macro "..name)
		simplemacros[name] = repl
		return ""
	end)
	--for name, repl in pairs(simplemacros) do code = code:gsub(name, repl, 1, true) end
	for name, repl in pairs(simplemacros) do
		local invoc = 0
		code = code:gsub("([^a-zA-Z0-9_])"..name.."([^a-zA-Z0-9_])", function(b, a)
			invoc = invoc+1
			return b..(repl:gsub("(_BRACE_[0-9]+_)", "%1"..invoc.."_"))..a
		end)
	end
	
	code = code:gsub("\\", "\n")
	
	local codet = {}
	local exprt = {}
	local parenLevel = 0
	for i = 1, #code do
		local c = code:sub(i, i)
		if c=="(" then
			parenLevel = parenLevel+1
		elseif c==")" then
			parenLevel = parenLevel-1
			if parenLevel==0 then
				table.insert(codet, evaluateExpression(table.concat(exprt)))
				exprt = {}
			end
		else
			if parenLevel==0 then table.insert(codet, c)
			else                  table.insert(exprt, c) end
		end
	end
	code = table.concat(codet)
	
	return code
end
local function fixCode(code)
	code = code:gsub(",", " ")
	code = code:gsub(":([^\\/])", ":\n%1")
	code = code:gsub("[ \t]+:", ":")
	code = code:gsub("%]", " %] ")
	code = code:gsub("%[", " %[ ")
	code = code:gsub("%*", " %* ")
	code = code:gsub("\n[ \t\r\n]*", "\n")
	code = code:gsub(" +", " ")
	local curScope = ""
	code = code:gsub("(%.?)([a-zA-Z_][a-zA-Z0-9_%.]*)(:?)", function(dot, name, colon)
		if directiveFunctions[name] then return dot..name..colon end
		if dot=="." then
			assert(curScope~="", "scoped label before any unscoped label: "..name)
			name = curScope.."."..name
		end
		if colon==":" and not name:find("^_BRACE_") then curScope = name:match("^[^%.]+") end
		return name..colon
	end)
	
	return code
end
local stringEscapes = { ["\\"] = "\\", ["n"] = "\n", ["r"] = "\r", ["t"] = "\t", ["0"] = "\0", ["\""] = "\"", ["\'"] = "\'", }
local function prefixCode(code, fn) -- fix strings, add line numbers
	local outt = {}
	local outnextnl = {}
	local linenum = 1
	local skipnl = false
	local function last() return outt[#outt] end
	local function out(c) assert(type(c)=="string"); table.insert(outt, c); end
	local function outn(n) out("$"..string.format("%02X", n)) out("\\") end
	local function outnext(c) assert(type(c)=="string"); table.insert(outnextnl, c); end
	local state = "code" -- code, comment, string, stringesc
	
	local lastbracelabel = 0
	local function bracelabel() lastbracelabel = lastbracelabel+1; return "_BRACE_"..lastbracelabel.."_"; end
	local bracestack = {}
	local bracehasmid = {}
	local lastnl = false
	
	out(".ln 1"); out("\n");
	for i = 1, #code do
		local c = code:sub(i, i)
		local cn = code:sub(i+1, i+1)
		
		if state=="code" then
			if     c=="\r"                              then
			elseif c=="\n" or c=="/"                    then
				linenum = linenum+1
				if not skipnl then out("\n") out(".ln "..linenum); out("\n"); end
				lastnl = true
				for _, v in ipairs(outnextnl) do
					if v=="\n" and skipnl then out("\\")
					else                       out(v)    end
				end; outnextnl = {};
				skipnl = false
			elseif c=="#"  or c==";" or (c=="/" and cn=="/") then state = "comment"
			elseif c=="\t" or c==" " then if (not lastnl) then out(" ") end
			elseif c=="\"" then state = "string" lastnl = false
			elseif c=="\\" then skipnl = true; out("\\");
			elseif c:find("^[a-zA-Z0-9_%.:%$%(%)%*,%[%]%+%-%*%/]$") then out(c); lastnl = false
			elseif c=="{" then
				table.insert(bracestack, bracelabel())
				if not lastnl then out(bracestack[#bracestack].."MID") end
				outnext(bracestack[#bracestack].."START:"); outnext("\n");
			elseif c=="}" then
				if not lastnl then out(bracestack[#bracestack].."START") end
				if not bracehasmid[#bracestack] then outnext(bracestack[#bracestack].."MID:"); outnext("\n"); end
				outnext(bracestack[#bracestack].."END:"); outnext("\n");
				bracehasmid[#bracestack] = nil
				bracestack[#bracestack] = nil
			elseif c=="|" then
				if not lastnl then out(bracestack[#bracestack].."END") end
				outnext(bracestack[#bracestack].."MID:"); outnext("\n");
				bracehasmid[#bracestack] = true
			else error("invalid char "..c) end
		elseif state=="comment" then
			if c=="\n" then state = "code" out("\n") lastnl = true end
		elseif state=="string" then
			if     c=="\\" then state = "stringesc"
			elseif c=="\"" then state = "code"
			else outn(c:byte()) end
		elseif state=="stringesc" then
			out(stringEscapes[c] or error("invalid escape "..c)); state = "string";
		end
	end
	assert(#bracestack==0, "unclosed brace")
	local code2 = table.concat(outt)
	
	return code2
end
local function fixFilename(fn)
	fn = fn:gsub("[^a-zA-Z0-9_]", "_")
	return fn
end
local function includeFile(fn)
	local code = readFile(fn)
	code = prefixCode(code, fn)
	local fnf = fixFilename(fn)
	code = ".fn "..fnf.."\n"..code
	code = code:gsub(".include ([^\r\n]+)", function(fn2)
		return "\n"..includeFile(fn2).."\n"..".fn "..fnf.."\n"
	end)
	return code
end
local function instrsFromArch(arch)
	local function arraySize(imms) local s = 1; for i = 1, #imms do s = s*(imms[i].val or error("invalid array size")) end; return s; end
	local instrs = {
		imm8  = function() return 0, true end,
		imm16 = function() return 0, true end,
		byte  = function() return 1, false end,
		word  = function() return 2, false end,
		["byte imm8"] = function() return 0, true end,
		["word imm16"] = function() return 0, true end,
		["byte [ imm8 ]" ] = function(imms) return arraySize(imms)  , false end,
		["byte [ imm16 ]"] = function(imms) return arraySize(imms)  , false end,
		["word [ imm8 ]" ] = function(imms) return arraySize(imms)*2, false end,
		["word [ imm16 ]"] = function(imms) return arraySize(imms)*2, false end,
	}
	local function addMnem(mnem, opcode)
		instrs[mnem] = opcode
		if mnem:find("%*") then instrs[mnem:gsub("%*", "%[").." ]"] = opcode end
	end
	for _, instr in ipairs(arch.instructions) do
		if instr.mnem then
			local mnem = instr.mnem
			mnem = mnem:gsub("([%*%+%-])", " %1 ")
			mnem = trim(mnem):gsub(" +", " ")
			addMnem(mnem, instr.opcode)
			local alias = aliases[trim(mnem)]
			if alias then for _, v in ipairs(alias) do addMnem(v, instr.opcode) end end
		end
	end
	return instrs
end
local function assembleFile(fn, arch)
	local code = includeFile(fn)
	code = preprocessCode(code)
	code = fixCode(code)
	local instrs = instrsFromArch(arch)
	local mem = assembleCode(code, instrs)
	return mem
end

local function mnemsFromArch(arch)
	local mnems = {}
	for _, instr in ipairs(arch.instructions) do
		if instr.mnem then
			mnems[instr.opcode] = instr.mnem
		end
	end
	return mnems
end
local function disassembleMemory(mem, arch)
	print("Disassembly:")
	local mnems = mnemsFromArch(arch)
	local addr = 0
	local function nextByte() local b = mem[addr]; addr = addr+1; return b; end
	local lastaddr = 0
	while addr<=0xFFFF do
		local opcode = nextByte()
		if opcode then
			local line = {}
			local mnem = mnems[opcode] or "???"
			table.insert(line, trim(mnem:gsub("imm[0-9]+", "")))
			local tlen = 1
			for lens in mnem:gmatch("imm([0-9]+)") do local len = tonumber(lens)/8
				if len==1 then
					table.insert(line, "$"..string.format("%02X", nextByte() or 0))
				elseif len==2 then
					table.insert(line, "$"..string.format("%04X", (nextByte() or 0)*256 + (nextByte() or 0)))
				else error("invalid imm len") end
				tlen = tlen + len
			end
			local lineb = {}
			for i = addr-tlen, addr-1 do
				table.insert(lineb, string.format("%02X", mem[i] or 0))
			end
			local lb = table.concat(lineb, " ")
			if lastaddr~=addr-tlen then print("...") end
			print(string.format("%04X", addr-tlen).." | "..(" "):rep(8-#lb)..lb.." | "..table.concat(line, " "))
			lastaddr = addr
		end
	end
	print()
end
local function memToHex(hex)
	local mem = {}
	local addr = 0
	for d in hex:gmatch("[0-9a-fA-F][0-9a-fA-F]") do
		mem[addr] = tonumber(d, 16)
		addr = addr+1
	end
	return mem
end
local function disassembleHex(hex, arch)
	disassembleMemory(memToHex(hex), arch)
end

local function printMemory(mem)
	print("Memory Dump:")
	local anynonempty = false
	local lastbase = -16
	local lastline = ""
	local numreps = 0
	local function closereps(base)
		if numreps~=0 then
			print("(repeated "..numreps.." more times, up to "..string.format("%04X", base+16)..")")
			numreps = 0
		end
	end
	for base = 0, 0xFFF0, 16 do
		local line = {}
		local nonempty = false
		for addr = base, base+15 do
			if addr%4==0 then table.insert(line, " ") end
			if mem[addr]==false then
				nonempty = true
				table.insert(line, "XX ")
			elseif mem[addr] then
				nonempty = true
				table.insert(line, string.format("%02X", mem[addr]).." ")
			else
				table.insert(line, "-- ")
			end
		end
		if nonempty then
			local l = table.concat(line)
			if l~=lastline or base~=lastbase+16 then
				closereps(base)
				if base ~= lastbase+16 then print("...") end
				print(string.format("%04X", base).." |"..l)
			else
				numreps = numreps+1
			end
			lastline = l
			lastbase = base
			anynonempty = true
		end
	end
	closereps(lastbase)
	if not anynonempty then print("Empty") end
	print()
end

local ts = ts or {
	call = function() end,
	eval = function() end,
}
ts.eval [[
	function commandShiftBrick(%x, %y, %z) { commandToServer('shiftBrick', %x, %y, %z); }
	function commandPlantBrick() { commandToServer('plantBrick'); }
]]
local function plantBrickAt(brickpos, pos)
	local dx, dy, dz = pos[1]-brickpos[1], pos[2]-brickpos[2], pos[3]-brickpos[3]
	ts.call("commandShiftBrick", dy, -dx, dz)
	ts.call("commandPlantBrick")
	brickpos[1], brickpos[2], brickpos[3] = pos[1], pos[2], pos[3]
end
local function buildMemory(mem, romsize, offset, len)
	offset = offset or 0
	
	local rombytes = romsize[1]*romsize[2]*romsize[3]/8
	if len and len>rombytes then error("rom not big enough to hold "..len.." bytes (holds "..rombytes..")") end
	if not len then
		for i = 0, 0xFFFF do
			if mem[i] and (i<offset or i>=offset+rombytes) then error("memory does not fit in rom at addr "..string.format("%04X", i)) end
		end
	end
	
	local brickpos = {0, 0, 0}
	for x = 0, romsize[1]-1 do
		for y = 0, romsize[2]-1 do
			for z = 0, romsize[3]-1 do
				local addr = offset + ((romsize[3]/8)*(x + y*romsize[1]) + math.floor(z/8))
				local pow = math.pow(2, z%8)
				local data = (addr>=offset and ((not len) or addr<offset+len) and mem[addr]) or 0
				local bit = math.floor(data/pow)%2
				if bit==1 then plantBrickAt(brickpos, {x, -y, z}) end
			end
		end
	end
end

local function strtovec(str) local v = {}; for word in str:gmatch("[^ \t\r\n]+") do table.insert(v, tonumber(word)) end; return v; end
function AssembleFile(fn, romsizes, offsets, lens) local offset = tonumber(offsets); local len = tonumber(lens); local romsize = strtovec(romsizes);
	local arch = arch8608
	local mem = assembleFile(fn, arch)
	printMemory(mem)
	if #romsize>0 then assert(#romsize==3, "incorrect rom size") end
	buildMemory(mem, romsize, offset, len)
	disassembleMemory(mem, arch)
end
ts.eval [[
	function AssembleFile(%fn, %romsize, %offset, %len) { luacall("AssembleFile", %fn, %romsize, %offset, %len); }
]]

if arg then AssembleFile(arg[1] or "../8608programs/test.asm", "16 16 8", "0", "256") end