from collections import namedtuple
import os.path

from tests.elf import ET, SHT
from tests.elfwriter import ElfSection, create_elf_file
from tests.dwarf import DW_AT, DW_FORM, DW_TAG


DwarfAttrib = namedtuple('DwarfAttrib', ['name', 'form', 'value'])
DwarfDie = namedtuple('DwarfAttrib', ['tag', 'attribs', 'children'])
DwarfDie.__new__.__defaults__ = (None,)


def _append_uleb128(buf, value):
    while True:
        byte = value & 0x7f
        value >>= 7
        if value:
            buf.append(byte | 0x80)
        else:
            buf.append(byte)
            break


def _append_sleb128(buf, value):
    while True:
        byte = value & 0x7f
        value >>= 7
        if (not value and not (byte & 0x40)) or (value == -1 and (byte & 0x40)):
            buf.append(byte)
            break
        else:
            buf.append(byte | 0x80)


def _compile_debug_abbrev(cu_die):
    buf = bytearray()
    code = 1
    def aux(die):
        nonlocal code
        _append_uleb128(buf, code)
        code += 1
        _append_uleb128(buf, die.tag)
        buf.append(bool(die.children))
        for attrib in die.attribs:
            _append_uleb128(buf, attrib.name)
            _append_uleb128(buf, attrib.form)
        buf.append(0)
        buf.append(0)
        if die.children:
            for child in die.children:
                aux(child)
    aux(cu_die)
    buf.append(0)
    return buf


def _compile_debug_info(cu_die, little_endian, bits):
    buf = bytearray()
    byteorder = 'little' if little_endian else 'big'

    buf.extend(b'\0\0\0\0')  # unit_length
    buf.extend((4).to_bytes(2, byteorder))  # version
    buf.extend((0).to_bytes(4, byteorder))  # debug_abbrev_offset
    buf.append(bits // 8)  # address_size

    die_offsets = []
    relocations = []
    code = 1
    decl_file = 1
    def aux(die, depth):
        nonlocal code, decl_file
        if depth == 1:
            die_offsets.append(len(buf))
        _append_uleb128(buf, code)
        code += 1
        for attrib in die.attribs:
            if attrib.name == DW_AT.decl_file:
                value = decl_file
                decl_file += 1
            else:
                value = attrib.value
            if attrib.form == DW_FORM.addr:
                buf.extend(value.to_bytes(bits // 8, byteorder))
            elif attrib.form == DW_FORM.data1:
                buf.append(value)
            elif attrib.form == DW_FORM.udata:
                _append_uleb128(buf, value)
            elif attrib.form == DW_FORM.sdata:
                _append_sleb128(buf, value)
            elif attrib.form == DW_FORM.string:
                buf.extend(value.encode())
                buf.append(0)
            elif attrib.form == DW_FORM.ref4:
                relocations.append((len(buf), value))
                buf.extend(b'\0\0\0\0')
            elif attrib.form == DW_FORM.sec_offset:
                buf.extend(b'\0\0\0\0')
            elif attrib.form == DW_FORM.flag_present:
                pass
            elif attrib.form == DW_FORM.exprloc:
                _append_uleb128(buf, len(value))
                buf.extend(value)
            else:
                assert False, attrib.form
        if die.children:
            for child in die.children:
                aux(child, depth + 1)
            buf.append(0)
    aux(cu_die, 0)

    unit_length = len(buf) - 4
    buf[:4] = unit_length.to_bytes(4, byteorder)

    for offset, index in relocations:
        buf[offset:offset + 4] = die_offsets[index].to_bytes(4, byteorder)
    return buf


def _compile_debug_line(cu_die, little_endian):
    buf = bytearray()
    byteorder = 'little' if little_endian else 'big'

    buf.extend(b'\0\0\0\0')  # unit_length
    buf.extend((4).to_bytes(2, byteorder))  # version
    buf.extend(b'\0\0\0\0')  # header_length
    buf.append(1)  # minimum_instruction_length
    buf.append(1)  # maximum_operations_per_instruction
    buf.append(1)  # default_is_stmt
    buf.append(1)  # line_base
    buf.append(1)  # line_range
    buf.append(1)  # opcode_base
    # Don't need standard_opcode_length

    def compile_include_directories(die):
        for attrib in die.attribs:
            if attrib.name != DW_AT.decl_file:
                continue
            dirname = os.path.dirname(attrib.value)
            if dirname:
                buf.extend(dirname.encode('ascii'))
                buf.append(0)
        if die.children:
            for child in die.children:
                compile_include_directories(child)
    compile_include_directories(cu_die)
    buf.append(0)

    decl_file = 1
    directory = 1
    def compile_file_names(die):
        nonlocal decl_file, directory
        for attrib in die.attribs:
            if attrib.name != DW_AT.decl_file:
                continue
            dirname, basename = os.path.split(attrib.value)
            buf.extend(basename.encode('ascii'))
            buf.append(0)
            # directory index
            if dirname:
                _append_uleb128(buf, directory)
                directory += 1
            else:
                _append_uleb128(buf, 0)
            _append_uleb128(buf, 0)  # mtime
            _append_uleb128(buf, 0)  # size
        if die.children:
            for child in die.children:
                compile_file_names(child)
    compile_file_names(cu_die)
    buf.append(0)

    unit_length = len(buf) - 4
    buf[:4] = unit_length.to_bytes(4, byteorder)
    header_length = unit_length - 6
    buf[6:10] = header_length.to_bytes(4, byteorder)
    return buf


def compile_dwarf(dies, little_endian=True, bits=64):
    if isinstance(dies, DwarfDie):
        dies = (dies,)
    assert all(isinstance(die, DwarfDie) for die in dies)
    cu_die = DwarfDie(DW_TAG.compile_unit, [
        DwarfAttrib(DW_AT.comp_dir, DW_FORM.string, '/usr/src'),
        DwarfAttrib(DW_AT.stmt_list, DW_FORM.sec_offset, 0),
    ], dies)

    return create_elf_file(ET.EXEC, [
        ElfSection(
            name='.debug_abbrev',
            sh_type=SHT.PROGBITS,
            data=_compile_debug_abbrev(cu_die),
        ),
        ElfSection(
            name='.debug_info',
            sh_type=SHT.PROGBITS,
            data=_compile_debug_info(cu_die, little_endian, bits),
        ),
        ElfSection(
            name='.debug_line',
            sh_type=SHT.PROGBITS,
            data=_compile_debug_line(cu_die, little_endian),
        ),
        ElfSection(
            name='.debug_str',
            sh_type=SHT.PROGBITS,
            data=b'\0',
        ),
    ], little_endian=little_endian, bits=bits)