import argparse
import enum
import lark
import os
import struct
import sys

import obj_pb2

_HERE = os.path.dirname(__file__)
GRAMMAR_FILE = os.path.join(_HERE, 'as.ebnf')

opcodes = {
        'nop'  : lambda: '0000',
        'load' : lambda p0, p1, p2: f'1{p0:x}{p1:x}{(p2 >> 1)&0xf:x}',
        'store': lambda p0, p1, p2: f'2{p0:x}{p1:x}{(p2 >> 1)&0xf:x}',
        'add'  : lambda p0, p1, p2: f'3{p0:x}{p1:x}{p2:x}',
        'sub'  : lambda p0, p1, p2: f'4{p0:x}{p1:x}{p2:x}',
        'or'   : lambda p0, p1, p2: f'5{p0:x}{p1:x}{p2:x}',
        'and'  : lambda p0, p1, p2: f'6{p0:x}{p1:x}{p2:x}',
        'not'  : lambda p0, p1: f'7{p0:x}{p1:x}0',
        'xor'  : lambda p0, p1, p2: f'8{p0:x}{p1:x}{p2:x}',
        'seth' : lambda p0, p1: f'9{p0:x}{p1&0xff:02x}',
        'shr'  : lambda p0, p1, p2: f'a{p0:x}{p1:x}{p2:x}',
        'mul'  : lambda p0, p1, p2: f'b{p0:x}{p1:x}{p2:x}',
        'cmp'  : lambda p0, p1: f'c{p0:x}{p1:x}0',
        'beq'  : lambda p0, p1: f'd{(p1 >> 1)&0xfff:03x}',
        'set'  : lambda p0, p1: f'e{p0:x}{p1&0xff:02x}',
        'bneq' : lambda p0, p1: f'f{(p1 >> 1)&0xfff:03x}',
        '.word': lambda p0: f'{p0:04x}',
        }


class AsTransformer(lark.Transformer):
    def __init__(self):
        self.addr = 0

    def statement(self, s):
        opcode, *params = s
        if params[0] is None:
            params = []
        addr = self.addr
        self.addr += 2
        return ('OP', addr, opcode, *params)

    def opcode(self, o):
        (o, ) = o
        return str(o)

    def reg_offset(self, r):
        if len(r) < 2:
            return r[0], 0
        return tuple(r[:])

    LABEL = str
    HEX_LITTERAL = lambda _, x: int(x[2:], 16)

    def label(self, l):
        (l,) = l
        return ('LBL', l, self.addr)

    def label_ref(self, l):
        (l,) = l
        return l

    def raw_word(self, w):
        (w,) = w
        addr = self.addr
        self.addr += 2
        return ('OP', addr, '.word', w)

    def register(self, r):
        (r,) = r
        if r == 'pc':
            return 14
        if r == 'lr':
            return 13
        if r == 'sp':
            return 12
        return int(r[1:])

    def immediate(self, i):
        (i,) = i
        return int(i)

    def symbol_def(self, s):
        (sym,) = s
        return ('SYM', s[0])

    start = list


def generate_ops(ops, labels, relocs):
    def filter_label(params, pc):
        for p in params:
            if isinstance(p, str):  # label ref
                if len(params) == 1:  # branch
                    yield 14 # pc
                    yield labels[p] - pc - 4
                else:  # set, allow relocs here
                    relocs.append((pc, p))
                    yield 0xff
            elif isinstance(p, tuple):  # reg offset
                yield p[0]
                yield p[1]
            else:
                yield p

    for op in ops:
        addr = op[0]
        fmt = opcodes[op[1]]
        params = filter_label(op[2:], addr)
        yield fmt(*params)


def larkparse(f, debug=False):
    with open(GRAMMAR_FILE) as g:
        asparser = lark.Lark(g.read())
    tree = asparser.parse(f.read())
    if debug:
        print(tree.pretty())
    lines = AsTransformer().transform(tree)
    labels = {l[1]: l[2] for l in lines if l[0] == 'LBL'}
    ops = [l[1:] for l in lines if l[0] == 'OP']
    syms = [(l[1], labels[l[1]]) for l in lines if l[0] == 'SYM']
    relocs = []
    return [int(x, 16) for x in generate_ops(ops, labels, relocs)], syms, relocs


# any record: >BHs, type, length, data
# header: >cccc, "pol0"
# text record: >BHs, name, text_len, text
# reloc record: >pHip, txt_rec_name, offset, target

def write_obj(fout, ops, syms, relocs):
    obj = obj_pb2.ObjFile(header=obj_pb2.Header(magic='pol0'))

    ends = [s[1] for s in syms[1:]] + [2 * len(ops)]
    data = b''.join(struct.pack('>H', op) for op in ops)

    for (name, begin), end in zip(syms, ends):
        section = obj_pb2.Section(name=name, text=data[begin:end])
        obj.sections.append(section)
        for offs, target in relocs:
            if offs >= begin and offs < end:
                section_off = offs - begin
                reloc = obj_pb2.Reloc(section=name, offset=section_off, target=target)
                obj.relocs.append(reloc)

    fout.write(obj.SerializeToString())



def parse_args():
    parser = argparse.ArgumentParser(description='Assemble.')
    group = parser.add_mutually_exclusive_group()
    group.add_argument('--vhdl', action='store_true',
                        help='output VHDL code')
    group.add_argument('--coe', action='store_true',
                        help='output a Vivado COE file')
    parser.add_argument('--debug', action='store_true',
                        help='print the AST')
    parser.add_argument('--input', '-c', type=argparse.FileType('r'),
                        default=sys.stdin, help='input file (default: stdin)')
    parser.add_argument('--output', '-o', type=argparse.FileType('wb'),
                        default=sys.stdout.buffer, help='output file')
    return parser.parse_args()


def main():
    args = parse_args()

    ops, syms, relocs = larkparse(args.input, debug=args.debug)

    if args.vhdl:
        print(',\n'.join(f'x"{x:04x}"' for x in ops))
        for addr, target in relocs:
            print(f'-- reloc: {target} @{addr}')
    elif args.coe:
        assert not relocs, "relocs not supported with coe output"
        print('memory_initialization_radix=16;')
        print('memory_initialization_vector=')
        code = ',\n'.join(f'{x:04x}' for x in ops)
        print(f'{code};')
    else:
        write_obj(args.output, ops, syms, relocs)


if __name__ == "__main__":
    main()