#!/usr/bin/env python3
import sys
import os
import shutil
import tarfile
import struct
from lib.manifest import Manifest, Component
from lib.elf import ELF64Image


def reset():
    if len(sys.argv) < 3:
        print("USAGE: {} reset <manifest-path>".format(sys.argv[0]))

    manifest_path = sys.argv[2]
    manifest = Manifest(manifest_path)
    manifest.reset()
    manifest.save()
    return 0


def add_headers():
    if len(sys.argv) < 3:
        print("USAGE: {} add-headers <manifest-path> <component-name> <source-directory> <dest-directory>".format(sys.argv[0]))
        return -1

    manifest_path = sys.argv[2]
    component_name = sys.argv[3]
    dest_path = sys.argv[4]
    src_path = sys.argv[5]

    manifest = Manifest(manifest_path)
    manifest.load()

    component = manifest.get_component(component_name)
    component.add_headers(src_path, dest_path)

    manifest.save()

    return 0


def add_binary():
    if len(sys.argv) < 3:
        print("USAGE: {} add-binary <manifest-path> <component-name> <source-file> <dest-directory>".format(sys.argv[0]))
        return -1

    manifest_path = sys.argv[2]
    component_name = sys.argv[3]
    dest_path = sys.argv[4]
    src_path = sys.argv[5]

    manifest = Manifest(manifest_path)
    manifest.load()

    component = manifest.get_component(component_name)
    component.add_binary(src_path, dest_path)

    manifest.save()
    return 0


def build_bsp():
    if len(sys.argv) < 4:
        print("USAGE: {} build-bsp <manifest-path> <bootstrap-program-path> <dest-path>".format(sys.argv[0]))
        return -1

    manifest_path = sys.argv[2]
    bootstrap_path = sys.argv[3]
    bsp_path = sys.argv[4]

    if os.path.exists(bsp_path):
        os.remove(bsp_path)

    bsp_file = tarfile.open(name=bsp_path, mode='x')

    manifest = Manifest(manifest_path)
    manifest.load()

    for n, c in manifest.get_all_components().items():
        for h in c.get_headers():
            header_src = h['src']
            header_dest = h['dest']
            while header_dest.startswith('/'):
                header_dest = header_dest[1:]

            for f in os.listdir(header_src):
                bsp_file.add(
                        os.path.join(header_src, f),
                        arcname=os.path.join(header_dest, f))

        for b in c.get_binaries():
            binary_src = b['src']
            binary_dest = b['dest']
            while binary_dest.startswith('/'):
                binary_dest = binary_dest[1:]
            binary_dest = os.path.join(binary_dest, os.path.basename(binary_src))

            bsp_file.add(binary_src, arcname=binary_dest)

    bsp_file.close()
    bootstrap_offset = os.path.getsize(bsp_path)

    bsp = open(bsp_path, mode='ab')
    prog = open(bootstrap_path, mode='rb')

    padding = bootstrap_offset % 0x1000
    bsp.write(b'\0' * padding)
    bootstrap_offset += padding
    prog_len = 0

    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(bootstrap_path) != 0:
        print('failed to parse bootstrap program')

    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()

    return 0


if len(sys.argv) < 2:
    print("USAGE: {} <command> [args...]".format(sys.argv[0]))
    exit(-1)

commands = {
        'reset': reset,
        'add-headers': add_headers,
        'add-binary': add_binary,
        'build-bsp': build_bsp,
}

cmd_name = sys.argv[1]
if cmd_name not in commands:
    print("FATAL: unrecognised command '{}'".format(cmd_name))
    exit(-1)

try:
    exit(commands[cmd_name]())
except RuntimeError as e:
    print('FATAL: {}'.format(e))
    exit(-1)