import os
import struct
import tarfile
from pathlib import Path
from elf import ELF64Image


class BSPWriteError(Exception):
    pass


class BSP:
    def __init__(self, out_path, **kwargs):
        if os.path.exists(out_path):
            os.remove(out_path)

        self.__bootstrap_exec_path = None
        self.__path = out_path
        self.__sysroot = kwargs.get('sysroot', None)

        out_dir = os.path.dirname(out_path)
        Path(out_dir).mkdir(parents=True, exist_ok=True)
        self.__tarfile = tarfile.open(name=out_path, mode='x')


    def add_file(self, src, dst):
        if self.__sysroot is not None:
            while src[0] == '/':
                src = src[1:]
            src = os.path.join(self.__sysroot, src)
        print(f'tarfile.add({src}, arcname={dst})')
        self.__tarfile.add(src, arcname=dst)


    def set_bootstrap_program(self, path):
        if self.__sysroot is not None:
            while path[0] == '/':
                path = path[1:]
            path = os.path.join(self.__sysroot, path)
        self.__bootstrap_exec_path = path


    def __write_trailer(self):
        bootstrap_offset = os.path.getsize(self.__path)
        print(f'bootstrap_offset was {bootstrap_offset}')
        bsp = open(self.__path, mode='ab')
        prog = open(self.__bootstrap_exec_path, mode='rb')

        padding = 0x1000 - (bootstrap_offset % 0x1000)
        bsp.write(b'\0' * padding)
        bootstrap_offset += padding
        prog_len = 0
        print(f'bootstrap_offset is now {bootstrap_offset}')

        while True:
            data = prog.read(1024)
            bsp.write(data)
            prog_len += len(data)

            if len(data) < 1024:
                break

        prog.close()

        prog_elf = ELF64Image()
        if prog_elf.load(self.__bootstrap_exec_path) != 0:
            raise BSPWriteError(f'{self.__bootstrap_exec_path} is not a valid ELF binary')

        for ph in prog_elf.ph_list:
            print('{}: {}/{} ({}/{} bytes) {}'.format(
                ph.type,
                ph.offset, ph.vaddr,
                ph.filesz, ph.memsz, ph.flags))

        trailer = struct.pack('>IQIQIQQQQQQQ',
                0xcafebabe,
                0, bootstrap_offset,
                bootstrap_offset, prog_len,
                prog_elf.text.offset, prog_elf.text.vaddr, prog_elf.text.memsz,
                prog_elf.data.offset, prog_elf.data.vaddr, prog_elf.data.memsz,
                prog_elf.entry)

        bsp.write(trailer)

        bsp.close()
        prog.close()


    def __del__(self):
        self.close()


    def close(self):
        self.__tarfile.close()
        if self.__bootstrap_exec_path is None:
            return

        self.__write_trailer()