import struct
from .objectfile import ObjectFile
from . import CompilerError

def align(x, m):
    while ((x % m) != 0):
        x = x + 1
    return x

def wrap_negative(x, bits):
    b = struct.unpack('<I', struct.pack('<i', x))[0]
    mask = (1 << bits) - 1
    return b & mask

reloc_map = {}

def reloc(t):
    def f(c):
        reloc_map[t] = c
    return f


@reloc('lit_add_8')
def apply_lit8(reloc, sym, section, reloc_value):
    assert sym.value % 4 == 0
    offset = (sym.value - (align(reloc_value + 2, 4)))
    assert offset in range(0, 1024, 4), str(offset)+str( self.dst.sections)
    rel8 = offset >> 2
    section.data[reloc.offset] = rel8


@reloc('wrap_new11')
def apply_wrap_new11(reloc, sym, section, reloc_value):
    offset = sym.value - (align(reloc_value, 2) + 4)
    assert offset in range(-2048, 2046, 2)
    imm11 = wrap_negative(offset >> 1, 11)
    section.data[reloc.offset] = (imm11 & 0xff)
    section.data[reloc.offset + 1] |= (imm11 >> 8) & 0x7


@reloc('rel8')
def apply_rel8(reloc, sym, section, reloc_value):
    assert sym.value % 2 == 0
    offset = sym.value - (align(reloc_value, 2) + 4)
    assert offset in range(-256, 254, 2), str(offset) + str(reloc)
    imm8 = wrap_negative(offset >> 1, 8)
    section.data[reloc.offset] = imm8


@reloc('bl_imm11_imm10')
def apply_bl_imm11(reloc, sym, section, reloc_value):
    assert sym.value % 2 == 0
    offset = sym.value - (align(reloc_value, 2) + 4)
    assert offset in range(-16777216, 16777214, 2), str(offset)
    imm32 = wrap_negative(offset >> 1, 32)
    imm11 = imm32 & 0x7FF
    imm10 = (imm32 >> 11) & 0x3FF
    s = (imm32 >> 24) & 0x1
    section.data[reloc.offset + 2] = imm11 & 0xFF
    section.data[reloc.offset + 3] |= (imm11 >> 8) & 0x7
    section.data[reloc.offset] = imm10 & 0xff
    section.data[reloc.offset + 1] |= ((imm10 >> 8) & 0x3) | (s << 2)

@reloc('b_imm11_imm6')
def apply_b_imm11_imm6(reloc, sym, section, reloc_value):
    assert sym.value % 2 == 0
    offset = sym.value - (align(reloc_value, 2) + 4)
    assert offset in range(-1048576, 1048574, 2), str(offset)
    imm32 = wrap_negative(offset >> 1, 32)
    imm11 = imm32 & 0x7FF
    imm6 = (imm32 >> 11) & 0x3F
    s = (imm32 >> 24) & 0x1
    section.data[reloc.offset + 2] = imm11 & 0xFF
    section.data[reloc.offset + 3] |= (imm11 >> 8) & 0x7
    section.data[reloc.offset] |= imm6
    section.data[reloc.offset + 1] |= (s << 2)

# ARM reloc!!
# TODO: move to target classes???
@reloc('b_imm24')
def apply_b_imm24(reloc, sym, section, reloc_value):
    assert sym.value % 4 == 0
    assert reloc_value % 4 == 0
    offset = (sym.value - (reloc_value + 8))
    rel24 = wrap_negative(offset >> 2, 24)
    section.data[reloc.offset+2] = (rel24 >> 16) & 0xFF
    section.data[reloc.offset+1] = (rel24 >> 8) & 0xFF
    section.data[reloc.offset+0] = rel24 & 0xFF


class Linker:
    """ Merges the sections of several object files and 
        performs relocation """
    def link(self, objs, layout={}):
        # Create new object file to store output:
        self.dst = ObjectFile()

        # Create sections with address:
        for section_name, address in layout.items():
            self.dst.get_section(section_name).address = address

        # First copy all sections into output sections:
        for iobj in objs:
            offsets = {}
            # Merge sections:
            for in_s in iobj.sections.values():
                out_s = self.dst.get_section(in_s.name)
                # TODO: align section in other way:
                while out_s.Size % 4 != 0:
                    out_s.add_data(bytes([0]))

                # Add new section:
                offsets[in_s.name] = out_s.Size
                out_s.add_data(in_s.data)


            # Merge symbols:
            for sym in iobj.symbols.values():
                out_s = self.dst.get_section(sym.section)
                value = offsets[sym.section] + out_s.address + sym.value
                self.dst.add_symbol(sym.name, value, sym.section)

            # Merge relocations:
            for reloc in iobj.relocations:
                offset = offsets[reloc.section] + reloc.offset
                self.dst.add_relocation(reloc.sym, offset, reloc.typ, reloc.section)

        # Perform relocations:
        for reloc in self.dst.relocations:
            # Lookup symbol:
            if reloc.sym not in self.dst.symbols:
                raise CompilerError('Undefined reference "{}"'.format(reloc.sym))
            sym = self.dst.symbols[reloc.sym]
            # patch up:
            section = self.dst.get_section(reloc.section)

            # Determine location in memory of reloc patchup position:
            reloc_value = section.address + reloc.offset

            if reloc.typ in reloc_map:
                f = reloc_map[reloc.typ]
                f(reloc, sym, section, reloc_value)
            else:
                raise NotImplementedError('Unknown relocation type {}'.format(reloc.typ))

        return self.dst