#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright © Huawei Technologies Co., Ltd. 2025. All rights reserved.
"""
功 能: 裁剪rtos
版权信息：华为技术有限公司，版本所有(C) 2025
"""

import os
import stat
import shutil
import filecmp
import re
import tempfile
import time

import yaml

from bmcgo import misc
from bmcgo.tasks.task import Task
from bmcgo.tasks import misc as Task_Misc
from bmcgo.utils.data_classes import StripConfig


ZIMAGE_FILE = "zImage"
IMAGE_NAME = "linux-"


class TaskClass(Task):
    def __init__(self, config, work_name=""):
        super().__init__(config, work_name)
        self.rtos_chain = self.config.cpu
        if self.config.rtos_kernel:
            self.rtos_chain += "_" + self.config.rtos_kernel
        self.rtos_tar = os.path.join(self.config.oct_rtos, "rtos_with_ko.tar.gz")
        self.rtos_cfg = os.path.join(self.config.temp_platform_build_dir, "rtos_sdk/rtos/")
        self.rtos_oct = self.get_manufacture_config("base/rtos_oct")
        if self.rtos_oct:
            self.rtos_dist_dir = f"{self.config.tools_path}/rtos-dist"
            self.download_file_old_sha256 = f"{self.config.tools_path}/rtos_dist_download_file_old.sha256"
            self.download_file_new_sha256 = f"{self.config.tools_path}/rtos_dist_download_file_new.sha256"
            self.rtos_dist_new_sha256 = f"{self.config.tools_path}/rtos_dist_new.sha256"
            self.skip_download = False
            self.skip_install = False
            if not os.path.exists(self.rtos_dist_dir):
                os.makedirs(self.rtos_dist_dir, exist_ok=True)
            self.rtos_oct_file = os.path.realpath(os.path.join(self.config.code_path, self.rtos_oct))

    def run(self):
        # 只有带oct配置的才需要生成裁剪
        # 构造oct裁剪工作目录
        self._prepare_env()
        if (os.path.isdir(os.path.join(self.rtos_cfg, "oct_initrd")) and
            os.path.isdir(os.path.join(self.rtos_cfg, "oct_rootfs"))):
            if self.rtos_oct:
                self._download_rtos_dist()
                self._install_rtos_dist()
            oct_instance = Oct(self.config, True, self.rtos_cfg, self.work_name)
            oct_instance.run()
            oct_instance = Oct(self.config, False, self.rtos_cfg, self.work_name)
            oct_instance.run()
        else:
            self._merge_hi1711sdk()

        self.run_command(f"cp {self.rtos_tar} {self.config.sdk_path}")
        self.success(f"rtos_with_ko.tar.gz裁剪完成")

    def _download_rtos_dist(self):
        self.info(f"移除下载路径: {self.rtos_dist_dir}")
        self.run_command(f"rm -rf {self.rtos_dist_dir}", ignore_error=True, sudo=True)
        self.run_command(f"rm -rf {self.download_file_old_sha256}", ignore_error=True, sudo=True)
        partner_tools_dir = f"{os.path.expanduser('~')}/rtos_dist"
        self.info(f"从缓存目录{partner_tools_dir}复制编译器工具")
        self.run_command(f"cp -rf {partner_tools_dir}/. {self.rtos_dist_dir}")
        self.info("下载rtos_dist工具结束")

    def _install_rtos_dist(self):
        self.skip_install = not self.check_need_install(
            self.rtos_dist_dir, Task_Misc.RTOS_DIST_SHA256_PATH, self.rtos_dist_new_sha256
        )
        if self.skip_install:
            self.info("rtos_dist版本匹配, 跳过安装")
            return 
        is_ubuntu = self.tools.is_ubuntu
        self.chdir(self.rtos_dist_dir)
        self.info("安装 dist rpm 包")
        for rpm in os.listdir("./"):
            if not os.path.isfile(rpm) or not rpm.endswith(".rpm"):
                continue
            self.info("安装 {}".format(rpm))
            if not is_ubuntu:
                self.run_command("rpm -ivh {}".format(rpm), sudo=True)
            else:
                self.pipe_command(["rpm2cpio {}".format(rpm), "sudo cpio -id -D /"])

        logname = os.getenv(misc.ENV_LOGNAME, None)
        if logname and logname != "root":
            user_group = f"{os.getuid()}:{os.getgid()}"
            self.run_command(f"chown {user_group} /opt/RTOS -R", sudo=True)
        self.run_command("cp -af {} {}".format(self.rtos_dist_new_sha256, Task_Misc.RTOS_DIST_SHA256_PATH), sudo=True)

    def _merge_hi1711sdk(self):
        self.info("开始构建rtos_with_ko.tar.gz")
        user_group = f"{os.getuid()}:{os.getgid()}"
        self.chdir(self.config.oct_rtos)
        rtos_tar = os.path.join(self.config.sdk_path, "rtos.tar.gz")
        sdk_tar = os.path.join(self.config.sdk_path, "hi1711sdk.tar.gz")
        # 解压SDK
        pack_dir = os.path.join(self.config.oct_rtos, "rtos_with_driver", "rootfs")
        self.run_command(f"rm -rf rtos_with_driver", sudo=True)
        os.makedirs(pack_dir)
        self.run_command(f"tar --xattrs --xattrs-include=* -xf {rtos_tar}", sudo=True)
        self.run_command(f"chown {user_group} {pack_dir}/.. -R", sudo=True)

        sdk_path = os.path.join(self.config.oct_rtos, "sdk")
        self.run_command(f"rm -rf {sdk_path}", sudo=True)
        os.makedirs(sdk_path)
        self.run_command("tar -xf {} -C {}".format(sdk_tar, sdk_path), sudo=True)
        ko_path = f"{pack_dir}/lib/modules/"
        self.run_command(f"chown {user_group} {sdk_path} -R", sudo=True)
        self.run_command("cp -dfr {}/. {}".format(sdk_path, ko_path))
        self.run_command(f"tar --xattrs --xattrs-include=* -czpf {self.rtos_tar} rtos_with_driver")

    def _prepare_env(self):
        python_link = shutil.which("python")
        if not python_link:
            python3_link = shutil.which("python3")
            python_link = python3_link[:-1]
            self.run_command(f"ln -s {python3_link} {python_link}", sudo=True)
        if not os.path.isdir(self.config.oct_rtos):
            os.makedirs(self.config.oct_rtos)
        cmd = f"source {self.config.rtos_path}/dlsetenv.sh -p {self.rtos_chain} --sdk-path={self.config.sysroot}"
        cmd += " > /dev/null"
        env_file = tempfile.NamedTemporaryFile()
        with os.fdopen(os.open(env_file.name, flags=os.O_WRONLY | os.O_CREAT | os.O_TRUNC, \
                mode=stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH), "w") as file_handler:
            file_handler.write("#!/bin/bash\n\n")
            file_handler.write(cmd)
            file_handler.write("\n")
            file_handler.write("printenv\n")
        setenv_file = tempfile.NamedTemporaryFile()
        self.pipe_command(["bash " + env_file.name], setenv_file.name)
        with open(setenv_file.name, "r") as file_handler:
            lines = file_handler.readlines()
        for line in lines:
            line = line.strip()
            if not line:
                continue
            chunks = line.split("=", 1)
            os.environ[chunks[0]] = chunks[1]
        os.environ["RTOS_ARM64_CHAIN"] = self.rtos_chain


class Oct(Task):
    def __init__(self, config, is_initrd, rtos_dir, work_name=""):
        super().__init__(config, work_name)
        self.is_initrd = is_initrd
        self.tag = "initrd" if is_initrd else "rootfs"
        self.oct_dir = os.path.join(rtos_dir, "oct_" + self.tag)
        self.policy_file = os.path.join(self.oct_dir, "policy.yml")
        with open(self.policy_file) as f:
            self.policy: dict = yaml.safe_load(f)
        if is_initrd:
            self.img_file = os.path.join(self.config.oct_rtos, self.tag + "_boot.cpio.gz")
        else:
            self.img_file = os.path.join(self.config.oct_rtos, "rtos_with_ko.tar.gz")
        self.img_temp = os.path.join(self.config.oct_rtos, self.tag + "_boot_temp.cpio.gz")
        self.work_dir = os.path.join(self.config.oct_rtos, self.tag)
        if is_initrd:
            self.pack_dir = os.path.join(self.work_dir, "unpack")
        else:
            self.pack_dir = os.path.join(self.work_dir, "rtos_with_driver", "rootfs")
        if not os.path.isdir(self.work_dir):
            os.makedirs(self.work_dir)
        self.tools.remove_path_and_create(self.pack_dir)

    def run(self):
        os.chdir(self.work_dir)
        if not self._oct_rtos():
            return
        self.pipe_command([f"zcat {self.img_temp}", f"cpio -id -D {self.pack_dir}"])
        self.chdir(self.pack_dir)
        # 删除不必要的文件
        self._run_policy()
        if self.is_initrd:
            self._clean_initrd()
        self._add_sdk_files()
        self._copy_hcc_lib()
        self._strip_image()
        if self.is_initrd:
            self._packet_initrd(self.img_file)
            self._mk_uimage()
        else:
            self._packet_rootfs(self.img_file)
        self._update_file_flag(self.work_dir)

    def _mk_uimage(self):
        self.chdir(self.config.oct_rtos)
        kernel_config = os.path.join(self.config.rtos_path, "bin", "kernel-config")
        cmds = [
            f"cat {kernel_config}",
            "grep Linux/arm64",
            "awk '{print $3}'"
        ]
        fullname, _ = self.tools.pipe_command(cmds, capture_output=True)
        image_name = IMAGE_NAME + fullname

        cmd = os.path.join(self.config.rtos_path, "bin/modules/bin/mkheader")
        cmd += f" -A arm64 -O linux -T kernel -C none -n {image_name} -d {ZIMAGE_FILE} uImage"
        self.run_command(cmd)

    def _check_if_file_changed(self, directory, oct_file):
        # oct文件会影响裁剪
        basename = os.path.basename(oct_file)
        dirname = os.path.dirname(oct_file)
        tmp_sha = f"{directory}/{basename}.sha.temp"
        new_sha = f"{directory}/{basename}.sha.new"
        old_sha = f"{directory}/{basename}.sha"
        self.pipe_command([f"sha256sum {oct_file}", "awk '{print $1}'"], tmp_sha)
        # 策略文件会影响裁剪
        policy = os.path.join(dirname, "policy.yml")
        self.pipe_command([f"sha256sum {policy}", "awk '{print $1}'"], tmp_sha, append=True)
        rtos_md5 = Task_Misc.BUILD_TOOLS_SHA256_PATH
        hi171x_sdk_md5 = Task_Misc.SDK_SHA256_PATH
        self.pipe_command([f"cat {rtos_md5} {hi171x_sdk_md5} {tmp_sha}"], new_sha)
        # 文件不存在时必然要构建
        if not os.path.isfile(old_sha):
            return True
        if filecmp.cmp(old_sha, new_sha):
            return False
        return True

    def _update_file_flag(self, directory):
        for file in os.listdir(directory):
            if not file.endswith(".sha.new"):
                continue
            file = os.path.join(directory, file)
            old_sha = file[:-4]
            if os.path.isfile(old_sha):
                os.unlink(old_sha)
            shutil.copyfile(file, old_sha)

    def _gen_initrd(self, conf, secconf):
        rtos_chain = os.environ.get("RTOS_ARM64_CHAIN")
        self.chdir(os.path.join(self.config.rtos_path, "bin"))
        cmd = f"./oct --dist=../{rtos_chain}/pool/RTOS.dist --conf={conf} --secconf={secconf}"
        cmd += " --nostrip"
        self.run_command(cmd)
        for file in os.listdir("."):
            if file.startswith("rtos-e"):
                self.run_command(f"ar -x {file}")
                self.run_command(f"rm {file}")

    def _oct_rtos(self):
        """
        返回True表示裁剪完成，返回False表示无需裁剪
        """
        cwd = os.getcwd()
        oct_config = self.policy.get("oct", {})
        conf = oct_config.get("conf")
        conf = os.path.join(self.oct_dir, conf)
        secconf = oct_config.get("secconf")
        if secconf and secconf != "no":
            secconf = os.path.join(self.oct_dir, secconf)
            secure_need = self._check_if_file_changed(cwd, secconf)
        else:
            secconf = "no"
            secure_need = False
        rootfs_need = self._check_if_file_changed(cwd, conf)
        if rootfs_need or secure_need:
            self.info("开始构建 " + self.img_temp)
            self._gen_initrd(conf, secconf)
            shutil.move("initrd.cpio.gz", self.img_temp)
            kernel_file = os.path.join(self.config.oct_rtos, ZIMAGE_FILE)
            shutil.move("zImage", kernel_file)
        else:
            self.info(f"缓存有效，无需裁剪 {self.img_temp}")
            return False
        self.chdir(cwd)
        return True

    def _rm_files_exclude(self, directory, roles):
        self.info("特定目录下不符合规则的文件")
        tmpdir = tempfile.TemporaryDirectory()
        real_files = []
        for root, dirs, files in os.walk(directory):
            for file in files:
                real_files.append(f"{root}/{file}")
            for file in dirs:
                real_files.append(f"{root}/{file}")
        for file in real_files:
            for role in roles:
                if not re.match(role, file):
                    continue
                dst_dir = os.path.join(tmpdir.name, os.path.dirname(file))
                if not os.path.isdir(dst_dir):
                    os.makedirs(dst_dir)
                self.run_command(f"cp -rP {file} {dst_dir}", command_echo=False)
        self.run_command(f"ls -al {tmpdir.name}", command_echo=False)
        self.run_command(f"rm -rf {directory}", command_echo=False)
        self.run_command(f"cp -rP {tmpdir.name}/{directory} .", command_echo=False)

    def _rm_files_exclude_busybox(self, directory):
        self.info("删除Initrd镜像特定目录下非busybox相关文件")
        cmds = [
            f"find {directory}",
            "xargs file",
            "grep busybox",
            "grep -v busybox-udhcp",
            "awk '{print $1}'",
            "awk -F ':' '{print $1}'",
        ]
        filelist, _ = self.pipe_command(cmds, capture_output=True, ignore_error=True)
        real_files = filelist.strip().split("\n")
        tmpdir = tempfile.TemporaryDirectory()
        for file in real_files:
            dst_dir = os.path.join(tmpdir.name, os.path.dirname(file))
            if not os.path.isdir(dst_dir):
                os.makedirs(dst_dir)
            self.run_command(f"cp -rP {file} {dst_dir}", command_echo=False)
        self.run_command(f"ls -al {tmpdir.name}", command_echo=False)
        self.run_command(f"rm -rf {directory}", command_echo=False)
        sub_dir = directory.split("/")[0]
        self.run_command(f"cp -rP {tmpdir.name}/{sub_dir} .", command_echo=False)

    def _copy_hcc_lib(self):
        self.info("从编译工具链复制libstdc++和libgcc_s库文件")
        compile_path = self.config.cross_compile_install_path
        self.tools.copy_file_with_prefix(f"{compile_path}/aarch64-target-linux-gnu/lib64/libstdc++.so", "lib64")
        self.tools.copy_file_with_prefix(f"{compile_path}/aarch64-target-linux-gnu/lib64/libgcc_s.so", "lib64")
        os.unlink("lib64/libstdc++.so.6.0.24-gdb.py")

    def _strip_image(self):
        strip_cmd = os.path.join(self.config.cross_compile_install_path, "bin", self.config.strip)
        strip_dir = os.path.relpath(".")
        tmp = tempfile.NamedTemporaryFile()
        strip_config = StripConfig(
            exclude_roles=[],
            include_roles=[],
            no_strip_file_output=tmp.name,
            strip_cmd=strip_cmd,
            strip_dir=strip_dir
        )
        self.tools.strip(strip_config)

    def _packet_initrd(self, img_file):
        cmds = ["find .", "sort", "cpio -H newc -o", "gzip -c"]
        self.pipe_command(cmds, out_file=img_file)
        file_stat = os.stat(img_file)
        # initrd最大5M
        max_size = 5 * 1024 * 1024
        if file_stat.st_size > max_size:
            raise misc.errors.BmcGoException(f"{img_file} is overide, {file_stat.st_size} > {max_size}, "
                                             + "packet_initrd fail!")
        seek = int(max_size / 1024 - 1)
        self.run_command(f"dd if=/dev/zero of={img_file} bs=1K count=1 seek={seek}")

    def _packet_rootfs(self, img_file):
        self.chdir(self.work_dir)
        self.run_command(f"tar -czpf {img_file} rtos_with_driver/rootfs")

    def _clean_initrd(self):
        # 清理initrd内容
        self._rm_files_exclude_busybox("bin")
        self._rm_files_exclude_busybox("sbin")
        self._rm_files_exclude_busybox("usr/sbin")

    def _add_sdk_files(self):
        sdk_file = os.path.join(self.config.sdk_path, "sdk_initrd.tar.gz")
        if not self.is_initrd:
            self.run_command(f"mkdir -p boot")
            self.run_command(f"cp {self.config.oct_rtos}/initrd_boot.cpio.gz boot/initrd_boot.cpio.gz")
            if not os.path.isfile("boot/uImage"):
                self.run_command(f"cp {self.config.oct_rtos}/uImage boot/")
            sdk_file = os.path.join(self.config.sdk_path, "sdk_rootfs.tar.gz")
        self.run_command("tar -xf {} -C .".format(sdk_file))

    def _run_policy(self):
        for file in self.policy.get("nouse_files", []):
            if not os.path.exists(file):
                continue
            if os.path.isdir(file):
                shutil.rmtree(file)
            else:
                os.unlink(file)
        for exclude in self.policy.get("rm_exclude_files", []):
            directory = exclude.get("dir")
            files = exclude.get("files")
            if not directory:
                raise misc.errors.BmcGoException(
                    f"OCT裁剪策略文件({self.policy_file})存在错误，rm_exclude_files节点找到为空的dir节点"
                )
            self._rm_files_exclude(directory, files)
