#!/usr/bin/python3

from enum import Enum

class MalformedArgument(Exception):
	pass

class ArgType(Enum):
	IMM = 1
	REG = 2
	VAL = 3
	LAB = 4

	def gettype(s):
		if s.isnumeric():
			return (ArgType.IMM, int(s))
		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):
		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):
	NOP = 0
	PUSH = 1, ArgType.VAL
	POP = 2, ArgType.REG
	ADD = 3, ArgType.VAL, ArgType.VAL, ArgType.VAL
	MUL = 4, ArgType.VAL, ArgType.VAL, ArgType.VAL
	DIV = 5, ArgType.VAL, ArgType.VAL, ArgType.VAL
	JL = 6, ArgType.LAB, ArgType.VAL, ArgType.VAL
	CLB = 7, ArgType.LAB
	SYS = 8, 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):
		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, labnum, regnum):
		self.opcode = opcode
		self.args = args
		for a in args:
			if a[0] == ArgType.REG:
				if a[1] < 0 or a[1] >= regnum:
					raise RangeCheckException(a[0],
					                          a[1],
					                          regnum)
			elif a[0] == ArgType.LAB:
				if a[1] < 0 or a[1] >= labnum:
					raise RangeCheckException(a[0],
					                          a[1],
					                          regnum)

	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):
		self.asm.append(Line(ins, args, self.labnum, self.regnum))

	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, labnum=16, regnum=16):
		self.asm = []
		self.labnum = labnum
		self.regnum = regnum