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')
class LitAdd8:
    def apply(self):
        pass


class Linker:
    def set_symbol(self, sym):
        self.dst.add_symbol(sym.name, sym.value, sym.section)

    def link(self, objs):
        # Create new object file to store output:
        self.dst = ObjectFile()
        self.dst.get_section('code').address = 0x08000000

        # 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)
                offsets[in_s.name] = out_s.address + out_s.Size
                out_s.add_data(in_s.data)
                # TODO: align section

            # Merge symbols:
            for sym in iobj.symbols.values():
                self.set_symbol(sym)

            # Merge relocations:
            for reloc in iobj.relocations:
                self.dst.add_relocation(reloc.sym, reloc.offset, reloc.typ, reloc.section)

        # Check that there are no more unresolved symbols:
        # 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)

            if reloc.typ == 'lit_add_8':
                assert sym.value % 4 == 0
                offset = (sym.value - (align(reloc.offset + 2, 4)))
                assert offset in range(0, 1024, 4)
                rel8 = offset >> 2
                section.data[reloc.offset] = rel8
            elif reloc.typ == 'wrap_new11':
                offset = sym.value - (align(reloc.offset, 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
            elif reloc.typ == 'rel8':
                assert sym.value % 2 == 0
                offset = sym.value - (align(reloc.offset, 2) + 4)
                assert offset in range(-256, 254, 2), str(offset) + str(reloc)
                imm8 = wrap_negative(offset >> 1, 8)
                section.data[reloc.offset] = imm8
            elif reloc.typ == 'bl_imm11_imm10':
                assert sym.value % 2 == 0
                offset = sym.value - (align(reloc.offset, 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)
            elif reloc.typ == 'b_imm11_imm6':
                assert sym.value % 2 == 0
                offset = sym.value - (align(reloc.offset, 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)
            else:
                raise NotImplementedError('Unknown relocation type {}'.format(reloc.typ))

        return self.dst