#! /usr/bin/env python3
# -*- 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.
'''
功    能：构建框架
修改记录：2021-10-11 创建
'''
import json
import importlib
import os
import time
import sys
import argparse
import threading
import signal
import shutil
import socket
import stat

from socket import socket, AF_INET, SOCK_STREAM, SOL_SOCKET, SO_REUSEADDR
from multiprocessing import Process, Queue
import yaml
from bmcgo import misc
from bmcgo import errors
from bmcgo.utils.perf_analysis import PerfAnalysis
from bmcgo.utils.config import Config
from bmcgo.logger import Logger
from bmcgo.utils.tools import Tools
from bmcgo.bmcgo_config import BmcgoConfig
from bmcgo.functional.deploy import BmcgoCommand as Deploy

bmcgo_install_path = os.path.split(os.path.realpath(__file__))[0]
args = None
conan_args = None
unknown_args = None
config: Config = None
# 注意端口号需要与V3的配置不同
WORK_SERVER_PORT = 62345
# 任务失败状态
TASK_STATUS_FAILED = "Failed"
# 用于通知主程序 server已经启动
q = Queue(1)
tool = Tools()
log = tool.log


class WsServerCommException(errors.BmcGoException):
    """
        与状态服务通信失败
    """

    def __init__(self, *arg, **kwarg):
        super(WsServerCommException, self).__init__(*arg, **kwarg)


class WorkStatusServer(Process):
    """
    全局任务状态管理器
    """
    def __init__(self):
        super().__init__()
        self.task_status_dict = {}
        self.start_time = time.time()
        self.fd = socket(AF_INET, SOCK_STREAM)
        self.fd.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        address = ("", WORK_SERVER_PORT)
        try:
            self.fd.bind(address)
            self.fd.listen(10)
            # 向主进程发送准备好的消息
            q.put(1)
        except Exception as e:
            log.error(f"任务状态错误: {e}")
            raise Exception(f"请手动运行此命令(killall 未安装请执行 apt install psmisc): killall python") from e
        self.has_work_failed = False

    def run(self):
        while True:
            self._accept()

    def _accept(self):
        cur_time = time.time()
        if int(cur_time - self.start_time)/60 > 60:
            self.start_time = time.time()
            log.warning(f"任务状态字典为: {self.task_status_dict}")
            log.warning("构建时长已超 60 分钟")

        cli_fd, _ = self.fd.accept()
        recv = cli_fd.recv(1024)
        recv_json = recv.decode()
        msg = json.loads(recv_json)
        act = msg["action"]
        resp = ""
        wn = msg["work_name"]
        if act == "set_if_not_exist":
            st = msg["work_status"]
            resp = "False"
            status = self.task_status_dict.get(wn)
            if status is None:
                self.task_status_dict[wn] = st
                resp = "True"
        elif act == "set":
            st = msg["work_status"]
            self.task_status_dict[wn] = st
            if st == TASK_STATUS_FAILED:
                self.has_work_failed = True
            resp = "True"
        elif act == "get":
            if self.has_work_failed:
                resp = TASK_STATUS_FAILED
            else:
                st = self.task_status_dict.get(wn)
                resp = "None" if st is None else st

        cli_fd.send(resp.encode())
        cli_fd.close()


class WorkStatusClient(Process):
    """
    全局任务状态管理器
    """
    def __init__(self):
        super().__init__()
        self.task_status_dict = {}

    # 获取任务状态
    def get(self, target_name, work_name, ignore_error=False):
        return self.__comm("{\"action\":\"get\", \"work_name\":\"%s/%s\"}" % (target_name, work_name), ignore_error)

    # 任务不存在时创建任务状态，否则返回False
    def set_if_not_exist(self, target_name, work_name, status, ignore_error=False):
        return self.__comm("{\"action\":\"set_if_not_exist\", \"work_name\":\"%s/%s\", \"work_status\":\"%s\"}" %
                           (target_name, work_name, status), ignore_error)

    # 设置任务状态
    def set(self, target_name, work_name, status, ignore_error=False):
        self.__comm("{\"action\":\"set\", \"work_name\":\"%s/%s\", \"work_status\":\"%s\"}" %
                    (target_name, work_name, status), ignore_error)

    def __comm(self, msg, ignore_error):
        fd = socket(AF_INET, SOCK_STREAM)
        fd.settimeout(2)
        while True:
            try:
                fd.connect(("", WORK_SERVER_PORT))
                fd.send(msg.encode())
                msg = fd.recv(1024)
                fd.close()
                return msg.decode()
            except Exception as e:
                if ignore_error:
                    log.info(str(e))
                    return None
                else:
                    raise WsServerCommException("与任务状态服务器通信失败") from e

ws_client = WorkStatusClient()


def wait_finish(target_name, wait_list, prev_work_name):
    if not wait_list:
        return True
    start_time = time.time()
    cnt = 0
    while True:
        finish = True
        time.sleep(0.1)
        for work_name in wait_list:
            cur_time = time.time()
            status = ws_client.get(target_name, work_name, ignore_error=True)
            if status is None:
                log.debug(f"任务{target_name}/{prev_work_name}执行失败，原因是未能获取到任务{target_name}/{work_name}的状态")
                ws_client.set(target_name, prev_work_name, TASK_STATUS_FAILED, ignore_error=True)
                return False
            if status == "Done":
                continue
            if status == TASK_STATUS_FAILED:
                log.debug(f"任务{target_name}/{prev_work_name}执行失败，原因其等待的任务{target_name}/{work_name}失败")
                ws_client.set(target_name, prev_work_name, TASK_STATUS_FAILED, ignore_error=True)
                return False
            finish = False
            # 每等待60s打印一次日志
            if int(cur_time - start_time) >= 60:
                start_time = time.time()
                cnt += 60
                log.info("目标 {} 正在等待任务: {}, 当前已等待 {} 秒".format(target_name, work_name, cnt))
            break
        if finish:
            return True


class WorkerScheduler(Process):
    '''
    '''
    def __init__(self, target_name, work, perf: PerfAnalysis):
        super().__init__()
        self.work = work
        self.target_name = target_name
        self.work_name = self.work["name"]
        self.klass = self.work.get("klass", "")
        self.perf = perf

    def load_class(self):
        if self.klass == "":
            return None
        # bmcgo的任务类名固定为TaskClass
        if self.klass.startswith("bmcgo"):
            work_path = self.klass
            class_name = "TaskClass"
        else:
            split = self.klass.split(".", -1)
            work_path = ".".join(split[:-1])
            class_name = split[-1]
        log.debug("工作路径: {}, 类名: {}".format(work_path, class_name))
        try:
            work_py_file = importlib.import_module(work_path)
            return getattr(work_py_file, class_name)
        except ModuleNotFoundError as e:
            ignore = self.work.get("ignore_not_exist", False)
            if ignore:
                log.warning(f"{self.klass} 已配置 ignore_not_exist 且为真, 跳过执行")
                return None
            else:
                raise e

    def run(self):
        ret = -1
        try:
            ret = self._run()
            if ret != 0:
                ws_client.set(self.target_name, self.work_name, TASK_STATUS_FAILED, ignore_error=True)
        except WsServerCommException as exc:
            msg = str(exc)
            log.debug(msg)
            ret = -1
        except Exception as exc:
            if os.environ.get("LOG"):
                import traceback
                log.info(traceback.format_exc())
            msg = str(exc)
            log.error(msg)
            ws_client.set(self.target_name, self.work_name, TASK_STATUS_FAILED, ignore_error=True)
            ret = -1
        if ret != 0:
            log.debug(f"任务名: {self.work_name}, 类名: {self.klass} 退出状态错误")
        return ret

    def _run(self):
        '''
        功能描述：执行work
        '''
        global config
        work_name = self.work_name
        log.debug(f"任务名: {self.work_name}, 类: {self.klass}) 已就绪")
        ret = wait_finish(self.target_name, self.work.get("wait"), work_name)
        if not ret:
            log.debug(f"等待任务 {self.work_name} 类 {self.klass} 发生错误")
            return -1
        target_config = self.work.get("target_config")
        config.deal_conf(target_config)
        # bmcgo的任务类名固定为TaskClass
        work_class = self.load_class()
        # 如果未指定类时，不需要执行
        if work_class is not None:
            work_x = work_class(config, work_name)
            # work配置项和target配置项
            work_config = self.work.get("work_config")
            work_x.deal_conf(work_config)
            ret = ws_client.set_if_not_exist(self.target_name, work_name, "Running")
            self.perf.add_data(work_name, "running")
            if not args.debug_frame and ret:
                # 创建进程并且等待完成或超时
                ret = work_x.run()
                if ret is not None and ret != 0:
                    return -1
            elif not args.debug_frame and not ret:
                # 不需要创建进程，等待任务执行完成即可
                wait_list = []
                wait_list.append(work_name)
                ret = wait_finish(self.target_name, wait_list, work_name)
                if not ret:
                    log.debug(f"等待任务 {self.work_name} 类 {self.klass} 发生错误")
                    return -1

            log.debug(f"任务 {work_name} 开始安装步骤")
            if not args.debug_frame:
                ret = work_x.install()
                if ret is not None and ret != 0:
                    return -1
            self.perf.add_data(work_name, "finish")

        # 创建子任务
        ret = exec_works(self.target_name, self.work.get("subworks"), work_name, self.perf)
        if not ret:
            ws_client.set(self.target_name, self.work_name, TASK_STATUS_FAILED, ignore_error=True)
            log.error(f"运行子任务 {self.work_name} 类 {self.klass}失败")
            return -1
        # 创建include_target子任务
        target_include = self.work.get("target_include")
        if target_include:
            ret = create_target_scheduler(work_name, target_include, self.perf)
            if not ret:
                ws_client.set(self.target_name, self.work_name, TASK_STATUS_FAILED, ignore_error=True)
                log.error(f"创建计划表 {target_include} 失败")
                return -1
        log.success(f"任务 {work_name} 完成")
        ws_client.set(self.target_name, self.work_name, "Done")
        return 0


class Worker():
    """
    任务执行器
    """
    work_name = ""
    target_name = ""

    def exec_work(self, target_name, work, perf: PerfAnalysis):
        try:
            return self._exec_work(target_name, work, perf)
        except Exception as exc:
            msg = str(exc)
            log.debug(msg)
            return -1

    def run(self, target_name, work, perf: PerfAnalysis):
        self.work_name = work["name"]
        self.target_name = target_name
        t = threading.Thread(target=self.exec_work, args=(target_name, work, perf))
        t.start()

    def _exec_work(self, target_name, work, perf: PerfAnalysis):
        ws = WorkerScheduler(target_name, work, perf)
        ws.start()
        try_cnt = 0
        while ws.is_alive():
            try_cnt += 1
            time.sleep(0.1)

            # 每个任务的超时时间为30分钟(18000次循环)，如果超时则失败
            if try_cnt > 42000:
                log.error(f"任务{self.target_name}/{self.work_name}执行超时(>30min)，强制退出")
                ws.kill()
                return -1
        if ws.exitcode is not None and ws.exitcode == 0:
            return 0
        return -1


def exec_works(target_name, work_list, prev_work_name, perf: PerfAnalysis):
    if not work_list:
        return True
    # 创建任务并等待完成
    wait_list = []
    for work in work_list:
        worker = Worker()
        worker.run(target_name, work, perf)
        wait_list.append(work["name"])
    return wait_finish(target_name, wait_list, prev_work_name)


def read_config(file_path: str):
    '''
    功能描述：读取json内容
    '''
    if not os.path.exists(file_path):
        raise errors.NotFoundException(f"{file_path} 路径不存在")

    # conan专用处理
    with open(file_path, "r") as fp:
        data = fp.read()
        if conan_args:
            if conan_args.upload_package:
                data = data.replace('${upload}', "true")
            else:
                data = data.replace('${upload}', "false")
            data = data.replace('${conan_package}', conan_args.conan_package)
        return yaml.safe_load(data)


def create_target_scheduler(target_alias, target_name, perf: PerfAnalysis):
    log.info(f"创建新目标 {target_name} 构建计划表")
    global config
    if args.debug_task is None:
        manifest_target = f"{config.code_path}/target/{target_name}.yml"
        bmcgo_target = os.path.join(bmcgo_install_path, "target", f"{target_name}.yml")
        if os.path.isfile(manifest_target):
            target_file = manifest_target
        elif os.path.isfile(bmcgo_target):
            target_file = bmcgo_target
        else:
            raise Exception(f"构建目标文件 [target_]{target_name}.yml 不存在")
        # 创建配置

        work_list = read_config(target_file)
    else:
        work_list = [
            {
                "name": "test_task",
                "klass": args.debug_task
            }
        ]
    if isinstance(work_list, dict):
        target_cfg = work_list.get("target_config", {})
        config.deal_conf(target_cfg)
        environments = work_list.get("environment", {})
        for key, value in environments.items():
            log.success(f"配置环境变量 {key}: {value}")
            os.environ[key] = value
    # 打印配置
    config.print_config()
    # 打印任务清单
    log.debug(f"任务列表:{work_list}")
    # 创建任务调度器
    if isinstance(work_list, dict):
        subworks = work_list.get("subworks")
        if not subworks or not isinstance(subworks, list):
            raise errors.BmcGoException(f"target文件{target_file}缺少subworks配置")
        return exec_works(target_alias, subworks, "TOP", perf)
    else:
        return exec_works(target_alias, work_list, "TOP", perf)


class Frame(object):
    def __init__(self, bconfig: BmcgoConfig):
        self.bconfig = bconfig
        https_proxy = os.environ.get("https_proxy")
        if https_proxy is not None and https_proxy != "":
            log.warning("检测到环境设置了https_proxy代理，可能导致网络资源访问失败.")
            log.warning("  如需关闭代理，请执行命令: export https_proxy=")
            log.warning("  如需关闭特定域名或IP的代理，可以设置no_proxy环境变量，也可将no_proxy设置到/etc/profile中.")
            log.warning("  no_proxy配置示例: export no_proxy=localhost,127.0.0.0/8,.huawei.com,.inhuawei.com")
        self.perf = PerfAnalysis(os.path.join(bconfig.manifest.folder, "output"))
        self.code_path = os.path.join(bconfig.manifest.folder, "build")
        sys.path.append(self.code_path)
        os.makedirs(misc.CACHE_DIR, exist_ok=True)
        Config.log_init()
        self.ws_server = None
        self.config = None
        self.need_deploy = False

    @staticmethod
    def record_build_time():
        date_file = f"{config.temp_path}/date/date.txt"
        os.makedirs(os.path.dirname(date_file), exist_ok=True)
        with os.fdopen(os.open(date_file, os.O_WRONLY | os.O_CREAT | os.O_TRUNC,
                               stat.S_IWUSR | stat.S_IRUSR), 'w') as file_handler:
            now = time.strftime("%H:%M:%S %b %-d %Y", time.localtime())
            file_handler.write(now)
            file_handler.close()

    def get_all_target(self):
        targets = {}
        dirname = os.path.join(bmcgo_install_path, "target")
        for file in os.listdir(dirname):
            if not file.endswith(".yml"):
                continue
            tgt_file = os.path.join(bmcgo_install_path, file)
            targets[file] = tgt_file
        dirname = os.path.join(self.code_path, "target")
        if not os.path.isdir(dirname):
            return targets
        for file in os.listdir(dirname):
            if not file.endswith(".yml"):
                continue
            tgt_file = os.path.join(dirname, file)
            targets[file] = tgt_file
        return targets

    def cmc_open_parse(self):
        global args, unknown_args, conan_args, config
        if args.target != "cmc_open" and args.target != "docker_build":
            return None
        parser = argparse.ArgumentParser(description="cmc_open info")
        parser.add_argument("-ud", "--upload_docker", help="是否上传docker镜像", action=misc.STORE_TRUE)
        parser.add_argument("-pn", "--partner_name", help="伙伴名称")
        local_time = time.strftime("%y%m%d%H%M", time.localtime())
        default_tag = f'v3_partner_{local_time}_{self.bconfig.current_branch}'
        parser.add_argument("-dtag", "--docker_tag", help="上传docker镜像tag", default=default_tag)
        docker_args, unknown_args = parser.parse_known_args(unknown_args)
        if docker_args.upload_docker and docker_args.partner_name is None:
            raise errors.BmcGoException("伙伴名称为空, 无法完成上传, 请指定伙伴名称!")
        return docker_args

    def kunpeng_publish_parse(self):
        global args, unknown_args, conan_args, config
        if args.target != "cmc_open" and args.target != "kunpeng_publish":
            return None
        os.environ['NOT_CLEAN_CONAN_CACHE'] = "True"
        parser = argparse.ArgumentParser(description="open_source conan repo info")
        parser.add_argument("-opr", "--open_conan_remote", help="开源/伙伴仓远端(场内远端)")
        kunpeng_publish_args, unknown_args = parser.parse_known_args(unknown_args)
        return kunpeng_publish_args

    def parse(self, argv=None):
        global args, unknown_args, conan_args, config
        parser = argparse.ArgumentParser(description="构建openUBMC",
            parents=[Config.argparser(self.code_path, self.bconfig.partner_mode)],
            add_help=False,
            formatter_class=argparse.RawTextHelpFormatter,
        )
        help_info = "构建目标，请查看build/target目录, 支持的目标："
        targets = self.get_all_target()
        for tgt, _ in targets.items():
            help_info += "\n" + tgt[:-4]

        parser.add_argument("-d", "--debug_frame", help=argparse.SUPPRESS, const="debug", action="store_const")
        if log.is_debug:
            parser.add_argument("--debug_task", help="调试任务，与目标描述文件的klass完全相同，如：bmcgo.tasks.task_download_buildtools",
                                default=None)
        else:
            parser.add_argument("--debug_task", help=argparse.SUPPRESS, default=None)

        parser.add_argument("-v", "--version", help="构建版本号，不指定时从manifest.yml读取", default="")
        parser.add_argument("-t", "--target", help=help_info, default="personal")
        parser.add_argument("--deploy", help="将hpm包部署至bmc设备", action=misc.STORE_TRUE)
        args, unknown_args = parser.parse_known_args(argv)
        if args.target.startswith("target_"):
            args.target = args.target[7:]
        log.info(f'已知参数: {argv}')
        log.info(f"调试框架: {args.debug_frame}, 构建参数: {args}")
        conan_index_options = []
        if args.target == "dependencies":
            parser = argparse.ArgumentParser(description="build target")
            parser.add_argument("-cp", "--conan_package", help="软件包名, 示例: kmc/21.1.2.B001", default="")
            parser.add_argument("-uci", "--upload_package", help="是否上传软件包", action=misc.STORE_TRUE)
            parser.add_argument("-o", "--options", help="组件特性配置, 示例: -o skynet:enable_luajit=True",
                                action='append')
            conan_args, unknown_args = parser.parse_known_args(unknown_args)
            if conan_args is None or conan_args.conan_package == "":
                raise errors.BmcGoException("软件包选项为空, 请输入软件包选项!")
            conan_index_options = conan_args.options
        docker_args = self.cmc_open_parse()
        kunpeng_publish_args = self.kunpeng_publish_parse()

        config = Config(self.bconfig, target=args.target)
        config.parse_args(argv)
        self.record_build_time()
        config.set_version(args.version)
        config.set_conan_index_options(conan_index_options)
        if docker_args:
            config.set_docker_info(docker_args)
        if kunpeng_publish_args:
            config.set_kunpeng_publish_info(kunpeng_publish_args)
        self.need_deploy = args.deploy
        self.config = config
        # 开启schema校验，也可参考personal目标配置禁用校验
        self.config.set_schema_need_validate(True)

    def sigint_handler(self, _, __):
        log.debug('关闭全局状态管理器，所有任务都将因与状态管理器通信失败而退出!')
        self.ws_server.kill()
        signal.raise_signal(signal.SIGKILL)

    def sigterm_handler(self, _, __):
        log.debug('接收到终止信号!, 退出!')
        signal.raise_signal(signal.SIGINT)

    def run(self):
        tool.sudo_passwd_check()
        shutil.rmtree(misc.CACHE_DIR)
        os.makedirs(misc.CACHE_DIR)
        os.chmod(misc.CACHE_DIR, 0o777)
        if shutil.which(misc.CONAN) is not None:
            tool.run_command("conan remove --locks")
        # 启动全局状态服务
        signal.signal(signal.SIGINT, self.sigint_handler)
        signal.signal(signal.SIGTERM, self.sigterm_handler)
        self.ws_server = WorkStatusServer()
        self.ws_server.start()
        succ = True
        # 初始化性能数据收集功能
        try:
            # 等待sever起来
            q.get()
            succ = create_target_scheduler(args.target, args.target, self.perf)
        except Exception as e:
            log.error(str(e))
            log.error(f"Error: 构建目标({args.target})失败，请参考“构建检查”和下载“全量构建日志（右上角-日志-全量构建日志）”检查错误！！！！！！！！")
            succ = False

        self.ws_server.kill()
        if os.environ.get("PERF") is not None:
            self.perf.render(args.target)
        if succ:
            log.success(f"任务 {args.target} 执行成功")
            hpm_path = os.path.join(self.config.output_path, f"rootfs_{self.config.board_name}.hpm")
            if self.need_deploy and os.path.isfile(hpm_path):
                deploy = Deploy(self.bconfig, ["-f", hpm_path])
                deploy.run()
            time.sleep(0.5)
        else:
            # 等待2秒给各子任务同步状态
            time.sleep(2)
            raise errors.BmcGoException(f"任务 {args.target} 执行失败")
        return 0
