#!/bin/python
import argparse
import string
from os.path import abspath, dirname

def padTo(args, addr):
    if (args[0] == 0):
        raise ValueError(".PADTO with label or to 0x0000 is not allowed")

    if (addr > args[0]):
        raise ValueError(".PADTO cannot pad to 0x{:04x} when starting from 0x{:04x}".format(args[0], addr))

    return (args[0] - addr) * [0]

labels = {}
constants = {}

instructions = [
    { "opcode": "LD", "params": [
        { "type": ["r", "r"], "format": lambda args, _: [0b01000000 | (args[0] << 3) | args[1]] },
        { "type": ["r", "8b"], "format": lambda args, _: [0b00000110 | (args[0] << 3), args[1]] },
        { "type": ["r", "(HL)"], "format": lambda args, _: [0b01000110 | (args[0] << 3)] },
        { "type": ["(HL)", "r"], "format": lambda args, _: [0b01110000 | args[1]] },
        { "type": ["(HL)", "8b"], "format": lambda args, _: [0b00110110, args[1]] },
        { "type": ["A", "(BC)"], "format": lambda args, _: [0b00001010] },
        { "type": ["A", "(DE)"], "format": lambda args, _: [0b00011010] },
        { "type": ["(BC)", "A"], "format": lambda args, _: [0b00000010] },
        { "type": ["(DE)", "A"], "format": lambda args, _: [0b00010010] },
        { "type": ["A", "(nn)"], "format": lambda args, _: [0b11111010, args[1] & 0xff, args[1] >> 8] },
        { "type": ["(nn)", "A"], "format": lambda args, _: [0b11101010, args[0] & 0xff, args[0] >> 8] },
        { "type": ["A", "(C)"], "format": lambda args, _: [0b11110010] },
        { "type": ["(C)", "A"], "format": lambda args, _: [0b11100010] },
        { "type": ["A", "(n)"], "format": lambda args, _: [0b11110000, args[1]] },
        { "type": ["(n)", "A"], "format": lambda args, _: [0b11100000, args[0]] },
        { "type": ["A", "(HL-)"], "format": lambda args, _: [0b00111010] },
        { "type": ["(HL-)", "A"], "format": lambda args, _: [0b00110010] },
        { "type": ["A", "(HL+)"], "format": lambda args, _: [0b00101010] },
        { "type": ["(HL+)", "A"], "format": lambda args, _: [0b00100010] },
        { "type": ["rr", "16b"], "format": lambda args, _: [0b00000001 | (args[0] << 4), args[1] & 0xff, args[1] >> 8] },
        { "type": ["(nn)", "SP"], "format": lambda args, _: [0b00001000, args[0] & 0xff, args[0] >> 8] },
        { "type": ["SP", "HL"], "format": lambda args, _: [0b11111001] },
        { "type": ["HL", "8b"], "format": lambda args, _: [0b11111000, args[1]] },
    ]},
    { "opcode": "PUSH", "params": [
        { "type": ["rr"], "format": lambda args, _: [0b11000101 | (args[0] << 4)] },
    ]},
    { "opcode": "POP", "params": [
        { "type": ["rr"], "format": lambda args, _: [0b11000001 | (args[0] << 4)] },
    ]},
    { "opcode": "ADD", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10000000 | (args[0])] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10000110] },
        { "type": ["8b"], "format": lambda args, _: [0b11000110, args[0]] },
        { "type": ["SP", "8b"], "format": lambda args, _: [0b11101000, args[1]] },

        # TODO: HL,rr (x9 ?)
    ]},
    { "opcode": "ADC", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10001000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10001110] },
        { "type": ["8b"], "format": lambda args, _: [0b11001110, args[0]] },
    ]},
    { "opcode": "SUB", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10010000 | (args[0])] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10010110] },
        { "type": ["8b"], "format": lambda args, _: [0b11010110, args[0]] },
    ]},
    { "opcode": "SBC", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10011000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10011110] },
        { "type": ["8b"], "format": lambda args, _: [0b11011110, args[0]] },
    ]},
    { "opcode": "CP", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10111000 | (args[0])] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10111110] },
        { "type": ["8b"], "format": lambda args, _: [0b11111110, args[0]] },
    ]},
    { "opcode": "INC", "params": [
        { "type": ["r"], "format": lambda args, _: [0b00000100 | (args[0] << 3)] },
        { "type": ["(HL)"], "format": lambda args, _: [0b00110100] },

        # THIS ONE IS SPECULATIVE
        { "type": ["rr"], "format": lambda args, _: [0b00000011 | (args[0] << 4)] },
    ]},
    { "opcode": "DEC", "params": [
        { "type": ["r"], "format": lambda args, _: [0b00000101 | (args[0] << 3)] },
        { "type": ["(HL)"], "format": lambda args, _: [0b00110101] },

        { "type": ["rr"], "format": lambda args, _: [0b00001011 | (args[0] << 4)] },
    ]},
    { "opcode": "AND", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10100000 | (args[0])] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10100110] },
        { "type": ["8b"], "format": lambda args, _: [0b11100110, args[0]] },
    ]},
    { "opcode": "OR", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10110000 | (args[0])] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10110110] },
        { "type": ["8b"], "format": lambda args, _: [0b11110110, args[0]] },
    ]},
    { "opcode": "XOR", "params": [
        { "type": ["r"], "format": lambda args, _: [0b10101000 | (args[0])] },
        { "type": ["(HL)"], "format": lambda args, _: [0b10101110] },
        { "type": ["8b"], "format": lambda args, _: [0b11101110, args[0]] },
    ]},
    { "opcode": "CCF", "params": [
        { "type": [], "format": lambda args, _: [0b00111111] },
    ]},
    { "opcode": "SCF", "params": [
        { "type": [], "format": lambda args, _: [0b00110111] },
    ]},
    { "opcode": "DAA", "params": [
        { "type": [], "format": lambda args, _: [0b00100111] },
    ]},
    { "opcode": "CPL", "params": [
        { "type": [], "format": lambda args, _: [0b00101111] },
    ]},
    { "opcode": "JP", "params": [
        { "type": ["16b"], "format": lambda args, _: [0b11000011, args[0] & 0xff, args[0] >> 8] },
        { "type": ["HL"], "format": lambda args, _: [0b11101001] },
        { "type": ["cc", "16b"], "format": lambda args, _: [0b11000010 | (args[0] << 3), args[1] & 0xff, args[1] >> 8] }
    ]},
    { "opcode": "JR", "params": [
        { "type": ["8b"], "format": lambda args, _: [0b00011000, args[0]] },
        { "type": ["16b"], "format": lambda args, addr: [0b00011000, (args[0] - addr - 2) & 0xff] },
        { "type": ["cc", "8b"], "format": lambda args, _: [0b00100000 | (args[0] << 3), args[1]] },
        { "type": ["C", "8b"], "format": lambda args, _: [0b00100000 | (3 << 3), args[1]] },
        { "type": ["cc", "16b"], "format": lambda args, addr: [0b00100000 | (args[0] << 3), (args[1] - addr - 2) & 0xff] },
        { "type": ["C", "16b"], "format": lambda args, addr: [0b00100000 | (3 << 3), (args[1] - addr - 2) & 0xff] }
    ]},
    { "opcode": "CALL", "params": [
        { "type": ["16b"], "format": lambda args, _: [0b11001101, args[0] & 0xff, args[0] >> 8] },
        { "type": ["cc", "16b"], "format": lambda args, _: [0b11000100 | (args[0] << 3), args[1] & 0xff, args[1] >> 8] }
    ]},
    { "opcode": "RET", "params": [
        { "type": [], "format": lambda args, _: [0b11001001] },
        { "type": ["cc"], "format": lambda args, _: [0b11000000 | (args[0] << 3)] },
    ]},
    { "opcode": "RETI", "params": [
        { "type": [], "format": lambda args, _: [0b11011001] },
    ]},
    { "opcode": "RST", "params": [
        { "type": ["n"], "format": lambda args, _: [0b11000111 | (args[0] << 3)] },
    ]},
    { "opcode": "DI", "params": [
        { "type": [], "format": lambda args, _: [0b11110011] },
    ]},
    { "opcode": "EI", "params": [
        { "type": [], "format": lambda args, _: [0b11111011] },
    ]},
    { "opcode": "NOP", "params": [
        { "type": [], "format": lambda args, _: [0b00000000] },
    ]},
    { "opcode": "HALT", "params": [
        { "type": [], "format": lambda args, _: [0b01110110] },
    ]},
    { "opcode": "STOP", "params": [
        { "type": [], "format": lambda args, _: [0b00010000, 0b00000000] },
    ]},
    { "opcode": "RLCA", "params": [
        { "type": [], "format": lambda args, _: [0b00000111] },
    ]},
    { "opcode": "RLA", "params": [
        { "type": [], "format": lambda args, _: [0b00010111] },
    ]},
    { "opcode": "RRCA", "params": [
        { "type": [], "format": lambda args, _: [0b00001111] },
    ]},
    { "opcode": "RRA", "params": [
        { "type": [], "format": lambda args, _: [0b00011111] },
    ]},
    { "opcode": "BIT", "params": [
        { "type": ["n", "r"], "format": lambda args, _: [0b11001011, 0b01000000 | (args[0] << 3) | args[1]] },
        { "type": ["n", "(HL)"], "format": lambda args, _: [0b11001011, 0b01000110 | (args[0] << 3)] },
    ]},
    { "opcode": "SET", "params": [
        { "type": ["n", "r"], "format": lambda args, _: [0b11001011, 0b11000000 | (args[0] << 3) | args[1]] },
        { "type": ["n", "(HL)"], "format": lambda args, _: [0b11001011, 0b11000110 | (args[0] << 3)] },
    ]},
    { "opcode": "RES", "params": [
        { "type": ["n", "r"], "format": lambda args, _: [0b11001011, 0b10000000 | (args[0] << 3) | args[1]] },
        { "type": ["n", "(HL)"], "format": lambda args, _: [0b11001011, 0b10000110 | (args[0] << 3)] },
    ]},
    { "opcode": "RLC", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00000000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00000110] },
    ]},
    { "opcode": "RL", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00010000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00010110] },
    ]},
    { "opcode": "RRC", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00001000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00001110] },
    ]},
    { "opcode": "RR", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00011000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00011110] },
    ]},
    { "opcode": "SLA", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00100000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00100110] },
    ]},
    { "opcode": "SWAP", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00110000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00110110] },
    ]},
    { "opcode": "SRA", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00101000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00101110] },
    ]},
    { "opcode": "SRL", "params": [
        { "type": ["r"], "format": lambda args, _: [0b11001011, 0b00111000 | args[0]] },
        { "type": ["(HL)"], "format": lambda args, _: [0b11001011, 0b00111110] },
    ]},
    { "opcode": ".DB", "params": [
     { "type": ["*"], "format": lambda args, _: args },
    ]},
    { "opcode": ".PADTO", "params": [
     { "type": ["16b"], "format": lambda args, addr: padTo(args,addr) },
    ]},
]

registers = {
    "A": 7,
    "B": 0,
    "C": 1,
    "D": 2, # Maybe ?
    "E": 3,
    "H": 4, # Maybe ?
    "L": 5, # Maybe ?
    # (HL) is not a register but is associated with 6

    "BC": 0,
    "DE": 1,
    "HL": 2, # Confirmed
    "AF": 3, # TODO: Only for PUSH & POP
    "SP": 3, # TODO: For everything except PUSH & POP
}

conditions = {
    # Not checked:
    "NZ": 0,
    "Z": 1,
    "NC": 2,
    "C": 3,
}

class Param:
    def __init__(self, value, labels):
        self.type, self.value = self.get_type_value(value.upper().strip(), labels)

    def get_type_value(self, input, labels):
        if input in ['A', 'B', 'C', 'D', 'E', 'H', 'L']:
            return ['r', input], registers[input]
        elif input in ['BC', 'DE', 'HL', 'SP', 'AF']:
            return ['rr', input], registers[input]
        elif len(input) == 4 and input[:2] == '0X':
            return ['8b'], int(input[2:], 16)
        elif len(input) == 6 and input[:2] == '0X':
            return ['16b'], int(input[2:], 16)
        elif len(input) == 8 and input[:3] == '(0X' and input[-1] == ')':
            return ['(nn)'], int(input[3:-1], 16)
        elif len(input) == 6 and input[:3] == '(0X' and input[-1] == ')':
            return ['(n)'], int(input[3:-1], 16)
        elif input in ['NZ', 'Z', 'NC', 'C']:
            return ['cc', input], conditions[input]
        elif input in ['0','1','2','3','4','5','6','7']:
            return ['n'], int(input)
        elif input in ["(HL)", "(BC)", "(DE)", "(C)", "(HL-)", "(HL+)"]:
            return [input], 0
        elif input.startswith('='):
            if labels is None:
                return ['16b'], 0
            else:
                return ['16b'], labels[input[1:]]
        else:
            raise ValueError("Invalid parameter ({})".format(input))


class Instruction:
    def __init__(self, value, labels):
        splitted = value.split(' ')
        self.opcode = splitted[0].upper().strip()
        self.params = [Param(param, labels) for param in splitted[1:] if param.strip()]

    def get_instruction_format(self):
        for instruction in instructions:
            if self.opcode == instruction['opcode']:
                for params in instruction['params']:
                    if len(params["type"]) == 1 and params["type"][0] == "*":
                        return params
                    if len(params["type"]) == len(self.params):
                        for i in range(len(params["type"])):
                            if params["type"][i] not in self.params[i].type:
                                break
                        else:
                            return params
                return None

    def to_bytes(self, address):
        instruction_format = self.get_instruction_format()['format']
        return instruction_format([param.value for param in self.params], address)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("input_file", help="The input file in gbasm")
    parser.add_argument("output_file", help="The output gb rom filename (e.g. basic.rom)")
    args = parser.parse_args()

    lines, _ = preprocess(args.input_file, 0)
    print("\n".join(lines))

    program = assemble(lines)

    output = open(args.output_file, "wb")
    output.write(program)
    output.close()

def isValidHexadecimal(s):
    return all(c in string.hexdigits for c in s) and (len(s) == 2 or len(s) == 4)


def getConstant(name):
    global constants
    if name.startswith("$") and len(name) > 2:
        if isValidHexadecimal(name[1:]):
            return "0x" + name[1:]
        else:
            return getConstant(constants[name[1:].upper()])
    return name


def preprocess(input_file, offset):
    f = open(input_file, "r")
    relative_path = dirname(abspath(input_file))

    global labels
    global constants
    starting_address = offset;
    lines = []
    # Preprocess
    for line in f:
        line_without_comment = line.split(';')[0].strip()

        words = line_without_comment.split(' ')

        if len(words) == 2 and words[0].upper() == ".INCLUDE":
            if not words[1].startswith('"') or not words[1].endswith('"'):
                raise ValueError("Invalid parameter for .INCLUDE ({})".format(words[1]))
            file = relative_path + "/" + words[1][1:-1]
            included_lines, starting_address = preprocess(file, starting_address)
            lines += included_lines
            continue

        if len(words) == 3 and words[0].upper() == ".DEFINE":
            if isValidHexadecimal(words[1]):
                    raise ValueError("\"{}\" is an invalid definition name since it could be mixed up with an hexadecimal value".format(words[1]))

            constants[words[1].upper()] = words[2]
            print(constants)
            continue
        
        if ':' in line_without_comment:
            splitted = line_without_comment.split(':')
            labels[splitted[0].strip().upper()] = starting_address
            line_without_comment = splitted[1].strip()

        if line_without_comment != '':
            words = line_without_comment.replace(",", " ").replace("0xFF00+", "").split(" ")

            for i in range(0, len(words)):
                if words[i].startswith("$") and len(words[i]) > 2:
                    words[i] = getConstant(words[i])

            res = " ".join(words).replace("$", "0x")
            lines.append(res)
            instruction = Instruction(res, None)
            print("Line: " + line)
            print("Instruction: " + instruction.opcode + " " + str(instruction.params))
            print("Valid: " + str(instruction.get_instruction_format()))
            print("Format: " + str(instruction))
            starting_address += len(instruction.to_bytes(starting_address))

    print(labels)

    return lines, starting_address

def assemble(lines):
    program = []
    # Compile
    for line in lines:
        instruction = Instruction(line, labels)
        print("Instruction: " + instruction.opcode + " " + str(instruction.params))
        print("Valid: " + str(instruction.get_instruction_format()))
        print("Format: " + str(instruction))

        program += instruction.to_bytes(len(program))
    return bytearray(program)


if __name__ == "__main__":
    main()