#!/usr/bin/env python3
# encoding=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 argparse
import json
from collections import defaultdict

import yaml

from bmcgo import misc
from bmcgo.errors import BmcGoException
from bmcgo.misc import CommandInfo
from bmcgo.utils.tools import Tools
from bmcgo.bmcgo_config import BmcgoConfig

tool = Tools("lock")
log = tool.log
command_info: CommandInfo = CommandInfo(
    group=misc.GRP_MISC,
    name="lock",
    description=["自动维护产品和平台lockfile"],
    hidden=False
)


def if_available(bconfig: BmcgoConfig):
    return bconfig.ibmc_sdk is not None or bconfig.manifest is not None


class BmcgoCommand:
    def __init__(self, bconfig: BmcgoConfig, *args):
        self.bconfig = bconfig
        parser = argparse.ArgumentParser(
            prog="bingo lock",
            description="自动维护产品和平台lockfile",
            add_help=True,
            formatter_class=argparse.RawTextHelpFormatter
        )
        group = parser.add_mutually_exclusive_group(required=True)
        # -c / --create：不带参数
        group.add_argument(
            '-c', '--create',
            action='store_true',
            help='创建新lockfile（无需指定名称），将全新生成lockfile'
        )
        # -u / --update：必须带一个组件名参数
        group.add_argument(
            '-u', '--update',
            help='更新lockfile指定的组件（需提供组件名）'
        )

        parser.add_argument(
            '-r', '--remote',
            help='远程中心仓'
        )

        parsed_args, _ = parser.parse_known_args(*args)
        self.create_flag = parsed_args.create
        self.update_component = parsed_args.update
        self.remote = parsed_args.remote
        self.component_revisions = {}

        if self.bconfig.manifest:
            self.build_path = os.path.join(self.bconfig.manifest.folder, "build")
            community_name = misc.community_name()
            self.lock_name = f"{community_name}.lock"
            self.openubmc_sdk_path = os.path.join(self.build_path, "sdk")
            self.openubmc_sdk_lock_name = "openubmc_sdk.lock"
            self.openubmc_sdk_lockfile_path = os.path.join(self.openubmc_sdk_path, self.openubmc_sdk_lock_name)
            self.oct_sdk_flag = False
            if os.path.exists(self.openubmc_sdk_path):
                self.oct_sdk_flag = True
        if self.bconfig.ibmc_sdk:
            self.lock_name = "ibmc_sdk.lock"
            self.build_path = self.bconfig.ibmc_sdk.folder
        self.subsys_dir = os.path.join(self.build_path, "subsys")
        self.subsys_stable_dir = os.path.join(self.subsys_dir, "stable")
        if os.path.exists(self.subsys_stable_dir):
            self.subsys_dir = self.subsys_stable_dir
        if not os.path.isdir(self.subsys_dir):
            raise BmcGoException(f"没有找到subsys目录：{self.subsys_dir}")
        self.lockfile_path = os.path.join(self.build_path, self.lock_name)
        self.temp_path = os.path.join(self.build_path, ".temp")
        self.temp_lockfile_path = os.path.join(self.temp_path, self.lock_name)
        if os.path.isdir(self.temp_path):
            tool.run_command(f"rm -rf {self.temp_path}")
        os.makedirs(self.temp_path)

    @staticmethod
    def conan_lock_add(requires, input_lock=None, output_lock=None):
        lock_cmd = f"conan lock add --requires='{requires}'"
        if input_lock:
            lock_cmd += f" --lockfile {input_lock}"
        if output_lock:
            lock_cmd += f" --lockfile-out {output_lock}"
        tool.run_command(lock_cmd, show_log=True)

    @staticmethod
    def conan_lock_remove(requires, build_require=False, input_lock=None, out_lock=None):
        if not build_require:
            lock_cmd = f"conan lock remove --requires='{requires}'"
        else:
            lock_cmd = f"conan lock remove --build-requires='{requires}'"
        if input_lock:
            lock_cmd += f" --lockfile {input_lock}"
        if out_lock:
            lock_cmd += f" --lockfile-out {out_lock}"
        tool.run_command(lock_cmd, show_log=True)

    @staticmethod
    def conan_lock_update(requires):
        lock_cmd = f"conan lock update --requires='{requires}'"
        tool.run_command(lock_cmd, show_log=True)

    @staticmethod
    def conan_lock_upgrade(requires, upgrade_requires):
        lock_cmd = f"conan lock upgrade --requires='{requires}' --update-requires={upgrade_requires}"
        tool.run_command(lock_cmd, show_log=True)

    @staticmethod
    def _get_manifest_conan_list(manifest_path):
        require_dict = {}
        with open(manifest_path, 'r') as f:
            manifest = yaml.safe_load(f)
        comps = manifest.get('dependencies')
        if not comps:
            return require_dict
        for comp in comps:
            conan = comp.get(misc.CONAN)
            require_name = conan.split("/")[0]
            require_dict[require_name] = conan
        return require_dict

    @staticmethod
    def _get_lockfile_conan_list(lockfile_path):
        require_list = []
        with open(lockfile_path, 'r') as f:
            lock_data = yaml.safe_load(f)
        comps = lock_data.get('requires')
        if not comps:
            return require_list
        for comp in comps:
            require_version = comp.split("#")[0]
            require_list.append(require_version)
        return require_list

    def run(self):
        cruuent_dir = os.getcwd()
        os.chdir(self.temp_path)
        if self.create_flag:
            self._create_lockfile()
            tool.run_command(f"cp -af {self.temp_lockfile_path} {self.lockfile_path}")
        if self.update_component:
            if cruuent_dir.endswith("/build/sdk"):
                self.lockfile_path = self.openubmc_sdk_lockfile_path
            if not os.path.isfile(self.lockfile_path):
                raise BmcGoException(f"{self.lockfile_path}不存在，请先执行bingo lock -c,创建lockfile")
            self._update_lockfile()
        tool.run_command(f"rm -rf {self.temp_path}")

    def conan_lock_create(self, requires=None, lockfile_out=None):
        lock_cmd = "conan lock create "
        if requires:
            lock_cmd += f"--requires='{requires}' "
        else:
            lock_cmd += ". "
        if self.remote:
            lock_cmd += f" -r {self.remote} "
        if lockfile_out:
            lock_cmd += f"--lockfile-out {lockfile_out} "
        lock_cmd += f"-pr profile.ini -u"
        tool.run_command(lock_cmd, show_log=True)

    def _conan_lock_get_require_by_pkgname(self, pkgname):
        self.conan_lock_create(requires=pkgname, lockfile_out="temp.lock")
        with open("temp.lock", 'r') as f:
            lock_data = json.load(f)
        if lock_data is None:
            return None
        for lock_require in lock_data["requires"]:
            if lock_require.startswith(pkgname.split("/")[0]):
                return lock_require
        return None

    def _generate_conanfile(self):
        os.chdir(self.temp_path)
        lock_py_dir = os.path.dirname(os.path.realpath(__file__))
        conan_file_path = os.path.join(lock_py_dir, "..", "tasks", "conan", "conanfile.py")
        if not os.path.isfile(conan_file_path):
            raise BmcGoException(f"没有找到conanfile:{conan_file_path}")
        tool.run_command(f"cp -af {conan_file_path} .")
        tool.run_command(f"sed -i 's/openubmc/product_lock/g' conanfile.py")
        log.success("生成conanfile.py成功！")

    def _generate_requires(self, comps=None):
        # 将组件写入manifest.yml
        new_fd = os.fdopen(os.open("manifest.yml", os.O_WRONLY | os.O_CREAT,
                                   stat.S_IWUSR | stat.S_IRUSR), 'w')
        new_fd.write("base:\n")
        new_fd.write(f"  version: 0.0.1@openubmc/stable\n")
        new_fd.write("dependencies:\n")
        if comps is None:
            comps = self._get_subsys_requires()
        log.debug("依赖列表：{}".format(comps))
        for comp in comps:
            new_fd.write("  - conan: \"{}\"\n".format(comp))
        log.success("生成manifest.yml成功！")

    def _get_subsys_requires(self, pkgname=None):
        comps = []
        # 遍历目标目录下的所有yml文件,获取组件全集
        for f in os.listdir(self.subsys_dir):
            if not f.endswith(".yml"):
                continue
            if f == "qemu.yml":
                continue
            with open(os.path.join(self.subsys_dir, f)) as fp:
                yml = yaml.safe_load(fp)
            deps = yml.get('dependencies')
            for dep in deps:
                conan = dep.get(misc.CONAN)
                if conan:
                    if pkgname and conan.startswith(pkgname):
                        return [conan]
                    comps.append(conan)
        return comps

    def _merge_product_requires(self, product_dir):

        # 遍历目录下的所有"manifest.yml"文件
        manifest_files = []
        for root, _, files in os.walk(product_dir):
            for file in files:
                if file == "manifest.yml":
                    manifest_files.append(os.path.join(root, file))
        # 查找manifest.yml 中带版本号的dependencies
        for manifest_file in manifest_files:
            self._add_product_requires(manifest_file)

    def _add_product_requires(self, manifest_file):
        with open(manifest_file, 'r') as f:
            manifest = yaml.safe_load(f)
        comps = manifest.get('dependencies')
        if not comps:
            return
        for comp in comps:
            conan = comp.get(misc.CONAN)
            if "/" in conan:
                log.success(f"{manifest_file} 存在自定义组件：{conan}")
                length = len(conan.split("@"))
                if length == 1:
                    pkgname = conan + "@openubmc/stable"
                if length == 2:
                    pkgname = conan
                lock_requires = self._conan_lock_get_require_by_pkgname(pkgname=pkgname)
                if lock_requires is not None:
                    self.conan_lock_add(lock_requires, input_lock=self.lock_name, output_lock=self.lock_name)

    def _oct_lockfile(self, inputfile, outputfile, require_dict):
        with open(inputfile, 'r') as f:
            lock_data = yaml.safe_load(f)
        lock_requires_list = []
        requires = lock_data["requires"]
        lock_requires_list.append(requires)
        for require in lock_data["requires"]:
            require_name = require.split("/")[0]
            require_version = require.split("#")[0]
            if require_dict.get(require_name) is None:
                self.conan_lock_remove(require_version, input_lock=inputfile, out_lock=outputfile)
        for require in lock_data["build_requires"]:
            require_name = require.split("/")[0]
            require_version = require.split("#")[0]
            if require_dict.get(require_name) is None:
                self.conan_lock_remove(require_version, True, input_lock=inputfile, out_lock=outputfile)

    def _create_lockfile(self):
        self._generate_conanfile()
        self._generate_requires()
        self.conan_lock_create(lockfile_out="subsys.lock")
        subsys_requiles_list = {}
        for item in self._get_subsys_requires():
            require_name = item.split("/")[0]
            subsys_requiles_list[require_name] = item
        self._oct_lockfile(inputfile="subsys.lock", outputfile="subsys.lock", require_dict=subsys_requiles_list)
        tool.run_command(f"cp -af subsys.lock {self.lock_name}")
        if self.bconfig.manifest:
            log.info(f"检测到在manifest仓，lockfile需要合并product下的自定义版本的组件")
            product_dir = os.path.join(self.build_path, "product")
            self._merge_product_requires(product_dir)
            if self.oct_sdk_flag:
                # 基于requiles_dict 裁剪出openubmc_sdk.lock
                sdk_manifest_path = os.path.join(self.openubmc_sdk_path, "manifest.yml")
                sdk_comp_list = self._get_manifest_conan_list(sdk_manifest_path)
                tool.run_command(f"cp -af subsys.lock {self.openubmc_sdk_lock_name}")
                self._oct_lockfile(inputfile=self.openubmc_sdk_lock_name, outputfile=self.openubmc_sdk_lock_name, require_dict=sdk_comp_list)
                tool.run_command(f"cp -af {self.openubmc_sdk_lock_name} {self.openubmc_sdk_path}")
                # openubmc.lock 删除 openubmc_sdk 组件
                require_list = self._get_lockfile_conan_list(self.openubmc_sdk_lock_name)
                for require in require_list:
                    self.conan_lock_remove(require, input_lock=self.lock_name, out_lock=self.lock_name)
                    self.conan_lock_remove(require, True, input_lock=self.lock_name, out_lock=self.lock_name)
        log.success("创建lockfile成功！")

    def _update_lockfile(self):
        dict_comps = self._load_lockfile()
        if self.update_component not in dict_comps:
            log.info(f"当前要更新的组件在lockfile中不存在,是否添加，确认输入Y")
            if input().upper() == "Y":
                subsys_pkg = self._get_subsys_requires(pkgname=self.update_component)[0]
                lock_requires = self._conan_lock_get_require_by_pkgname(pkgname=subsys_pkg)
                if lock_requires is not None:
                    self.conan_lock_add(lock_requires, self.lockfile_path, self.lockfile_path)
        else:
            log.info(f"当前要更新的组件在lockfile中存在版本：")
            n = 1
            dict_section = {}
            for version in dict_comps[self.update_component]:
                log.success(str(n) + "  " + version)
                dict_section[n] = version
                n += 1
            log.info("请选择你要维护的编号：")
            number_section = int(input())
            update_component = dict_section.get(number_section)
            if update_component not in dict_comps[self.update_component]:
                raise BmcGoException(f"选择版本错误！！")
            log.success("请选择动作\nA:更新版本  B:版本不变，仅更新revision")
            action_section = input().upper()
            if action_section == "A":
                subsys_pkg = self._get_subsys_requires(pkgname=self.update_component)[0]
                lock_requires = self._conan_lock_get_require_by_pkgname(pkgname=subsys_pkg)
            elif action_section == "B":
                if not self.remote:
                    raise BmcGoException(f"更新组件revision, 请指定remote, 示例: bingo lock -u vpd -r openubmc_opensource")
                subsys_pkg = update_component.split("#")[0]
                self.component_revisions = self._get_revisons(subsys_pkg)
                self._download_revisions(subsys_pkg)
                if len(self.component_revisions) > 20:
                    log.warning("警告:当前要更新的组件revisions已超过20个，请联系maintainer维护版本号！！！")
                log.info("当前要更新的组件存在如下revisions：")
                for k, v in self.component_revisions.items():
                    log.success("revision " + k)
                    log.info(v["message"])
                log.info("请输入你要更新的revision：")
                revision = input()
                lock_requires = subsys_pkg + "#" + revision + "%" + str(self.component_revisions[revision]["timestamp"])
            if lock_requires is not None:
                self.conan_lock_remove(update_component, input_lock=self.lockfile_path, out_lock=self.lockfile_path)
                self.conan_lock_add(lock_requires, self.lockfile_path, self.lockfile_path)
        log.success("更新lockfile成功！")

    def _get_revisons(self, com_reference):
        ret = tool.run_command(f"conan list {com_reference}#* -r {self.remote} -f=json", capture_output=True)
        data = json.loads(ret.stdout)
        return data.get(self.remote).get(com_reference).get("revisions")

    def _download_revisions(self, subsys_pkg):
        for rev in self.component_revisions.keys():
            tool.run_command(f"conan download {subsys_pkg}#{rev} -r {self.remote}")
            ret = tool.run_command(f"conan cache path {subsys_pkg}#{rev}", capture_output=True)
            if ret.returncode == 0:
                conandata_file = os.path.join(ret.stdout.strip(), "conandata.yml")
                with open(conandata_file, 'r') as f:
                    conf = yaml.safe_load(f)
                version = subsys_pkg.split("@")[0].split("/")[1]
                if conf.get("sources").get(version).get("message"):
                    self.component_revisions[rev]["message"] = conf.get("sources").get(version).get("message")
                else:
                    self.component_revisions[rev]["message"] = ""

    def _load_lockfile(self):
        # 打开并加载 JSON 文件
        with open(self.lockfile_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
        # 使用 defaultdict(list) 自动初始化列表
        component_versions = defaultdict(list)

        # 遍历所有依赖类型（requires, build_requires 等）
        for dep_type in ['requires']:
            if dep_type in data and isinstance(data[dep_type], list):
                for full_ref in data[dep_type]:
                    # 格式: name/version@channel#hash%timestamp
                    # 我们只需要 name 和 version
                    try:
                        # 提取 name/version 部分（在 @ 之前）
                        name_version = full_ref.split('@')[0]
                        name, version = name_version.split('/', 1)
                        component_versions[name].append(full_ref.split('%')[0])
                    except ValueError:
                        log.error(f"警告：无法解析依赖项 '{full_ref}'，跳过。")

        return dict(component_versions)  # 转为普通 dict 方便使用
