#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# SPDX-FileCopyrightText: Copyright TF-RMM Contributors.
#

"""
Script creating a bundle from app binaries and RMM binary

The script prepends the app binaries to the RMM binary, and generates a branch
instruction at the beginning of the binary file. This way the RMM code can start
running when the execution reaches the beginning of the bundled binary.

A bundled RMM binary has the following structure:
```
    +------------------------------+                          -+
    |     BL rmm_bin_offset        | Generated by this script  |
    |                              |                           |
    | Rest of the header unchanged |                           |
    |..............................|                           |
    |                              |                           + app_bin_file_1
    |      App binary content      |                           |
    |                              |                           |
    +------------------------------+                          -+
    |     Unchanged bin header     |                           |
    |..............................|                           |
    |                              |                           + app_bin_file_2
    |      App binary content      |                           |
    |                              |                           |
    +------------------------------+                          -+
    |                              |                           |
                  ...                                         ...
    |                              |                           |
    +------------------------------+                          -+
    |     Unchanged bin header     |                           |
    |..............................|                           |
    |                              |                           + app_bin_file_n
    |      App binary content      |                           |
    |                              |                           |
    +------------------------------+                          -+
    |                              |                           |
    |      RMM binary content      |                           +rmm_bin
    |                              |                           |
    +------------------------------+                          -+
```
"""

from argparse import ArgumentParser
import logging
import struct
import sys

logger = None


def initial_branch_instruction(offset):
    """Generate the initial branch instruction to jump to RMM text"""
    assert offset > 0
    assert offset % 4 == 0
    imm = offset // 4
    assert imm < (1 << 26)  # imm can be at most 25 bits
    template = 0x94000000
    # Use struct to make sure that the result is a 4 byte integer in
    # little-endian byte order
    return struct.pack("<I", template | imm)


def main():
    """Main function of the script"""

    parser = ArgumentParser(
        description="Create a bundle from the app and RMM binaries."
    )
    parser.add_argument(
        "app_bin_files",
        metavar="APP_BIN_FILE",
        type=str,
        nargs="+",
        help="input application data file(s) for bin generation",
    )
    parser.add_argument(
        "--out-bin",
        metavar="FILE",
        type=str,
        required=True,
        help="the output bin file generated by gen_app_bin.py",
    )
    parser.add_argument(
        "--rmm-bin",
        metavar="FILE",
        type=str,
        required=True,
        help="the RMM bin input file for bin generation",
    )
    parser.add_argument(
        "--log-file-name",
        metavar="FILE",
        type=str,
        required=False,
        default="",
        help="write logs to 'FILE' as well",
    )

    args = parser.parse_args()

    global logger
    logger = logging.getLogger()
    logger.setLevel(logging.DEBUG)
    fmt = logging.Formatter("%(levelname)s: %(message)s")
    console_hdl = logging.StreamHandler(sys.stdout)
    console_hdl.setFormatter(fmt)
    logger.addHandler(console_hdl)

    if args.log_file_name:
        file_hdl = logging.FileHandler(args.log_file_name, mode="w")
        file_hdl.setFormatter(fmt)
        logger.addHandler(file_hdl)

    apps_size = 0
    app_bin_contents = []

    # Collect the contents of the app bin files and concatenate them in a list.
    for app_bin_file_name in args.app_bin_files:
        with open(app_bin_file_name, "rb") as app_bin_file:
            app_bin_content = app_bin_file.read()
            apps_size += len(app_bin_content)
            app_bin_contents.append(app_bin_content)

    # Create the bundled bin file
    with open(args.out_bin, "wb") as out_file:
        # Write the starting branch instruction
        out_file.write(initial_branch_instruction(apps_size))
        # for the first entry, the Initial branch instruction is added in place
        # the first 4 bytes of the padding in the app header.
        start_offset = 4
        for app_bin_content in app_bin_contents:
            out_file.write(app_bin_content[start_offset:])
            # For the rest of the files, write the full header
            start_offset = 0

        # Add the RMM bin file to the bundle
        with open(args.rmm_bin, "rb") as rmm_bin_file:
            out_file.write(rmm_bin_file.read())

    logger.info(
        f"{args.out_bin} was successfully created. Added {len(args.app_bin_files)} app(s)."
    )
    logger.info(f"The offset of the RMM bin is {apps_size} (0x{apps_size:x}) bytes")


if __name__ == "__main__":
    main()