#!/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 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 (ArgType.IMM, int(s))
		elif s[0] == "-" and s[1:].isnumeric():
			return (ArgType.IMM, word_2c(int(s[1:])))
		elif s[0] == 'r' and s[1:].isnumeric():
			return (ArgType.REG, int(s[1:]))
		elif s[0] == 'l' and s[1:].isnumeric():
			return (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[0] == ArgType.REG or t[0] == ArgType.IMM:
				return t
			else:
				return None
		elif t[0] == 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. """
	NOP = 0
	PUSH = 1, ArgType.VAL
	POP = 2, ArgType.REG
	ADD = 3, ArgType.REG, ArgType.VAL, ArgType.VAL
	MUL = 4, ArgType.REG, ArgType.VAL, ArgType.VAL
	DIV = 5, ArgType.REG, ArgType.VAL, ArgType.VAL
	SDIV = 6, ArgType.REG, ArgType.VAL, ArgType.VAL
	JL = 7, ArgType.LAB, ArgType.VAL, ArgType.VAL
	CLB = 8, ArgType.LAB
	SYS = 9, ArgType.VAL

	def __init__(self, opcode, *args):
		if opcode > 0x7F or opcode < 0:
			raise OpcodeException(opcode)

		self.opcode = opcode
		self.argtypes = args

	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

encoding_types = {
#           start      mask   A  B
	2: (0x7F,       0xC0, 7),
	3: (0xFFF,      0xE0, 12),
	4: (0x1FFFF,    0xF0, 16),
	5: (0x3FFFFF,   0xF8, 21),
	6: (0x7FFFFFF,  0xFC, 26),
	7: (0xFFFFFFFF, 0xFE, 36),
#	A : number of bits in start byte
#	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, opcode, args):
		self.opcode = opcode
		self.args = args

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

	def __call__(self):
		b = bytes([self.opcode])
		for a in self.args:
			l = 2 if a[1] < 0x80 else None
			if a[0] == ArgType.REG:
				b = b + encode_pseudo_utf8(a[1],1,l)
			else:
				b = b + encode_pseudo_utf8(a[1],0,l)
		return b + bytes([0])

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 = getattr(Instruction, line[0].upper())
		except Exception as e:
			raise InstructionNotFoundException(line[0])

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

	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