#!/usr/bin/python3

from enum import Enum

class MalformedArgument(Exception):
	pass

def word_2c(w):
	""" Negate a non-negative integer using 32 bit two's compliment.

	:param w: An integer in two's compliment. A non-negative Python
	integer will work.
	:return: The negation of the integer stored as a two's compliment
	integer.
	"""
	return (~w + 1) & 0xFFFFFFFF
def ti(w):
	""" Explicitly transform integer into two's compliment representation.

	:param w: A python integer.
	:return: The integer in two's compliment.
	"""
	return w if w >= 0 else word_wc(-w)
def from_2c(w):
	""" Turn two's compliment word into Python integer.

	:param w: An integer in 32 bit twos compliment.
	:return: The integer as a proper Python string.
	"""
	if (w >> 31) & 1 == 0:
		return w
	return -word_2c(w)

class Argument:
	""" Class of arguments. Not used directly: It is used to store
	intermediate information during the assembly process. """

	def __init__(self, argtype, val, sign=False):
		""" Initialize an argument.

		:param argtype: Type of the argument (instance of ArgType).
		:param val: Python integer value of the argument.
		:param sign: If the argument should be treated as signed.
		Otherwise, the integer will be interpreted in execution
		as an unsigned integer.
		"""
		self.at = argtype
		self.sign = sign
		self.val = val
	def __str__(self):
		return f'({self.at}, {self.sign}, {self.val})'
	def high_bits(self):
		""" Returns the high bits that the argument would have
		    in the opcode. """
		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.

		    Valid parameters are:

		    * `r` followed by a nonnegative integer (register)
		    * `l` followed by a nonnegative integer (label)
		    * any integer (immediate value)

		    :param s: String representing the argument.
		    :return: The Argument object representing the argument.
		    :raises MalformedArgument:
		"""
		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:])), True)
		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.

		    :param s: String argument representing an argument.
		    :return: The Argument class containing the object, or None
		    if the string does not fit the type of self. """
		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):
	""" Exception thrown when arguments to an instruction are of the
	    incorrect length. """
	def __init__(self, opcode, insargs, argtypelen):
		self.opcode = opcode
		self.insargs = insargs
		self.argtypelen = argtypelen
	def __str__(self):
		return f'''\
arguments {self.insargs} to opcode {self.opcode} not of length {self.argtypelen}\
'''
class TypecheckException(Exception):
	""" Exception thrown when an argument to an instruction are of the
	    incorrect type. """
	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 what function is used to compile the instruction
	    (some instructions are actually versions of other instructions). """
	NOP = 0, "_render_default"
	PUSH = 1, "_render_default", ArgType.VAL
	POP = 2, "_render_default", ArgType.REG
	ADD = 3, "_render_default", ArgType.REG, ArgType.VAL, ArgType.VAL
	MOV = "ADD", "_render_mov", ArgType.REG, ArgType.VAL
	MUL = 4, "_render_default", ArgType.REG, ArgType.VAL, ArgType.VAL
	DIV = 5, "_render_default", ArgType.REG, ArgType.VAL, ArgType.VAL
	SDIV = "DIV", "_render_change_args", ArgType.REG, ArgType.VAL, ArgType.VAL
	SYS = 6, "_render_default", ArgType.VAL
	CLB = 7, "_render_default", ArgType.LAB
	JL = 8, "_render_default", ArgType.LAB, ArgType.VAL, ArgType.VAL
	JLS = "JL", "_render_change_args", ArgType.LAB, ArgType.VAL, ArgType.VAL
	JLE = 9, "_render_default", ArgType.LAB, ArgType.VAL, ArgType.VAL
	JLES = "JLE", "_render_change_args", ArgType.LAB, ArgType.VAL, ArgType.VAL
	JE = 10, "_render_default", ArgType.LAB, ArgType.VAL, ArgType.VAL
	J = "JE", "_render_j", ArgType.LAB
	JNE = 11, "_render_default", ArgType.LAB, ArgType.VAL, ArgType.VAL

	def __int__(self):
		""" Returns the opcode associated with the Instruction.
		    If it is a virtual instruction, it will resolve the string
		    name of the opcode and return its opcode. """
		if type(self.opcode) is int:
			return self.opcode
		return int(Instruction[self.opcode])

	def __init__(self, opcode, renderfun, *args):
		""" Initialize an Instruction. Do not call this function: it is
		    used to make enum values. To add a new instruction, modify
		    the Instruction enum.
		    This function sometimes takes string arguments because
		    certain values may not be loaded until later.

		    :param opcode: Opcode of the instruction, or a string
		    containing the case-sensitive name of the instruction from
		    which this instruction derives from.

		    :param renderfun: a string with the name of a function
		    in the class that returns the instruction opcode.

		    :param *args: Type of each argument to the instruction.
		    The amount of arguments denotes the amount of instructions.
		"""

		if type(opcode) is int and (opcode > 0x7F or opcode < 0):
			raise OpcodeException(opcode)

		self.opcode = opcode
		self.argtypes = args
		self.render = getattr(self, renderfun)

	def typecheck(self, sargs):
		""" Pass arguments to the instruction and check if the
		    arguments are correct.

		    :param sargs: List of arguments to the instruction
		    as strings.
		    :return: List of arguments (as Argument objects).
		    :raises TypeCheckException:
		    :raises TypecheckLenException:
		"""
		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

	def _render_mov(self, args):
		args = [args[0], args[1], Argument(ArgType.IMM, 0)]
		return Instruction[self.opcode].render(args)

	def _render_j(self, args):
		args = [args[0], Argument(ArgType.IMM, 0),
		        Argument(ArgType.IMM, 0)]
		return Instruction[self.opcode].render(args)

	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 _render_default(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