#!/usr/bin/python3
#coding:utf-8
#
# Copyright (c) 2021-2022 Amlogic, Inc. All rights reserved.
#
# SPDX-License-Identifier: MIT
#

r"""
OTA Packer: A Packer for Amlogic OTA image.
It takes multiple payload files, packs them into a single OTA image with a
header describing each payload. The header and payloads are aligned to 512
bytes. The payload types can be kernel, dsp firmware, bluetooth firmware, or
wifi firmware. Each payload is signed using mcuboot's imgtool before being
included in the OTA image.
The usage is as follows:

python3 ota_img_packer.py
    -f      The file path of the payload.
    -p      The partition name of the payload.
    -t      The target directory to store the final OTA image.
    -o      The output file name of the OTA image.
"""

import click
import subprocess
import os
import re
from dataclasses import dataclass, fields
from typing import List, Optional, Any, Tuple, Dict

class Field:
    def __init__(self, val: Optional[Any]=None, length: int=0):
        self.val = val
        self.length : int = length

"""
The OTA package image layout is:
    +--------+----------------------+
    |        |   OTA image header   |
    | Header |----------------------+
    |        |   Payload header     |
    +--------+----------------------+
    |          Payload data         |
    +-------------------------------+
    |            ... ...            |
    +-------------------------------+

OTA image header represent that the full ota package information which abstract
by class OTAImgHeader.
Payload header represent that individual payload data information which abstract
by class PayloadHeader.
"""

MCUBOOT_PARA={
    "hd_size" : 0x200,      # MCUBoot Head Size
    "al_size" : 0x1000,     # MCUBoot Head Alignment Size
}

OTA_HEADER_SIZE=0x200

@dataclass
class PayloadHeader:
    img_type: Field = Field(0, 4)
    img_name: Field = Field("", 16)
    img_size: Field = Field(0, 4)
    off_in_payload: Field = Field(0, 4)

@dataclass
class OTAImgHeader:
    magic: Field        = Field("AMLOGIC_UPGRADE",int(16))
    header_size: Field  = Field(int(OTA_HEADER_SIZE), int(4))
    package_size: Field = Field(int(0), int(4))     # header + data
    img_desc_table_offset: Field    = Field(int(0), int(4))
    payload_offset: Field           = Field(int(OTA_HEADER_SIZE), int(4))
    image_count: Field              = Field(0, int(4))

class Platform(object):
    """
    Encapsulates the hardware and software platform configurations
    retrieved from the system environment.

    This class serves as a centralized interface to validate and resolve
    project directories, target architectures, and build environments.
    """

    def __init__(self) -> None:
        self.board      = os.environ.get("BOARD")
        self.arch       = os.environ.get("ARCH")
        self.kernel     = os.environ.get("KERNEL")
        self.product    = os.environ.get("PRODUCT")
        self.scr_dir    = os.path.dirname(os.path.abspath(__file__))
        self.prj_dir    = os.path.dirname(self.scr_dir)

    def get_board(self) -> str:
        if self.board is None:
            click.secho("Please executing `source scripts/env.sh` firstly!",
                fg="red")
            raise Exception("ERROR")
        return self.board

    def get_arch(self) -> str:
        if self.arch is None:
            click.secho("Please executing `source scripts/env.sh` firstly!",
                fg="red")
            raise Exception("ERROR")
        return self.arch

    def get_kernel(self) -> str:
        if self.kernel is None:
            click.secho("Please executing `source scripts/env.sh` firstly!",
                fg="red")
            raise Exception("ERROR")
        return self.kernel

    def get_product(self) -> str:
        if self.product is None:
            click.secho("Please executing `source scripts/env.sh` firstly!",
                fg="red")
            raise Exception("ERROR")
        return self.product

    def get_prf_dir(self) -> str:
        return self.prj_dir

    def get_board_dir(self) -> str:
        return f"{self.prj_dir}/boards/{self.get_arch()}/{self.get_board()}"

    def get_out_dir(self) -> str:
        return (f"{self.prj_dir}/output/{self.get_arch()}-{self.get_board()}-"
                f"{self.get_product()}/{self.get_kernel()}")

class PartParser(object):
    """
    Parses partition configurations by resolving dependencies between build
    artifacts, preprocessor directives, and partition source files.

    This class extracts structured metadata from C-based partition tables by
    leveraging the designated C cross-compiler for macro expansion.
    """

    def __init__(self) -> None:
        self.platform = Platform()
        self.autoconf = f"{self.platform.get_out_dir()}/autoconf.h"
        self.cmakeche = f"{self.platform.get_out_dir()}/CMakeCache.txt"
        self.partfile = f"{self.platform.get_board_dir()}/partition_tbl.c"
        self.compiler = ""
        self.partname = []

    def extract_names(self, lines: List[str]) -> List[str]:
        """
        Extract partition name attributes from localized C-struct code blocks.

        param lines: A list of code strings representing distinct token blocks.
        return: A list of extracted partition name literals.
        """
        return [
            match.group(1)
            for line in lines
            if ".name" in line
            if (match := re.search(r'\.name\s*=\s*"([^"]+)"', line))
        ]

    def parser(self) -> List[str]:
        """
        Execute the pipeline to resolve, preprocess, and parse partition tables.

        The method extracts the C compiler path from the CMake cache, executes
        the compiler preprocessor (`-E`) to expand conditional macros using
        autoconf definitions, tracks brace scope depth to isolate structs, and
        extracts partition identifiers.

        return: An ordered sequence of partition names including the implicit
        bootloader.
        raises subprocess.CalledProcessError: If sub-process execution
        terminates abnormally.
        """
        blocks: List[str] = []
        buf: List[str] = []
        depth = 0

        try:
            with open(self.cmakeche, 'r', encoding='utf-8') as f:
                for line in f:
                    if line.startswith("CMAKE_C_COMPILER:"):
                        self.compiler = line.split('=', 1)[1].strip()
                        break
        except IOError:
            self.compiler = ""

        # Construct sub-process arguments for macro expansion via GCC/Clang
        # preprocessor
        cmd = [self.compiler, "-E", "-include", self.autoconf, self.partfile]
        ret = subprocess.run(cmd, stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE, text=True, check=True)
        lines = ret.stdout.splitlines()

        # Tokenize and track scope depth to isolate structured definition
        # boundaries
        for line in lines:
            if line.lstrip().startswith("#"):
                continue
            if "{" in line:
                depth += line.count("{")
            if depth > 0:
                buf.append(line)
            if "}" in line:
                depth -= line.count("}")
                if depth == 0:
                    blocks.append(" ".join(buf))
                    buf = []

        # Extract name properties and normalize the sequencing structure
        self.partname = self.extract_names(blocks)
        self.partname.insert(0, 'bootloader')

        return self.partname

    def listp(self) -> None:
        """
        Output the compiled partition configuration matrix to stdout.

        Iterates through the verified partition dataset and renders 
        formatted structural records via standard terminal streams.
        """
        click.secho("[ List Partition Table ]", fg="green")
        for idx, name in enumerate(self.parser()):
            click.secho(f"\t[{idx}] : {name}", fg="yellow")
        print()

class OTAImgPacker(object):
    """
    A Packer for Amlogic OTA image. It takes multiple payload files, packs them
    into a single OTA image with a header describing each payload. The header
    and payloads are aligned to 512 bytes. The payload types can be kernel,
    dsp firmware, bluetooth firmware, or wifi firmware. Each payload is signed
    using mcuboot's imgtool before being included in the OTA image.
    """

    def __init__(self,
                 ota_inputs: List[Tuple[str, str]],
                 img: str,
                 target_dir: str,
                 list: bool) -> None:
        self.ota_inputs     = ota_inputs         # [(file, partition)]
        self.ota_outputs    = {"output_dir":target_dir, "output_img":img}
        self.ota_tool_dir   = os.path.dirname(
                                os.path.dirname(os.path.abspath(__file__)))

        """
        This variable represents an extraction of each payload file that
        constitutes full OTA package image.
        [
            {
                "name"  : payload name,
                "path"  : payload file path,
                "header": PayloadHeader(),
            },
            ...,
        ]
        """
        self.payload_info: List[Dict[str, Any]]    = []
        self.ota_img_header             = OTAImgHeader()

        # The offset represent that the offset from payload data.
        self.offset: int                = 0
        self.mcuboot_header_size        = MCUBOOT_PARA['hd_size']
        self.mcuboot_align_size         = MCUBOOT_PARA['al_size']
        self.garbage_file: List[str]    = []
        self.list                       = list
        if self.list:
            self.partition              = PartParser()

    def mcuboot_imgtool(self, in_payload: str, out_signed: str,
            slot_size: str) -> bool:
        """ Sign the payload using mcuboot's imgtool. """
        cmd = [
            "python3",
            f"{self.ota_tool_dir}/lib/mcuboot/scripts/imgtool.py",
            "sign",
            "--key", f"{self.ota_tool_dir}/lib/mcuboot/root-rsa-2048.pem",
            "--header-size", f"{MCUBOOT_PARA['hd_size']}",
            "--align", "4",
            "--slot-size", slot_size,
            "--pad",
            "--version", "1.0.0",
            "--pad-header",
            in_payload,
            out_signed
        ]
        result = subprocess.run(cmd, check=True)
        if result.returncode == 0:
            return True
        return False

    def advance(self, val: int):
        """ Update the payload file offset. """

        self.offset += val

    def fill_payload_header(self,
            payload: Dict[str, Any],
            img_type: int,
            img_name: str,
            img_size: int) -> None:
        """ Fill the payload header with the given information. """

        # Fill payload header
        payload["header"].img_type.val = img_type   # payload type(-f)
        payload["header"].img_name.val = img_name   # payload name(-p)
        payload["header"].img_size.val = img_size   # payload size(calculate)
        payload["header"].off_in_payload.val = self.offset # payload offset

        curret_payload_size = int(payload["header"].img_size.val)
        self.advance(curret_payload_size)   # Update the offset

        # Fill OTA package header
        # Update the package size, which initial value is zero. The updated value
        # is the current package value + payload size.
        package_size_val = self.ota_img_header.package_size.val or 0
        val = curret_payload_size + int(package_size_val)
        self.ota_img_header.package_size.val = val

    def pack_ota_img(self, debug: bool=False) -> bool:
        """ Packaging all payload artifacts into an OTA package image. """

        out_dir=self.ota_outputs.get('output_dir')
        out_img=self.ota_outputs.get('output_img')
        ota_img = f"{out_dir}/{out_img}"

        ota_payload = f"{self.ota_outputs.get('output_dir')}/ota_payload"
        self.garbage_collect(ota_payload)

        ota_img_handle = open(ota_img, "wb")
        ota_payload_handle = open(ota_payload, "wb")

        payload_header = b''

        # Fill the payload header one by one firstly and then serialized it to
        # a binary file.
        # The payload data also can be serialized into a binary file that handle
        # is the ota_payload_handle.
        for i, payload in enumerate(self.payload_info):
            pd_type = payload["name"]
            pd_size = os.path.getsize(payload["path"])

            if self.list:
                pd_part = self.partition.parser()
                if pd_type in pd_part:
                    pd_indx = pd_part.index(pd_type)
                else:
                    self.error(f"The {pd_type} is not in partition table,"
                        f"please check the supported partition table "
                        f"with `-l` option")
            else:
                pd_indx = i

            self.fill_payload_header(payload, pd_indx, pd_type, pd_size)

            if debug:
                click.secho("[ Fill payload header ]", fg="green")
                click.echo(f"\t{click.style('File', fg='yellow')}: "
                        f"{payload.get('path')} \n\t"
                    f"{click.style('Payload Type', fg='yellow')}: "
                        f"{pd_type} {pd_indx}\n\t"
                    f"{click.style('Payload Size', fg='yellow')}: "
                        f"{hex(pd_size)}\n")

            # Serialize payload header + data.
            for f in fields(payload["header"]):
                # The f.name is item of class. The field_obj.val is the
                # connect which should be serialized.
                field_obj = getattr(payload.get("header"), f.name)

                serialied = self.serialize_field(payload["name"], field_obj)
                payload_header += serialied

            # Write the payload data into payload file.
            with open(payload['path'], "rb") as pf:
                ota_payload_handle.write(pf.read())

            # Update the image count
            self.ota_img_header.image_count.val += int(1) # type: ignore

        # Fill the ota image header
        package_size_val = self.ota_img_header.package_size.val or 0

        # The 512B is the size of header which contain both OTA header and
        # payload header.
        val = int(OTA_HEADER_SIZE) + int(package_size_val)
        self.ota_img_header.package_size.val = int(val)         # Package size
        self.ota_img_header.img_desc_table_offset.val = int(36) # image offset

        # Completely generate the payload binary file.
        ota_payload_handle.close()

        ota_header = b''

        # Serialize the ota image header
        for f in fields(self.ota_img_header):
            field_obj = getattr(self.ota_img_header, f.name)
            serialized = self.serialize_field(f.name, field_obj)
            ota_header += serialized

        # Write the ota image header
        ota_img_handle.write(ota_header)
        # Write the payload header
        ota_img_handle.write(payload_header)
        ota_img_handle.close()

        # Pad the remaining space with 0xFF.
        pad_data = b''
        with open(ota_img, "rb") as handle:
            pad_data = handle.read()
            pad_len = OTA_HEADER_SIZE - len(pad_data)

            if pad_len > 0:
                pad_data += b'\xFF' * pad_len
            else:
                self.error("Pad failed!")

        # Write payload data to ota image
        ota_img_handle = open(ota_img, "wb")
        ota_img_handle.write(pad_data)
        with open(ota_payload, "rb") as handle:
            ota_img_handle.write(handle.read())
        ota_img_handle.close()

        click.secho("[ OTA package successful! ]", fg="green")
        self.garbage_clean()
        return True

    def serialize_field(self, name: str, field_obj: Field) -> bytes:
        """
        Serialized the field_obj.val, the first parameter only as a debug
        information. """

        val = field_obj.val
        length = field_obj.length
        if isinstance(val, int):
            max_val = 2 ** (length * 8) - 1
            if val > max_val:
                raise ValueError(f"{name}: {val} exceeds max {max_val}")
            return val.to_bytes(length, byteorder='little', signed=False)
        elif isinstance(val, str):
            if val.isdigit():
                num = int(val)
                max_val = 2 ** (length * 8) - 1
                if num > max_val:
                    raise ValueError(f"{name}: {val} exceeds max {max_val}")
                return num.to_bytes(length, byteorder='little', signed=False)
            else:
                data = val.encode('utf-8')
                if len(data) > length:
                    raise ValueError(f"{name}: string too long ({val})")
                return data.ljust(length, b'\x00')
        else:
            raise TypeError(f"{name}: unsupported type: {type(val)}")

    def error(self, msg: str):
        click.secho("[ Error ]", fg="reg")
        raise Exception(msg)

    def garbage_collect(self, garbage: str):
        """
        This function collects all temporary files generated during execution.
        """
        self.garbage_file.append(garbage)

    def garbage_clean(self):
        """
        This function purges all temporary files identified by garbage collector
        """
        for f in self.garbage_file:
            if os.path.exists(f):
                os.remove(f)

    def pack(self, debug: bool=False):
        """ Pack the OTA image with the payload files. """

        # Sign input files using mcuboot tool to prepare them as payload
        # artifacts.
        for pyload_file, part_name in self.ota_inputs:
            mcuboot_target = f"{self.ota_outputs.get('output_dir')}/{part_name}"

            self.garbage_collect(mcuboot_target)

            # Preparation of payload image components.
            self.payload_info.append(
                {
                    "name": part_name,
                    "path":mcuboot_target,
                    "header": PayloadHeader()
                })

            # The mt means mcuboot target
            mt_size         = os.path.getsize(pyload_file)
            mt_signed_size  = mt_size + MCUBOOT_PARA['hd_size']
            slot_size       = ((mt_signed_size + MCUBOOT_PARA['al_size'] - 1) \
                & ~MCUBOOT_PARA['al_size'] - 1) + MCUBOOT_PARA['al_size']

            if debug:
                click.secho("[ MCUBOOT Sign ]", fg="green")
                print(f"\t{click.style('File', fg='yellow')}: {pyload_file}\n\t"
                    f"{click.style('Size', fg='yellow')}: {hex(mt_size)}\n\t"
                    f"{click.style('Signed Size', fg='yellow')}: "
                        f"{hex(mt_signed_size)}\n\t"
                    f"{click.style('Slot Size', fg='yellow')}: "
                        f"{hex(slot_size)}\n")

            self.mcuboot_imgtool(pyload_file, mcuboot_target, hex(slot_size))

        # Packagimg all payload artifacts into an OTA package image.
        self.pack_ota_img(debug)

@click.command()
@click.option("-f", "--file_name", multiple=True, required=True,
            help="File names")
@click.option("-p", "--partition", multiple=True, required=True,
            help="Partitions")
@click.option("-t", "--target_dir", required=True,
    help="Target directory")
@click.option("-o", "--target_name", required=True,
    help="Output file name")
@click.option('-d', "--debug", is_flag=True, help='Enable debug log')
@click.option('-l', "--list", is_flag=True, help='List the valid partition')
def main(file_name: List[str],
         partition: List[str],
         target_dir: str,
         target_name: str,
         debug: bool,
         list: bool):
    """
    Amlogic OTA packer. It takes multiple payload files, packs them into a
    single OTA image with a header describing each payload.
    """
    if len(file_name) != len(partition):
        raise click.UsageError("The number of files (-f) must match the " \
        "number of partitions (-p).")

    # pairs [Tuple(file_path: str, partition_name: str), ...]
    pairs = [(os.path.abspath(f), p) for f, p in zip(file_name, partition)]

    if debug:
        click.secho("[ Input Parameters ]", fg="green")
        for f, p in pairs:
            click.echo(f"\t{click.style('File', fg='yellow')}: "
                f"{f}\n\t{click.style('Partition', fg='yellow')}: {p}\t")
        click.echo(f"\t{click.style('Output Directory', fg='yellow')}: "
            f"{os.path.abspath(target_dir)}")
        click.echo(f"\t{click.style('OTA Image Name', fg='yellow')}: "
            f"{target_name}\n")

        if list:
            part = PartParser()
            part.listp()

    ota = OTAImgPacker(pairs, target_name, os.path.abspath(target_dir), list)
    ota.pack(debug)

if __name__ == "__main__":
    main()