import logging
import struct
from .objectfile import ObjectFile
from . import CompilerError
from .bitfun import encode_imm32

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


@reloc('ldr_imm12')
def apply_ldr_imm12(reloc, sym, section, reloc_value):
    assert sym.value % 4 == 0
    assert reloc_value % 4 == 0
    offset = (sym.value - (reloc_value + 8))
    U = 1
    if offset < 0:
        offset = -offset
        U = 0
    assert offset < 4096
    section.data[reloc.offset+2] |= (U << 7)
    section.data[reloc.offset+1] |= (offset >> 8) & 0xF
    section.data[reloc.offset+0] = offset & 0xFF

@reloc('adr_imm12')
def apply_adr_imm12(reloc, sym, section, reloc_value):
    assert sym.value % 4 == 0
    assert reloc_value % 4 == 0
    offset = (sym.value - (reloc_value + 8))
    U = 2
    if offset < 0:
        offset = -offset
        U = 1
    assert offset < 4096
    offset = encode_imm32(offset)
    section.data[reloc.offset+2] |= (U << 6)
    section.data[reloc.offset+1] |= (offset >> 8) & 0xF
    section.data[reloc.offset+0] = offset & 0xFF

@reloc('absaddr32')
def apply_absaddr32(reloc, sym, section, reloc_value):
    assert sym.value % 4 == 0
    assert reloc_value % 4 == 0
    offset = sym.value
    section.data[reloc.offset+3] = (offset >> 24) & 0xFF
    section.data[reloc.offset+2] = (offset >> 16) & 0xFF
    section.data[reloc.offset+1] = (offset >> 8) & 0xFF
    section.data[reloc.offset+0] = offset & 0xFF


class Linker:
    """ Merges the sections of several object files and 
        performs relocation """
    def __init__(self):
        self.logger = logging.getLogger('Linker')

    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)
                self.logger.debug('{} {}({})'.format(offsets[in_s.name], iobj, in_s.name))

            # 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)
                self.logger.debug('{} at 0x{:08X} in section {}'.format(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