#! /usr/bin/env python3

"""
Update U-Boot to the supplied version if the supplied version is newer.
"""

import hashlib
import logging
import os
import re
import shutil
import subprocess
import sys
from datetime import datetime
from logging.handlers import RotatingFileHandler
from pathlib import Path
from sysinfo import gatherRAMInfo, getMMDCRegister

RETURN_SUCCESS = 0
RETURN_NOTING_TO_DO = 0
RETURN_COULD_NOT_READ_UBOOT = 5
RETURN_FAILED_UPDATE = 11

PATH_UBOOT_VERSION_INFO = "/var/run/u-boot_version.txt"
PATH_UBOOT_IMAGE_S = "/boot/u-boot.imx-solo"
PATH_UBOOT_IMAGE_DL = "/boot/u-boot.imx-duallite"
RAM_THRESHOLD: int = 300  # in MB
SOLO_MMDC_REGISTER_VALUE: str = "0x83190000"
DUAL_LITE_MMDC_REGISTER_VALUE: str = "0x841a0000"
UBOOT_VERSION_REGEX = re.compile(
    r"^U-Boot ([\d.]+)-([\w_]+)\.([\w-]+)\+(\w+) \((.+)\)$"
)
VERSION_TIMESTAMP_FORMAT_STRING = "%b %d %Y - %H:%M:%S %z"

CMD_FLASH_ERASE = "flash_erase /dev/mtd0 0 0"
CMD_KOBS = "kobs-ng init -x -v"
UBOOT_VERSION_WITH_MMDC_REGISTER_STR = "May 22 2025 - 04:59:49 +0000"
# "2025-05-22 04:59:49+00:00"
UBOOT_VERSION_WITH_MMDC_REGISTER_DATETIME = datetime.strptime(
    UBOOT_VERSION_WITH_MMDC_REGISTER_STR, VERSION_TIMESTAMP_FORMAT_STRING
)

# defined by the YOCTO build
LOGFILE_OWNER_UID: int = 1000
LOGFILE_OWNER_GID: int = 1000

logger = logging.getLogger("uboot-auto-updater")


def is_file_empty(file_path: Path) -> bool:
    """
    Check if file exists and is empty.

    :return: boolean indicating if the file is empty or does not exist.
    """
    if not file_path.exists():
        return True

    return file_path.stat().st_size == 0


def extract_datetime_from_version(version: str) -> [datetime | None]:
    """
    Extract datetime from version string.

    :param version: The full version string.
    :return: datetime or None
    """
    match = UBOOT_VERSION_REGEX.match(version)

    if not match:
        logger.error(f"Could not match version string: {version}")
        return None

    return datetime.strptime(match.group(5), VERSION_TIMESTAMP_FORMAT_STRING)


def read_installed_uboot_version() -> [datetime | None]:
    """
    Return the version timestamp of the installed uboot bootloader.

    Reads the version information from the file `/var/run/u-boot_version.txt` and
    parses the timestamp from the version information.

    :return: timestamp of the version or None if the version could not be determined.
    """
    version_file_path = Path(PATH_UBOOT_VERSION_INFO)

    if not version_file_path.exists() or not version_file_path.is_file():
        logger.error(f"U-Boot version file '{version_file_path}' does not exist.")
        return None

    with version_file_path.open("r") as version_file:
        content = version_file.read()
        return extract_datetime_from_version(content)


def read_supplied_uboot_version(uboot_path: str) -> [datetime | None]:
    """
    Return the version timestamp of the supplied uboot image.

    :param uboot_path: path to uboot
    :return: timestamp of the version or None if the version could not be determined.
    """
    command = f"strings {uboot_path} | egrep 'U-Boot [0-9]{{4}}' | head -n 1"

    uboot_image_path = Path(uboot_path)

    if not uboot_image_path.exists() or not uboot_image_path.is_file():
        logger.error(f"U-Boot image '{uboot_image_path}' does not exist.")
        return None

    result = subprocess.run(command, shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        logger.error(f"Command failed with return code: {result.returncode}")
        logger.error(f"Error message:\n\t{result.stderr}")
        return None

    return extract_datetime_from_version(result.stdout.strip())


def gather_CPU_core_info() -> int:
    """
    Extract the CPU information and returns the amount of CPU cores.

    :return: amount of CPU cores
    """
    CAT = "/bin/cat"
    CPU_INFO = "/proc/cpuinfo"
    DEFAULT_CPU_CORES = 1
    PATTERN = r"processor[ \t]*:[ \t]+(\d+)"

    try:
        process = subprocess.Popen(
            [CAT, CPU_INFO], stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        stdout, stderr = process.communicate()

        decodedStdOut = stdout.decode()
        results = re.findall(PATTERN, decodedStdOut)
        return len(results) if results else None
    except (
        ValueError,
        TypeError,
        FileNotFoundError,
        PermissionError,
        subprocess.SubprocessError,
    ) as e:
        logger.error(
            f"Failed to read CPU core infos, continue with {DEFAULT_CPU_CORES} CPU cores as default."
        )
        logger.error(f"Error message:\n\t{e}")
        return DEFAULT_CPU_CORES


def flash_bootloader(uboot_path: str):
    """
    Flash a new version of the uboot bootloader.

    :param uboot_path: path to u-boot
    :return:
    """
    logger.info("Starting bootloader flashing process...")

    flash_erase_result = subprocess.run(
        CMD_FLASH_ERASE, shell=True, capture_output=True, text=True
    )

    if flash_erase_result.returncode != 0:
        logger.error(
            f"Command failed with return code: {flash_erase_result.returncode}"
        )
        logger.error(f"Error message:\n\t{flash_erase_result.stderr}")
        # continue the flashing process, the next command might not fail, and the system
        # should be in a consistent state

    process = subprocess.Popen(
        f"{CMD_KOBS} {uboot_path}",
        shell=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        cwd="/tmp",
    )

    stdout, stderr = process.communicate()

    if process.returncode != 0:
        logger.error(f"Command failed with return code: {process.returncode}")
        logger.error(f"Error message:\n\t{stderr}")
        return False

    return True


def hash_file(
    file_path: Path, hash_algorithm: str = "sha256", chunk_size: int = 8192
) -> str:
    """
    Create a hash over a file.

    :param file_path: Path to the file
    :param hash_algorithm: Hash algorithm to be used
    :param chunk_size: Chunk size, used to reduce load on memory
    :return: Returns the hash as a string
    """
    # Create a new hash object using the specified algorithm
    hash_func = hashlib.new(hash_algorithm)

    # Open the file in binary mode and read it in chunks
    with Path(file_path).open("rb") as file:
        while chunk := file.read(chunk_size):
            hash_func.update(chunk)

    # Return the hexadecimal digest of the hash
    return hash_func.hexdigest()


def update_uboot_boot_partition() -> bool:
    """
    Update U-Boot boot partition.

    :return: boolean depending on success
    """
    file_list: list[str] = [
        "boot.scr",
        "boot_selection.scr",
    ]
    ref_boot_partition_path: Path = Path("/boot")
    uboot_boot_partition_path: Path = Path("/boot/uboot")

    logger.info("Start U-Boot boot partition update...")

    if not ref_boot_partition_path.exists():
        logger.error(f"Path '{ref_boot_partition_path}' does not exist.")
        return False

    if not uboot_boot_partition_path.exists():
        logger.error(f"Path '{uboot_boot_partition_path}' does not exist.")
        return False

    for file in file_list:
        ref_file_path: Path = ref_boot_partition_path / file
        file_path: Path = uboot_boot_partition_path / file

        ref_file_hash = hash_file(ref_file_path)
        file_hash = hash_file(file_path)

        logger.debug(f"Checking file '{file_path}'.")

        if ref_file_hash != file_hash:
            logger.info(
                f"File hashes for the file '{file_path}' do not match, replacing it."
            )

            shutil.copy(ref_file_path, file_path)
        else:
            logger.info(f"File hashes for the file '{file_path}' match, skipping.")

    return True


def main():
    """
    Check if an uboot bootloader update is necessary.

    Compares the version of the installed uboot bootloader
    to the version supplied with the Linux image.
    """
    setup_logging()

    logger.debug("Start U-Boot auto update procedure.")

    # The service gets sometimes triggered twice. At the first time, the version file
    # is still empty. Catch this case as early as possible.
    if is_file_empty(Path(PATH_UBOOT_VERSION_INFO)):
        logger.warning("U-Boot version file is empty, exiting.")
        return RETURN_NOTING_TO_DO

    installed_version = read_installed_uboot_version()
    if installed_version is None:
        logger.error("Installed version could not be determined. Aborting.")
        return RETURN_FAILED_UPDATE

    CPUcores = gather_CPU_core_info()
    ram = gatherRAMInfo(logger)
    mmdc_register = getMMDCRegister(logger)
    if not mmdc_register:
        # check if the uboot version is too old where register is not provided yet
        logger.info(
            f"MMDC register not defined given {installed_version} where version {UBOOT_VERSION_WITH_MMDC_REGISTER_DATETIME} is required."
        )
        if UBOOT_VERSION_WITH_MMDC_REGISTER_DATETIME > installed_version:
            if flash_bootloader(PATH_UBOOT_IMAGE_S):
                logger.info(
                    f"Flashing the bootloader {PATH_UBOOT_IMAGE_S} so that MMDC register can be read"
                )
                return RETURN_SUCCESS
        else:
            logging.error("Could not get mmdc_register from uboot.")
            return RETURN_COULD_NOT_READ_UBOOT

    if CPUcores == 1:
        logger.info("Use solo")
        uboot_path = PATH_UBOOT_IMAGE_S
    else:
        if ram > RAM_THRESHOLD and mmdc_register == SOLO_MMDC_REGISTER_VALUE:
            logger.info("Special edge case use solo")
            uboot_path = PATH_UBOOT_IMAGE_S
        else:
            logger.info("Use dual lite")
            uboot_path = PATH_UBOOT_IMAGE_DL

    logger.info(f"CPU cores: {CPUcores}, RAM: {ram}MB, mmdc register: {mmdc_register}")
    logger.info(
        f"Further used u-boot file '{Path(uboot_path).name}' as supplied version."
    )

    supplied_version = read_supplied_uboot_version(uboot_path)
    if supplied_version is None:
        logger.error("Supplied version could not be determined. Aborting.")
        return RETURN_FAILED_UPDATE

    logger.info(
        f"Installed version: {installed_version} / Supplied version: {supplied_version}"
    )

    if supplied_version > installed_version:
        if flash_bootloader(uboot_path):
            logger.info(
                f"Flashing the bootloader from v{installed_version} to the new version "
                f"v{supplied_version} was successful."
            )
    elif (
        ram < RAM_THRESHOLD
        and mmdc_register == SOLO_MMDC_REGISTER_VALUE
        and CPUcores == 2
    ):
        logger.info(
            f"Flashing the bootloader {uboot_path} for solo->dual lite transition"
        )
        if flash_bootloader(uboot_path):
            logger.info("Flash successful")
    else:
        logger.info("U-Boot version is up to date.")

    if update_uboot_boot_partition():
        logger.info("Finished updating the U-Boot boot partition.")
    else:
        logger.error("Failed updating the U-Boot boot partition.")

    logger.info("Finished U-Boot auto update procedure.")

    return RETURN_SUCCESS


def setup_logging():
    """
    Set up the logging facility.
    """
    log_directory: Path = Path(
        os.environ["ECP_WORK_FILES_DIR"] + "/log/"
        if os.environ.get("ECP_WORK_FILES_DIR") is not None
        else "/mnt/persistent-data/uboot/"
    )
    handlers = []

    if log_directory.is_dir():
        log_file_path: Path = log_directory / "uboot-auto-updater.log"

        if not log_file_path.exists():
            log_file_path.touch()

        if (
            log_file_path.stat().st_uid != LOGFILE_OWNER_UID
            or log_file_path.stat().st_gid != LOGFILE_OWNER_GID
        ):
            os.chown(log_file_path, LOGFILE_OWNER_UID, LOGFILE_OWNER_GID)

        # (more or less 105 lines of logs) 4094 chars === 4kB / 64000 === 64kB
        # / 64kB * (1 + 2) = 192kb disk size
        f_handler = RotatingFileHandler(log_file_path, maxBytes=64000, backupCount=2)
        f_handler.setLevel(logging.INFO)
        handlers.append(f_handler)

    s_handler = logging.StreamHandler(stream=sys.stdout)
    s_handler.setLevel(logging.DEBUG)
    handlers.append(s_handler)

    logging.basicConfig(
        level=logging.DEBUG,
        format="%(asctime)s [%(levelname)s] %(message)s",
        handlers=handlers,
    )


if __name__ == "__main__":
    main()
