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

import click
import subprocess
import os
import re
from typing import Tuple, List

"""
Firmware Serialization and Structural Modernization Pipeline

Converts heterogeneous raw firmware binary payloads into C-compliant arrays.
The pipeline automates byte-stream extraction, resolves variable naming
conflicts via lexical refactoring, and appends static immutability qualifiers to
optimize target non-volatile memory allocation.

Usage:
    python gen_firmware.py
            -i <INPUT_FILE_1> [-i <INPUT_FILE_2> ...] -o <OUTPUT_C_FILE>
"""

class FWLoader(object):
    """
    Manages the serialization and symbol modernization of firmware binaries.

    This class facilitates the conversion of raw binary payload streams into
    C-compliant byte arrays leveraging system utilities, followed by regular
    expression-based lexical refactoring to enforce immutable type qualifiers.
    """
    def __init__(self, in_path: Tuple[str, ...], out_path: str,
                from_dir: bool) -> None:
        self.in_path    = list(in_path)
        self.out_path   = out_path
        self.fw_name    = [os.path.basename(f).replace('.', '_')
                            for f in self.in_path if from_dir]
        self.from_dir   = from_dir

    def is_file_exit(self, in_file: str) -> bool:
        """
        Validate the existence, structural integrity, and magnitude of a target
        file.
        """
        if (os.path.exists(in_file)
                and os.path.isfile(in_file)
                and os.path.getsize(in_file) > 0):
            return True
        return False

    def bin_to_array(self, in_file: str, out_file: str) -> None:
        """
        Serialize a raw binary payload into a C-style hexadecimal array via xxd.

        Iterates an external sub-process invocation to parse binary blocks into
        textual arrays, appending the result into the specified compilation
        unit.
        """
        if not self.is_file_exit(in_file):
            self.error(f"The input file({in_file}) is not exist")

        try:
            with open(out_file, "ab") as f:
                subprocess.run(
                    ["xxd", "-i", in_file],
                    stdout=f,
                    check=True
                )
        except subprocess.CalledProcessError as e:
            self.error(f"xxd conversion failed: {e}")

    def update_symbol(self, arrays: str, expect: str) -> str:
        """
        Refactor generated C array identifiers to conform to expected interface
        signatures.

        Employs regular expression matching to locate default variable naming
        patterns established by xxd and overwrites them with predefined target
        tokens.
        """
        name = re.search(rf".*unsigned char\s*(\w+{expect})\[\].*", arrays)
        if name and name.group(1):
            ret = re.sub(rf"\b{re.escape(name.group(1))}\b", expect, arrays)
        else:
            self.error(f"Can not find the expenct symbol {expect}")

        lens = re.search(rf".*unsigned int\s*(\w+{expect}_len).*", ret)
        if lens and lens.group(1):
            ret = re.sub(rf"\b{re.escape(lens.group(1))}\b", f"{expect}_len",
                            ret)
        else:
            self.error(f"Can not find the expenct symbol {expect}_len")

        return ret

    def update_symbol_type(self, arrays: str, before: str, after: str) -> str:
        """
        Perform global lexical replacement on data type declaration specifiers.

        Typically executed to append immutable storage qualifiers to array
        structures to maximize runtime flash optimization.
        """
        return re.sub(rf"\b{re.escape(before)}\b", after, arrays)

    def find_firmware_from_dir(self, firmware_dir: List[str]) -> List[str]:
        """
        Scan target directories to isolate valid firmware assets.

        Iterates through provided directory paths, validates their structural
        existence, and executes a filtered search to aggregate paths of discrete
        binary and text payload assets while excluding build scripts.

        :param firmware_dir: A list of directory paths designated for scanning.
        :return: An aggregated sequence of resolved firmware file paths.
        """

        # Validate structural existence of input directory parameters
        for fw in firmware_dir:
            if not os.path.isdir(fw):
                self.error("Invalid input parameter!")

        file_list: list[str] = []

        # Execute filtered traversal across verified directory scopes
        for fw in firmware_dir:
            cmd = (
                rf'find {fw} \( -name "*.txt" -o -name "*.bin" \)'
                rf' ! -name "CMakeLists.txt"')

            result = subprocess.run(
                cmd,
                shell=True,
                capture_output=True,
                text=True,
                check=True
            )

            # Append resolved file invariants to the centralized dataset
            file_list.extend(result.stdout.splitlines())

        return file_list

    def loader(self):
        """
        Execute the centralized firmware modernization and extraction pipeline.

        Purges historical build artifacts, aggregates underlying source binaries
        via serialization pipelines, and dynamically modernizes structural
        variable attributes via discrete textual parsing.
        """
        if os.path.isfile(self.out_path):
            os.system(f"rm -f {self.out_path}")
        elif not os.path.isdir(os.path.dirname(self.out_path)):
            os.makedirs(os.path.dirname(self.out_path))

        if self.from_dir:
            self.in_path = self.find_firmware_from_dir(self.in_path)
            self.fw_name = [os.path.basename(f).replace('.', '_')
                            for f in self.in_path]

        # Append individual binary assets sequentially into a consolidated
        # payload
        for fw in self.in_path:
            self.bin_to_array(fw, self.out_path)

        # Load the original arrays
        with open(self.out_path, 'r', encoding="utf-8") as f:
            arrays = f.read()

        # Execute structural token optimization across registered symbol spaces
        for fw in self.fw_name:
            arrays = self.update_symbol(arrays, fw)

        # Enforce constant allocation paradigms to target read-only sections
        arrays = self.update_symbol_type(arrays, "unsigned",
                                        '__attribute__((used)) const unsigned')

        # Atomic commit of the fully refactored system definition back to disk
        with open(self.out_path, 'w', encoding="utf-8") as f:
            f.write(arrays)

    def error(self, msg: str):
        click.secho(f"[FW ERROR] {msg}", fg="red")
        raise Exception()


class FwClickUsage(click.Command):
    """
    Overwrites the default runtime interface formatting to provide standardized
    system guidance.
    """
    def format_usage(self, ctx: click.Context,
                    formatter: click.formatting.HelpFormatter):
        formatter.write_usage(
            "gen_firmware.py",
            "-i/I FILE_1/DIR_1 [-i/I FILE_2/DIR_2 ... ] -o OUTPUT"
        )

@click.command(cls=FwClickUsage)
@click.option("-i", "--fw_in", multiple=True, help="Firmware binary file path")
@click.option("-I", "--fw_in_dir", multiple=True,
            help="Firmware binary file directory")
@click.option("-o", "--fw_out", required=True,
    help="Firmware output file path")
def main(fw_in: Tuple[str, ...], fw_out: str, fw_in_dir: Tuple[str, ...]):
    if fw_in and fw_in_dir:
        raise click.UsageError(
            "Options -i and -I are mutually exclusive."
        )

    if fw_in:
        fw = FWLoader(fw_in, fw_out, False)
    elif fw_in_dir:
        fw = FWLoader(fw_in_dir, fw_out, True)
    else:
        raise click.BadParameter("The invalid parameters!")
    fw.loader()

if __name__ == "__main__":
    main()