#!/usr/bin/python3

from enum import Enum

class MalformedArgument(Exception):
	pass

def word_2c(w):
	""" Negate a non-negative integer using two's compliment. """
	return (~w + 1) & 0xFFFFFFFF
def ti(w):
	""" Explicitly transform integer into two's compliment representation. """
	return w if w >= 0 else word_wc(-w)
def from_2c(w):
	""" Turn two's compliment word into Python integer. """
	if (w >> 31) & 1 == 0:
		return w
	return -word_2c(w)

class Argument:
	def __init__(self, argtype, val, sign=False):
		self.at = argtype
		self.sign = sign
		self.val = val
	def __str__(self):
		return f'({self.at}, {self.sign}, {self.val})'
	def high_bits(self):
		return int(self.sign) << 1 | (self.at == ArgType.REG)

class ArgType(Enum):
	""" Class denoting the type of an argument to an instruction. """

	IMM = 1
	""" Immediate values are ones that must be numbers (positive or negative). """

	REG = 2
	""" Type of registers. """

	VAL = 3
	""" Type that denotes either immediate values or registers. """

	LAB = 4
	""" Type of labels. """

	def gettype(s):
		""" Parses the type of the argument represented as a string
		    and returns a tuple with the first the first element being
		    the type and the second element being the integer value of
		    the argument.
		"""
		if s.isnumeric():
			return Argument(ArgType.IMM, int(s))
		elif s[0] == "-" and s[1:].isnumeric():
			return Argument(ArgType.IMM, word_2c(int(s[1:])))
		elif s[0] == 'r' and s[1:].isnumeric():
			return Argument(ArgType.REG, int(s[1:]))
		elif s[0] == 'l' and s[1:].isnumeric():
			return Argument(ArgType.LAB, int(s[1:]))
		else:
			raise MalformedArgument(s)

	def typecheck(self, s):
		""" Parses the type of the string and returns it if it fits
		    the type of the enum value. """
		t = ArgType.gettype(s)
		if self == ArgType.VAL:
			if t.at == ArgType.REG or t.at == ArgType.IMM:
				return t
			else:
				return None
		elif t.at == self:
			return t
		else:
			return None

class OpcodeException(Exception):
	pass
class TypecheckLenException(Exception):
	def __init__(self, opcode, insargs, argtypelen):
		self.opcode = opcode
		self.insargs = insargs
		self.argtypelen = argtypelen
	def __str__(self):
		return f'arguments {insargs} to opcode {self.opcode} not length {self.argtypelen}'
class TypecheckException(Exception):
	def __init__(self, argtype, sarg, i, opcode):
		self.argtype = argtype
		self.sarg = sarg
		self.i = i
		self.opcode = opcode
	def __str__(self):
		return f'opcode {self.opcode} has invalid value {self.sarg} (expected {self.argtype} in position {self.i}'

class Instruction(Enum):
	""" Class of microcode instructions. The first number is the opcode
	    and the suceeding values are the types of each of the
	    arguments. The first argument is the opcode and the second
	    argument is if the argument is a signed variant of another
	    opcode. """
	NOP = 0, False
	PUSH = 1, False, ArgType.VAL
	POP = 2, False, ArgType.REG
	ADD = 3, False, ArgType.REG, ArgType.VAL, ArgType.VAL
	MUL = 4, False, ArgType.REG, ArgType.VAL, ArgType.VAL
	DIV = 5, False, ArgType.REG, ArgType.VAL, ArgType.VAL
	SDIV = "DIV", True, ArgType.REG, ArgType.VAL, ArgType.VAL
	SYS = 6, False, ArgType.VAL
	CLB = 7, False, ArgType.LAB
	JL = 8, False, ArgType.LAB, ArgType.VAL, ArgType.VAL
	JLS = "JL", True, ArgType.LAB, ArgType.VAL, ArgType.VAL
	JLE = 9, False, ArgType.LAB, ArgType.VAL, ArgType.VAL
	JLES = "JLE", True, ArgType.LAB, ArgType.VAL, ArgType.VAL
	JE = 10, False, ArgType.LAB, ArgType.VAL, ArgType.VAL
	JNE = 11, False, ArgType.LAB, ArgType.VAL, ArgType.VAL

	def __int__(self):
		return self.opcode

	def __init__(self, opcode, signed_instruction, *args):
		if type(opcode) is int and (opcode > 0x7F or opcode < 0):
			raise OpcodeException(opcode)

		self.opcode = opcode
		self.argtypes = args
		if signed_instruction:
			self.render = self._render_change_args
		else:
			self.render = self._default_render

	def typecheck(self, sargs):
		""" Pass arguments to the instruction and check if the
		    arguments are correct. """
		rargs = []
		if len(sargs) != len(self.argtypes):
			raise TypecheckLenException(self.opcode, sargs,
			                            len(self.argtypes))
		for i in range(0, len(sargs)):
			t = self.argtypes[i].typecheck(sargs[i])
			if t is None:
				raise TypecheckException(self.argtypes[i],
				                         sargs[i],
				                         i, self.opcode)
			rargs.append(t)
		return rargs

	# The following will be called using OPCODE.render() instead of being
	# called directly.

	def _render_change_args(self, args):
		for i in range(0,len(args)):
			if args[i].at != ArgType.LAB:
				args[i].sign = True
		return Instruction[self.opcode].render(args)

	def _default_render(self, args):
		b = bytes([self.opcode])
		for a in args:
			l = 2 if a.val < 0x80 else None
			bex = encode_pseudo_utf8(a.val, a.high_bits(), l)
			b = b + bex
		return b + bytes([0])

encoding_types = {
#           start      mask   B
	2: (0x7F,       0xC0, 7),
	3: (0xFFF,      0xE0, 12),
	4: (0x1FFFF,    0xF0, 17),
	5: (0x3FFFFF,   0xF8, 22),
	6: (0x7FFFFFF,  0xFC, 27),
	7: (0xFFFFFFFF, 0xFE, 32),
#	B : Total number of bits excluding high bits
}

class InvalidNumberException(Exception):
	pass
class InvalidLengthException(Exception):
	pass
def encode_pseudo_utf8(n, high_bits, to):
	if n < 0:
		raise InvalidNumberException(n)
	if to is None or to < 0:
		for k in sorted(encoding_types):
			if n <= encoding_types[k][0]:
				to = k
				break
		if to is None:
			raise InvalidNumberException(n)
	if to > 8 or to < 0:
		raise InvalidLengthException(to)
	elif to == 1:
		if n < 0x80:
			return bytes([n])
		else:
			raise InvalidNumberException(n,to)

	(maxval, start_byte, n_tot) = encoding_types[to]
	if n > maxval or high_bits > 15:
		raise InvalidNumberException(n, high_bits)
	n = n | (high_bits << n_tot)
	all_bytes = []
	for i in range(0, to - 1):
		all_bytes.append(0x80 | (n & 0x3F))
		n >>= 6
	all_bytes.append(start_byte | n)
	return bytes(reversed(all_bytes))

class RangeCheckException(Exception):
	pass
class Line:
	def __init__(self, ins, args):
		self.ins = ins
		self.args = args

	def check_line(self, lablen, reglen):
		for a in self.args:
			if a.at == ArgType.REG:
				if a.val < 0 or a.val >= reglen:
					raise RangeCheckException(a.at,
					                          a.val,
					                          reglen)
			elif a.at == ArgType.LAB:
				if a.val < 0 or a.val >= lablen:
					raise RangeCheckException(a.at,
					                          a.val,
					                          reglen)

	def __call__(self):
		return self.ins.render(self.args)

class InstructionNotFoundException(Exception):
	pass
class Program:
	def asm_push_line(self, ins, args):
		l = Line(ins, args)
		l.check_line(self.lablen, self.reglen)
		self.asm.append(l)

	def parse_asm_line(self, line):
		line = line.split()
		line[0] = line[0].casefold()
		try:
			# TODO: is there no better way to do this in Python?
			ins = Instruction[line[0].upper()]
		except Exception as e:
			raise InstructionNotFoundException(line[0])

		args_w_type = ins.typecheck(line[1:])
		self.asm_push_line(ins, args_w_type)

	def parse_lines(self, lines):
		for l in lines:
			self.parse_asm_line(l)

	def __call__(self):
		b = bytes()
		for line in self.asm:
			b = b + line()
		return b

	def __init__(self, lablen=16, reglen=16):
		self.asm = []
		self.lablen = lablen
		self.reglen = reglen