#!/usr/bin/env python

####################################################################################################
#
#   asm-4004 - Assembler for the Intel 4004 Processor
#
#   Copyright 2020 Michael Riegert
#
#   Permission is hereby granted, free of charge, to any person obtaining a copy of this software
#   and associated documentation files (the "Software"), to deal in the Software without
#   restriction, including without limitation the rights to use, copy, modify, merge, publish,
#   distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in all copies or
#   substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING
#   BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
#   NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
#   DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
#   OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
####################################################################################################

# Written with Python 3.9

import os, sys
import argparse

####################################################################################################
####################################################################################################

# instructions supported by the 4004
instr = {
        'nop': [0x0, 0x0],
        'jcn': [0x1],
        'jnt': [0x1, 0x1],
        'jc' : [0x1, 0x2],
        'jz' : [0x1, 0x4],
        'jt' : [0x1, 0x9],
        'jnc': [0x1, 0xA],
        'jnz': [0x1, 0xC],
        'fim': [0x2],
        'src': [0x2],
        'fin': [0x3],
        'jin': [0x3],
        'jun': [0x4],
        'jms': [0x5],
        'inc': [0x6],
        'isz': [0x7],
        'add': [0x8],
        'sub': [0x9],
        'ld' : [0xA],
        'xch': [0xB],
        'bbl': [0xC],
        'ldm': [0xD],
        'wrm': [0xE, 0x0],
        'wmp': [0xE, 0x1],
        'wrr': [0xE, 0x2],
        'wpm': [0xE, 0x3],
        'wr0': [0xE, 0x4],
        'wr1': [0xE, 0x5],
        'wr2': [0xE, 0x6],
        'wr3': [0xE, 0x7],
        'sbm': [0xE, 0x8],
        'rdm': [0xE, 0x9],
        'rdr': [0xE, 0xA],
        'adm': [0xE, 0xB],
        'rd0': [0xE, 0xC],
        'rd1': [0xE, 0xD],
        'rd2': [0xE, 0xE],
        'rd3': [0xE, 0xF],
        'clb': [0xF, 0x0],
        'clc': [0xF, 0x1],
        'iac': [0xF, 0x2],
        'cmc': [0xF, 0x3],
        'cma': [0xF, 0x4],
        'ral': [0xF, 0x5],
        'rar': [0xF, 0x6],
        'tcc': [0xF, 0x7],
        'dac': [0xF, 0x8],
        'tcs': [0xF, 0x9],
        'stc': [0xF, 0xA],
        'daa': [0xF, 0xB],
        'kbp': [0xF, 0xC],
        'dcl': [0xF, 0xD],
        }

# instructions that take a
instr_hex = ['jcn', 'inc', 'isz', 'add', 'sub', 'ld', 'xch', 'bll', 'ldm']

# instructions that are 16-bit and dont address 12-bits
instr_dos = ['jcn', 'fim', 'isz']

# shortened versions of jcn
instr_short = ['jnt', 'jc', 'jz', 'jt', 'jnc', 'jnz']

# instructions that address a pair
instr_ev = ['fim', 'src', 'fin', 'jin']

# instructions that take in data
instr_data = ['ldm', 'bbl']

# instructions that are 16-bit and do address 12-bits
instr_jump = ['jun', 'jms']

# unsupported instructions
instr_unsupported = ['wpm']

# list of registers
reg_list = ['r0', 'r1', 'r2', 'r3', 'r4', 'r5', 'r6', 'r7', 'r8', 'r9', 'r10', 'r11', 'r12', 'r13', 'r14', 'r15', 'p0', 'p1', 'p2', 'p3', 'p4', 'p5', 'p6', 'p7']

symbols = {}

####################################################################################################
####################################################################################################

# Returns the hex value of a register given the register string
def reg_to_hex(reg):
    try:
        if (int(reg) < 16):
            return int(reg)
    except:
        pass
    if (reg not in reg_list):
        print(reg, "is not a register", file=sys.stderr)
        exit(1)

    if 'r' in reg:
        return int(reg.replace('r',''))
    elif 'p' in reg:
        return (int(reg.replace('p','')) << 1)

def symbol_to_hex(sym):
    try:
        return int(sym)
    except:
        pass

    if sym in reg_list:
        return reg_to_hex(sym)

    if isinstance(sym, int):
        return sym
    if "0x" in str(sym):
        if int(str(sym).replace("0x",""), 16) > 0xFFF:
            print("Warning,", sym, "may be too large", file=sys.stderr)
        return int(str(sym).replace("0x",""), 16)
    elif "#" in str(sym):
        if int(str(sym).replace("#",""), 10) > 4095:
            print("Warning,", sym, "may be too large", file=sys.stderr)
        return int(str(sym).replace("#",""), 10)
    else:
        try:
            return symbols[sym]
        except:
            print(sym, "is not a symbol, integer, or address!", file=sys.stderr)
            exit(1)

def parse_symbol(line):
    inst = line.replace(",", "").split(' ')
    try:
        dec = instr[inst[0]][:]
    except:
        return inst[0]
    if inst[0] in instr_dos or inst[0] in instr_jump or inst[0] in instr_short:
        return [-1, -1]
    else:
        return [-1]

def parse_line(line):
    inst = line.replace(",", "").split(' ')
    try:
        dec = instr[inst[0]][:]
    except:
        return inst[0]

    if len(dec) == 1 and inst[0] in instr_data:
        try:
            dec = [(dec[0] << 4) + int(inst[1] if inst[1] != '' else 0)]
        except:
            dec = [(dec[0] << 4) + 0]
    elif len(dec) == 1 and inst[0] in instr_jump:
        dec = [(dec[0] << 4) + (int(symbol_to_hex(inst[1])) >> 8)]
    elif len(dec) == 2:
        dec = [(dec[0] << 4) + symbol_to_hex(dec[1])]
    elif len(dec) == 1 and inst[0] in instr_dos:
        dec = [(dec[0] << 4) + symbol_to_hex(inst[1])]
    elif len(dec) == 1 and inst[0] in instr_ev:
        dec = [(dec[0] << 4) + ((reg_to_hex(inst[1]) >> 1) << 1)]
    else:
        dec = [(dec[0] << 4) + reg_to_hex(inst[1])]

    if inst[0] in ['jin', 'src']:
        dec[0] += 1

    if inst[0] in instr_unsupported:
        print(line, " is unsupported on the 4004", file=sys.stderr)
        exit(1)

    if inst[0] in instr_dos:
        try:
            dec.append(symbol_to_hex(inst[2]))
        except:
            print("Location required: ", line, file=sys.stderr)
            exit(1)

    elif inst[0] in instr_short:
        try:
            dec.append(symbol_to_hex(inst[1]))
        except:
            print("Location required: ", line, file=sys.stderr)
            exit(1)

    elif inst[0] in instr_jump:
        try:
            dec.append(symbol_to_hex(inst[1])%256)
        except:
            print("Location required: ", line, file=sys.stderr)
            exit(1)

    try:
        return bytes(dec)
    except:
        print(dec, " is too large!", file=sys.stderr)
        exit(1)

def reset_rom(num, rom):
    if rom == 0:
        return num
    else:
        return num%(2**(8*(rom > 0))) + rom * (2**(8))

def print_symbols(symbols):
    print("Symbols:")
    for sym in symbols:
        print("\t", sym, "\t\t", hex(symbols[sym]))

def parse_text(args, text):
    # remove blank lines
    text = [x for x in text if x != '']

    # reset rom position
    rom_pos = 0x0

    if args.verbose:
        print()

    # First pass to check symbols
    for line in text[:]:
        line = line.split(';')[0].lstrip()
        if line.replace(" ","") == '':
            continue

        parsed = parse_symbol(line)
        if isinstance(parsed, str):
            if "=" in parsed:
                if "*" in parsed:
                    # print(parsed, symbol_to_hex(parsed.split("=")[1]))
                    rom_pos = symbol_to_hex(parsed.split("=")[1])
                else:
                    symbols[parsed.split("=")[0]] = symbol_to_hex(parsed.split("=")[1])
            elif ".byte" in parsed:
                rom_pos += symbol_to_hex(line.split(" ")[1])
            else:
                symbols[parsed] = rom_pos
            continue
        else:
            rom_pos += len(parsed)

    if args.verbose:
        print_symbols(symbols)
        print()

    rom_pos = 0

    code = []

    # Second pass to translate
    for line in text[:]:

        verbose = args.verbose or (args.line is not None and (rom_pos >= args.line - args.surr and rom_pos <= args.line + args.surr))

        # Remove comments and skip empty lines
        line = line.split(';')[0].lstrip()
        if line.replace(" ","") == '':
            continue

        parsed = parse_line(line)

        if isinstance(parsed, str):
            if "*=" in parsed:
                if verbose:
                    print("\t", parsed)

                if symbol_to_hex(parsed[2:]) <= rom_pos:
                    print("Warning,", parsed, "is too short", file=sys.stderr)
                for i in range(rom_pos, symbol_to_hex(parsed[2:])):
                    nop_parse = parse_line("nop")
                    code.append(nop_parse.hex())
                    output_file.write(nop_parse)
                    rom_pos += 1
                continue

            elif "." in parsed:
                if parsed == ".byte":
                    if verbose:
                        print("\t", line.split(";")[0].lstrip())
                    for i in range(0, int(line.split(" ")[1])):
                        nop_parse = parse_line("nop")
                        code.append(nop_parse.hex())
                        output_file.write(nop_parse)
                        rom_pos += 1
                    continue
                elif parsed == ".end":
                    pass
            elif parsed == "$" or parsed == ".end":
                break

            if verbose:
                print("\t", parsed)

            continue

        if args.verbose or (args.line is not None and (rom_pos >= args.line - args.surr and rom_pos <= args.line + args.surr)):
            if len(parsed) == 2:
                print(" \033[31m$"f"{(rom_pos%2**16):03X}\033[0m", "\t\033[36m", parsed.hex()[:2], parsed.hex()[2:], "\033[0m\t", line)
            else:
                print(" \033[31m$"f"{(rom_pos%2**16):03X}\033[0m", "\t\033[36m", parsed.hex(), "\033[0m\t", line)

        rom_pos += len(parsed)

        if rom_pos > 2**8:
            print("ROM too large!", file=sys.stderr)
            print(rom_pos, " / 256 instructions", file=sys.stderr)
            exit(1)

        code.append(parsed.hex())
        output_file.write(parsed)

    if args.verbose and args.m:
        print()
        print("".join(code))
        print()

    elif args.m:
        print("".join(code))

####################################################################################################
####################################################################################################

parser = argparse.ArgumentParser(description='Assembler for Intel 4004 Processor')
parser.add_argument('-v', '--verbose', action='store_true', help="Verbose output")
parser.add_argument('-m', action='store_true', help="Print binary to stdout and file")
parser.add_argument('-i', '--input', help='Input files', action='append')
parser.add_argument('-o', '--output', help='Output file (Not compatible with multiple input files)')
parser.add_argument('--line', help='Print out lines surrounding hex location in binary file')
parser.add_argument('--surr', help='Number of surrounding lines for --line', type=int, default=2)

args = parser.parse_args()

if len(args.input) > 1 and args.output is not None:
    print("Multifile input doesn't support named output", file=sys.stderr)
    exit(1)

for file in args.input:
    if len(args.input) > 1 or args.output is None:
        args.output = ".".join(file.split(".")[:-1]) + ".bin"

    if args.line is not None:
        args.line = int(args.line, 16)

    if args.verbose:
        print()
        print("Input file:", file)
        print("Output file:", args.output)

    try:
        input_file = open(file, 'r')
    except:
        print("Unable to open file: ", file, file=sys.stderr)
        exit(1)

    try:
        output_file = open(args.output, 'wb')
    except:
        print("Unable to open file: ", args.output, file=sys.stderr)
        exit(1)

    text = input_file.read().lower().replace("\t", " ").split('\n')[:-1]

    parse_text(args, text)

    if args.verbose:
        print()
