#!/usr/bin/env python
# coding: utf-8
# Copyright (c) 2024 Huawei Technologies Co., Ltd.
# openUBMC is Licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#         http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
import os
import stat
import subprocess
import logging
import multiprocessing
import shutil

from qemu_schema import load_config, print_config

QEMU_TEMP = "temp/qemu_temp"


def gzip_cpio_map_dir(input_dir, cpio_path):
    # 创建openUBMC软件包解压路径,并解压openUBMC软件包
    cpio_name = os.path.basename(cpio_path).rsplit('.', 1)[0]
    subprocess.run(['/bin/rm', '-rf', f'{input_dir}/cpio'], check=True)
    os.makedirs(f'{input_dir}/cpio', exist_ok=True)
    shutil.copy(cpio_path, f'{input_dir}/cpio')
    cur_path = os.getcwd()
    os.chdir(f'{input_dir}/cpio')
    subprocess.run(['/bin/gzip', '-d', cpio_path.split("/")[-1]], check=True, stdout=None, stderr=None)
    os.system(f"/bin/cpio -idm < {cpio_name}")
    os.chdir(cur_path)


class DockerRun:
    def __init__(self, cfg, image_name):
        self.cfg = cfg
        self.check_image_exists(image_name)

    @staticmethod
    def check_docker_installed():
        """检查是否安装了Docker"""
        try:
            subprocess.run(["/usr/bin/docker", "--version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            logging.info("Docker 已安装，继续执行...")
        except subprocess.CalledProcessError as e:
            logging.warning("当前环境未安装docker，请安装docker后再执行")
            raise OSError("未识别到Docker，请安装Docker后再试") from e

    @staticmethod
    def pull_image(image_name):
        """从指定镜像仓下载镜像"""
        try:
            subprocess.run(["/usr/bin/docker", "pull", image_name], check=True)
            logging.info(f"镜像 {image_name} 下载完成")
        except subprocess.CalledProcessError as e:
            logging.warning(f"下载镜像 {image_name} 时出错：{e}")

    @staticmethod
    def delete_image(image_name):
        """删除Docker镜像"""
        try:
            subprocess.run(["/usr/bin/docker", "rmi", image_name], check=True)
            logging.info(f"镜像 {image_name} 已删除")
        except subprocess.CalledProcessError as e:
            logging.error(f"删除镜像 {image_name} 时出错：{e}")

    @staticmethod
    def run_docker_container(container_name, image_name, mount_path,
                             mapports):
        """运行Docker容器"""
        try:
            command = [
                "/usr/bin/docker", "run", "-itd", "--name", f"{container_name}",
                "-v", f"{mount_path}:/root/workspace/qemu_bin",
                "-p", f"{mapports['ssh_port']}:2202", "-p", f"{mapports['telnet_port']}:2302",
                "-p", f"{mapports['ipmi_port']}:6230/udp", "-p", f"{mapports['https_port']}:4430", 
                "-p", f"{mapports['snmp_port']}:1610/udp", "-p", f"{mapports['console_port']}:4444", 
                image_name, "/bin/bash", "/root/workspace/qemu_bin/start.sh"
            ]
            subprocess.run(command, check=True)
        except subprocess.CalledProcessError as e:
            logging.error(f"运行容器时出错：{e}")
    
    def check_image_exists(self, image_name):
        """检查Docker镜像是否存在"""
        try:
            subprocess.run(["/usr/bin/docker", "inspect", image_name],
                            check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            logging.info(f"镜像 {image_name} 已存在，继续执行...")
            return True
        except subprocess.CalledProcessError:
            logging.warning(f"镜像 {image_name} 不存在，将从指定镜像仓下载...")
            self.pull_image(image_name)
            return False

    def run(self, operation, image_name):
        self.check_docker_installed()
        if operation == "update":
            if self.check_image_exists(image_name):
                self.delete_image(image_name)
            self.pull_image(image_name)
            logging.info("镜像更新完成，但不会拉起容器")
        elif operation == "delete":
            if self.check_image_exists(image_name):
                self.delete_image(image_name)
            else:
                logging.info(f"镜像 {image_name} 不存在，无需删除")
        else:
            logging.error("无效的操作参数，请使用 'create'、'update' 或 'delete'")


class Runvemake:
    def __init__(self, cfg=None):
        if cfg is None:
            cfg = {}
    
        # 容器配置
        container_cfg = cfg.get("container_cfg", {})
        self.container = container_cfg.get("container", "hi1711")
        self.map_dir = os.path.join(container_cfg.get("inner_path", ""), "output", "data")
        self.qemu_bin = container_cfg.get("qemu_setup_dir", "")
        self.vm = container_cfg.get("container", "hi1711")  # 默认值hi1711
        self.wsl = container_cfg.get("wsl_enabled", False)
        self.docker = container_cfg.get("docker_enabled", False)
        self.operation = container_cfg.get("docker_operator", "")
 
        self.nums = container_cfg.get("docker_num", 1) if self.docker else 0  # 仅在docker启用时有效
        self.docker_images = container_cfg.get("docker_image", 
                    "swr.cn-north-4.myhuaweicloud.com/openubmc/qemu:latest") if self.docker else ""

        if self.docker:
            self.docker_run = DockerRun(cfg, self.docker_images)

        # Qemu启动依赖配置
        dependencies_cfg = cfg.get("start_dependencies", {})
        self.qemu_setup_dir = dependencies_cfg.get("qemu_setup_dir", f"{QEMU_TEMP}/qemu_start_temp")
        self.inner_path = dependencies_cfg.get("inner_path", "output/packet/inner")
        self.cpio_path = dependencies_cfg.get("cpio_path", "output/packet/inner/openUBMC_25.00.00.01_qemu.cpio.gz")
        self.zImage_path = dependencies_cfg.get("zImage_path", "output/packet/inner/zImage_openUBMC")
        self.dtb = os.path.join(self.qemu_setup_dir, dependencies_cfg.get("dtb", "hi1711_9p8c.dtb"))
        self.qemu = os.path.join(self.qemu_setup_dir, dependencies_cfg.get("qemu", "qemu-system-aarch64-release"))
        self.images = dependencies_cfg.get("inner_path", "")
        self.remote_qemu = dependencies_cfg.get("remote_qemu", False)
        self.remote_cpio = dependencies_cfg.get("remote_cpio", False)
        self.download_qemu_url = dependencies_cfg.get("download_qemu",
                        "https://repo.openubmc.cn/25.03/tools/qemu_setup.zip ")
        self.download_cpio_url = dependencies_cfg.get("download_cpio", 
                        "https://repo.openubmc.cn/25.03/firmware/openUBMC-CMT_25.00.00.01_qemu.zip ")
        
        if not os.path.exists(self.inner_path):
            os.makedirs(self.inner_path)

        if not os.path.exists(self.qemu_setup_dir):
            os.makedirs(self.qemu_setup_dir)

        # 端口映射配置
        mapports = cfg.get("mapports", {})
        self.ssh_mapport = mapports.get("ssh", 10022)
        self.https_mapport = mapports.get("https", 10443)
        self.ipmi_mapport = mapports.get("ipmi", 10623)
        self.snmp_mapport = mapports.get("snmp", 10161)
        self.telnet_mapport = mapports.get("telnet", 10023)
        self.qemu_monitor_console = mapports.get("monitor", 14444)
        
        # Qemu性能配置
        qemu_cfg = cfg.get("qemu_cfg", {})
        self.cpus = qemu_cfg.get("core_num", 2)
        self.memory = qemu_cfg.get("memory", 2)

    def download_qemu(self):
        subprocess.run(
            ["/bin/wget", "--timestamping",
              self.download_qemu_url, "--no-check-certificate"],
            cwd=QEMU_TEMP,
            check=True
        )
        self.unzip_qemu()
        
    def unzip_qemu(self):
        if not os.path.exists(f"{QEMU_TEMP}/download"):
            os.makedirs(f"{QEMU_TEMP}/download")
        subprocess.run(
            ["/bin/unzip", "-o", self.download_qemu_url.split("/")[-1], "-d", "download"],
            cwd=QEMU_TEMP,
            check=True
        )

    def download_cpio(self):
        # 从URL中提取文件名
        filename = os.path.basename(self.download_cpio_url)

        # 确保文件名以.cpio.gz结尾
        if not filename.endswith('.cpio.gz'):
            logging.warning(f"警告: 下载URL中的文件名{filename}不以.cpio.gz结尾")

        subprocess.run(
            ["/bin/wget", "--timestamping",
            "-O", filename,
            self.download_cpio_url, "--no-check-certificate"],
            cwd=QEMU_TEMP,
            check=True
        )

    def check_dependencies(self):
        if self.remote_qemu:
            self.download_qemu()
            subprocess.run(["/bin/cp",
                             f"{QEMU_TEMP}/download/{self.dtb.split('/')[-1]}", self.dtb], check=True)
            subprocess.run(["/bin/cp", 
                            f"{QEMU_TEMP}/download/{self.qemu.split('/')[-1]}", self.qemu], check=True)
            subprocess.run(["/bin/cp",
                             f"{QEMU_TEMP}/download/{self.zImage_path.split('/')[-1]}", self.zImage_path], check=True)
        
        if not os.path.exists(self.dtb):
            if not os.path.exists(f'{QEMU_TEMP}/download/{self.dtb.split("/")[-1]}'):
                if not os.path.exists(f'{QEMU_TEMP}/{self.download_qemu_url.split("/")[-1]}'):
                    self.download_qemu()
                else:
                    self.unzip_qemu()
            subprocess.run(["/bin/cp", f'{QEMU_TEMP}/download/{self.dtb.split("/")[-1]}', self.dtb], check=True)
        
        if not os.path.exists(self.qemu):
            if not os.path.exists(f'{QEMU_TEMP}/download/{self.qemu.split("/")[-1]}'):
                if not os.path.exists(f'{QEMU_TEMP}/{self.download_qemu_url.split("/")[-1]}'):
                    self.download_qemu()
                else:
                    self.unzip_qemu()
            subprocess.run(["/bin/cp", f'{QEMU_TEMP}/download/{self.qemu.split("/")[-1]}', self.qemu], check=True)
        
        if not os.path.exists(self.zImage_path):
            if not os.path.exists(f'{QEMU_TEMP}/download/{self.zImage_path.split("/")[-1]}'):
                if not os.path.exists(f'{QEMU_TEMP}/{self.download_qemu_url.split("/")[-1]}'):
                    self.download_qemu()
                else:
                    self.unzip_qemu()
            subprocess.run(["/bin/cp", f'{QEMU_TEMP}/download/{self.zImage_path.split("/")[-1]}', 
                            self.zImage_path], check=True)
            
        if self.remote_cpio or not os.path.exists(self.cpio_path):
            self.download_cpio()
            subprocess.run(["/bin/cp", f'{QEMU_TEMP}/{self.cpio_path.split("/")[-1]}',
                             self.cpio_path], check=True)

    def setenv(self):
        subprocess.run(["/bin/rm", "-rf", f"{self.map_dir}/*"], check=True)

        # 创建9p共享路径
        gzip_cpio_map_dir(self.map_dir, self.cpio_path)
        subprocess.run(["/bin/cp", "-r", f"{self.map_dir}/cpio/data", self.map_dir], check=True)
        subprocess.run(["/bin/cp", "-r", f"{self.map_dir}/cpio/mockdata", self.map_dir], check=True)
        subprocess.run(["/bin/cp", "-r", f"{self.map_dir}/cpio/opt/bmc/pram", self.map_dir], check=True)

    def set_docker_mount(self):
        # 将启动所需的文件放置到一个临时目录中，便于启动Docker
        # 在temp目录下创建一个docker_temp目录, 其下创建多个以数字开头的文件夹
        path = f"{QEMU_TEMP}/docker_temp"
        if os.path.exists(path):
            subprocess.run(["/bin/rm", "-rf", path], check=True)
        os.makedirs(path)
        for i in range(self.nums):
            cur_path = os.path.join(path, f'qemu_container{i}')
            os.makedirs(cur_path)
            output_data_path = os.path.join(cur_path, 'output', 'data')
            os.makedirs(output_data_path)
            # 拷贝所需文件到该目录下
            subprocess.run(["/bin/cp", "-r", f"{self.map_dir}/mockdata", output_data_path], check=True)
            subprocess.run(["/bin/cp", "-r", f"{self.map_dir}/data", os.path.join(cur_path, "data")], check=True)
            subprocess.run(["/bin/cp", "-r", f"{self.map_dir}/pram", os.path.join(cur_path, "pram")], check=True)
            subprocess.run(["/bin/cp", self.qemu, os.path.join(cur_path, "qemu-system-aarch64")], check=True)
            subprocess.run(["/bin/chmod", "+x", os.path.join(cur_path, "qemu-system-aarch64")], check=True)
            subprocess.run(["/bin/cp", self.zImage_path, os.path.join(cur_path, "zImage")], check=True)
            subprocess.run(["/bin/cp", self.dtb, os.path.join(cur_path, "hi1711_9p8c.dtb")], check=True)
            subprocess.run(["/bin/cp", self.cpio_path, os.path.join(cur_path, "initrd.cpio.gz")], check=True)
            
            # 制作启动脚本
            qemu_run = 'cd /root/workspace/qemu_bin && ./qemu-system-aarch64 ' \
                        '-machine hi1711 ' \
                        '-nographic ' \
                        '-smp 4 ' \
                        '-m 8G ' \
                        '-kernel ./zImage ' \
                        '-dtb ./hi1711_9p8c.dtb ' \
                        '-append "root=/dev/ram0 ip=dhcp console=ttyS0 ' \
                        'rdinit=/init memmap=1886M@0x87200000 ramdisk_size=30720 earlycon=serial-mm,0x08710000 '\
                        'printk.time=y ignore_loglevel ekbox=0x1f000$0x84be0000 kbox_mem=2432k@0x84900000" ' \
                        '-initrd ./initrd.cpio.gz ' \
                        '-netdev user,id=eth0,hostfwd=tcp::2202-:22,hostfwd=tcp::4430-:443,'\
                        'hostfwd=tcp::2302-:23,hostfwd=udp::6230-:623,hostfwd=udp::1610-:161 ' \
                        '-fsdev local,security_model=passthrough,id=fsdev1,path=./pram ' \
                        '-device virtio-9p-device,id=fs1,fsdev=fsdev1,mount_tag=pramshare ' \
                        '-fsdev local,security_model=passthrough,id=fsdev0,path=./data ' \
                        '-device virtio-9p-device,id=fs0,fsdev=fsdev0,mount_tag=hostshare '\
                        '-monitor telnet::4444,server,nowait'
            flags = os.O_WRONLY | os.O_CREAT
            modes = stat.S_IWUSR
            with os.fdopen(os.open(f'{cur_path}/start.sh', flags, modes), 'w+') as f:
                f.write(qemu_run)
            os.chmod(f"{cur_path}/start.sh", stat.S_IXOTH)
        pass

    def run_docker(self):
        self.setenv()
        self.set_docker_mount()
        if self.operation == "create":
            self.docker_run.check_docker_installed()
            containers = []
            for i in range(self.nums):
                containers.append({"name": f"qemu_container{i}", "image": f"{self.docker_images}"})
            mapport = {
                "ssh_port": self.ssh_mapport,
                "telnet_port": self.telnet_mapport,
                "ipmi_port": self.ipmi_mapport,
                "https_port": self.https_mapport,
                "snmp_port": self.snmp_mapport,
                "console_port": self.qemu_monitor_console
            }
            # 创建一个进程池，拉起多个容器
            processes = []
            workspace = os.getcwd()
            for index, container in enumerate(containers):
                mount = workspace + f"/{QEMU_TEMP}/docker_temp/qemu_container{index}"
                p = multiprocessing.Process(target=self.docker_run.run_docker_container,
                                            args=(container["name"], container["image"], 
                                                  mount, mapport))
                processes.append(p)
                p.start()
                for key, value in mapport.items():
                    mapport[key] = value + 10000

            # 等待所有进程完成
            for p in processes:
                p.join()

            logging.info("所有容器已启动完成。")    
        elif self.operation == "update" or self.operation == "delete":
            self.docker_run.run(self.operation, self.docker_images)

    def run_wsl(self):
        self.setenv()
        # 启动qemu命令
        # 是否是1711版本
        if self.container == "hi1711":
            # 配置环境
            os.chmod(self.qemu, stat.S_IXOTH)
            qemu_run = [
                self.qemu, '-machine', self.container, '-nographic', 
                '-smp', str(self.cpus), '-m', f'{self.memory}G',
                '-kernel', self.zImage_path,
                '-dtb', self.dtb,
                '-append', 'root=/dev/ram0 ip=dhcp console=ttyS0 rdinit=/init memmap=1886M@0x87200000 '
                           'ramdisk_size=30720 earlycon=serial-mm,0x08710000 printk.time=y ignore_loglevel '
                           'ekbox=0x1f000$0x84be0000 kbox_mem=2432k@0x84900000',
                '-initrd', self.cpio_path,
                '-netdev', f'user,id=eth0,hostfwd=tcp::{self.ssh_mapport}-:22,'\
                f'hostfwd=tcp::{self.https_mapport}-:443,hostfwd=tcp::{self.telnet_mapport}-:23,'
                f'hostfwd=udp::{self.ipmi_mapport}-:623,hostfwd=udp::{self.snmp_mapport}-:161',
                '-fsdev', f'local,security_model=passthrough,id=fsdev1,path=output/data/pram',
                '-device', 'virtio-9p-device,id=fs1,fsdev=fsdev1,mount_tag=pramshare',
                '-fsdev', f'local,security_model=passthrough,id=fsdev0,path=output/data/data',
                '-device', 'virtio-9p-device,id=fs0,fsdev=fsdev0,mount_tag=hostshare',
                '-fsdev', f'local,security_model=passthrough,id=fsdev2,path=output/data/mockdata',
                '-device', 'virtio-9p-device,id=fs2,fsdev=fsdev2,mount_tag=mockshare',
                '-monitor', f'telnet::{self.qemu_monitor_console},server,nowait'
            ]
            subprocess.run(qemu_run, check=True)

    def run(self):
        # 启动vemake
        self.check_dependencies()
        if self.wsl:
            logging.info("使用本地启动模式")
            self.run_wsl()
        elif self.docker:
            logging.info("使用docker启动模式")
            self.run_docker()
        else:
            logging.error('未定义的使用场景')
            raise SystemExit('未定义的使用场景')


# 提供脚本功能
if __name__ == "__main__":
    if not os.path.exists(QEMU_TEMP):
        os.makedirs(QEMU_TEMP)

    logging.basicConfig(
        level=logging.INFO,
        filename=f'{QEMU_TEMP}/qemu_start.log',
        filemode='w', 
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    # 示例日志记录
    logging.info("QEMU 启动日志开始记录")

    qemu_shells = "build/works/packet/qemu_shells"
    cfg = load_config(f"{qemu_shells}/config.json", f"{qemu_shells}/schema.json")
    print_config(cfg)
    runvemake = Runvemake(cfg)
    runvemake.run()