#!/usr/bin/env python3
#* Copyright (c) 2020 by UNISOC Technology(Shanghai), Inc.  ALL RIGHTS RESERVED.
#*
#* These coded instructions, statements, and computer programs are the
#* copyrighted works and confidential proprietary information of
#* UNISOC Technology(Shanghai), Inc. They may not be modified, copied, reproduced,
#* distributed, or disclosed to third parties in any manner, medium, or form,
#* in whole or in part, without the prior written consent of UNISOC Technology(Shanghai), Inc.

from __future__ import print_function
from sys import argv, exit, stderr
from argparse import ArgumentParser, FileType, Action
from os import fstat
from struct import pack
from hashlib import sha1
import sys

def set_pad_file(f, padding):
    pad = (padding - (f.tell() & (padding - 1))) & (padding - 1)
    f.write(pack(str(pad) + 'x'))

def create_header(args):
    magic = 'UNISOCBOOTIMAGE!'.encode()
    args.output.write(pack('16s', magic))
    args.output.write(pack('5I',
        getfilesize(args.kernel),                          # size in bytes
        args.base + args.kernel_offset,                 # physical load addr
        getfilesize(args.ramdisk),                         # size in bytes
        args.base + args.ramdisk_offset,                # physical load addr
        args.pagesize))                                  # flash page size we assume
    args.output.write(pack('I', getfilesize(args.dt)))             # device tree

    set_pad_file(args.output, args.pagesize)
    return 0

def getfilesize(f):
    if f is None:
        return 0
    try:
        print('size ', fstat(f.fileno()).st_size)
        return fstat(f.fileno()).st_size
    except OSError:
        return 0

def write_padded_file(f_out, f_in, padding):
    if f_in is None:
        return
    f_out.write(f_in.read())
    set_pad_file(f_out, padding)


def parse_int(x):
    return int(x, 0)


def parse_cmdline():
    parser = ArgumentParser()
    parser.add_argument('--kernel', help='path to the kernel', type=FileType('rb'),
                        required=True)
    parser.add_argument('--ramdisk', help='path to the ramdisk', type=FileType('rb'))
    parser.add_argument('--base', help='base address', type=parse_int, default=0x10000000)
    parser.add_argument('--kernel_offset', help='kernel offset', type=parse_int, default=0x00008000)
    parser.add_argument('--ramdisk_offset', help='ramdisk offset', type=parse_int, default=0x01000000)
    parser.add_argument('--pagesize', help='page size', type=parse_int,
                        choices=[2**i for i in range(11,15)], default=2048)
    parser.add_argument('--dt', help='device tree size', type=FileType('rb'))
    parser.add_argument('-o', '--output', help='output file name', type=FileType('wb'),
                        required=True)
    return parser.parse_args()


def create_file(args):
    write_padded_file(args.output, args.kernel, args.pagesize)
    write_padded_file(args.output, args.ramdisk, args.pagesize)
    write_padded_file(args.output, args.dt, args.pagesize)


def main():
    args = parse_cmdline()
    create_header(args)
    create_file(args)

if __name__ == '__main__':
    main()