# -*- coding: utf-8 -*-
import time
import traceback
from tenacity import stop_after_attempt, wait_fixed
from tenacity import retry as retry_adorn

import utils.common.log as logger
from utils.common.exception import HCCIException
from utils.business.param_util import ParamUtil
from plugins.DistributedStorage.common.UpgradeHotPatchOperate import UpgradeHotPatchOperate
from plugins.DistributedStorage.common.UpgradeOperate import UpgradeOperate
from plugins.DistributedStorage.common.constants import SERVICE_NAME
from plugins.DistributedStorage.common.constants import UpgradeStates
from plugins.DistributedStorage.common.constants import HotpatchStates
from plugins.DistributedStorage.common.constants import RollBackStates
from plugins.DistributedStorage.common.constants import UpgradeStateKey
from plugins.DistributedStorage.common.base import TestCase
from plugins.DistributedStorage.common.PublicHandleNew import RestPublicMethod
from plugins.DistributedStorage.Upgrade.scripts.impl.TC_Upgrade_Confirm import UpgradeConfirm
from plugins.DistributedStorage.Upgrade.scripts.impl.TC_Post_Upgrade_Check import PostUpgradeCheck
from plugins.DistributedStorage.Upgrade.scripts.impl.TC_Rollback_Upgrade_Pkg import UpgradePkg


class RollBackPkg(TestCase):
    def __init__(self, project_id, pod_id, fs_args, condition=None, metadata=None, **kwargs):
        super(RollBackPkg, self).__init__(project_id, pod_id)
        self.project_id = project_id
        self.pod_id = pod_id
        self.region_id = fs_args.get("region_id")
        self.more_args = kwargs
        self.condition = condition
        self.metadata = metadata
        self.check_timeout = 90
        self.node_list = list()
        self.master_node = fs_args.get("master_node")
        self.master_client = None
        self.slaver_node = fs_args.get("slaver_node")
        self.opr = UpgradeOperate(fs_args)
        self.user_name = fs_args["user_name"]
        self.password = fs_args["password"]
        self.remote_path = "/tmp/upgrade_tmp_hotpatch"
        self.fs_args = fs_args
        self.fs_args["upgrade_type"] = "rollback"
        self.rollback_key = self.fs_args["rollback_key"]
        self.rest_opr = RestPublicMethod(project_id, pod_id, fs_args)

    def procedure(self):
        logger.info('Start rollback task.')
        try:
            status_code, error_code, error_des = self.opr.try_login(self.user_name, self.password)
            if status_code != 200 or error_code != 0:
                err_msg = "Failed to login, Detail:[status:%s,code:%s], error:%s" % (status_code, error_code, error_des)
                logger.error(err_msg)
                raise Exception(err_msg)
            if not self.fs_args.get("hot_patch_tag"):
                self.rollback_pkg_procedure()
            else:
                self.get_host_list()
                self.roll_back_hotpatch_and_product()

        except HCCIException as e:
            logger.error('rollback upgrade pkg failed, details:{}'.format(e))
            logger.error(traceback.format_exc())
            raise e
        except Exception as e:
            logger.error('rollback upgrade pkg failed, details:{}'.format(e))
            logger.error(traceback.format_exc())
            raise HCCIException(620011, str(e))
        finally:
            if self.fs_args.get("hot_patch_tag") and self.master_client:
                UpgradeHotPatchOperate.del_script_om_fsm(self.master_client)
                del self.master_client
            logger.info('procedure end...')

    def roll_back_hotpatch_and_product(self):
        """
        回滚产品大包和热补丁包
        回滚前判断当前升级状态：升级大包，升级大包完成（热补丁未升级），升级热补丁，升级热补丁完成
        """
        roll_back_tasks = [
            "rollback_pkg_procedure",
            "roll_back_post_check",
            "restore_product_repo",
            "restore_product_db",
            "roll_back_post_check",
            "rollback_pkg_procedure"
        ]
        self.master_client, master_node = self.get_master_fsm_node()
        self.upload_script_to_master_fsm(master_node)
        product_upgrade_state = self.get_product_upgrade_state()
        hotpatch_upgrade_state = self.get_hotpatch_upgrade_state()
        if not self.fs_args.get("rollback_product_tag"):
            logger.info("Only rollback hotpatch pkg.")
            roll_back_tasks = roll_back_tasks[:-1]
        if product_upgrade_state <= UpgradeStates.PRODUCT_UPGRADE_COMMIT_START:
            # 大包升级确认前
            self.del_ignore_hosts()
            self.run_roll_back_tasks(RollBackStates.ROLLBACK_PRODUCT_START,
                                     roll_back_tasks)
        elif product_upgrade_state <= UpgradeStates.DEL_IGNORE_NODES_START and \
                hotpatch_upgrade_state < HotpatchStates.HOTPATCH_UPGRADE_START:
            self.del_ignore_hosts()
            self.run_roll_back_tasks(RollBackStates.RESTORE_PRODUCT_REPO_DATA_START,
                                     roll_back_tasks)
        else:
            self.run_roll_back_tasks(RollBackStates.ROLLBACK_HOTPATCH_START,
                                     roll_back_tasks)

    def retry(self):
        logger.info('Start retry rollback task.')
        try:
            status_code, error_code, error_des = self.opr.try_login(self.user_name, self.password)
            if status_code != 200 or error_code != 0:
                err_msg = "Failed to login, Details:[status:%s,code:%s]%s" % (status_code, error_code, error_des)
                logger.error(err_msg)
                raise Exception(err_msg)
            if not self.fs_args.get("hot_patch_tag"):
                self.rollback_pkg_procedure()
            else:
                self.get_host_list()
                self.roll_back_hotpatch_and_product_retry()
        except HCCIException as e:
            logger.error('rollback pkg failed, details:{}'.format(e))
            logger.error(traceback.format_exc())
            raise e
        except Exception as e:
            logger.error('rollback pkg failed, details:{}'.format(e))
            logger.error(traceback.format_exc())
            raise HCCIException(620011, str(e))
        finally:
            if self.fs_args.get("hot_patch_tag") and self.master_client:
                UpgradeHotPatchOperate.del_script_om_fsm(self.master_client)
                del self.master_client
            logger.info('procedure end...')

    def run_roll_back_tasks(self, start_task, roll_back_tasks):
        """
        0：补丁回滚完成完成
        1: 补丁回滚确认检查
        2：补丁回滚确认完成
        3：恢复大包数据库完成
        4：恢复大包软件仓完成
        5：回滚大包完成
        """
        for index, task in enumerate(roll_back_tasks[start_task:]):
            self.update_rollback_state(start_task + index)
            getattr(self, task)()

    def roll_back_hotpatch_and_product_retry(self):
        self.master_client, master_node = self.get_master_fsm_node()
        self.upload_script_to_master_fsm(master_node)
        current_upgrade_state = self.get_rollback_state()
        roll_back_tasks = [
            "rollback_pkg_retry",
            "roll_back_post_check",
            "restore_product_repo",
            "restore_product_db",
            "roll_back_post_check",
            "rollback_pkg_procedure"
        ]
        if not self.fs_args.get("rollback_product_tag"):
            logger.info("Only rollback hotpatch pkg.")
            roll_back_tasks = roll_back_tasks[:-1]
        if current_upgrade_state <= RollBackStates.ROLLBACK_HOTPATCH_START:
            self.run_roll_back_tasks(RollBackStates.ROLLBACK_HOTPATCH_START,
                                     roll_back_tasks)
        elif current_upgrade_state <= RollBackStates.ROLLBACK_HOTPATCH_POST_CHECK_START:
            self.run_roll_back_tasks(RollBackStates.ROLLBACK_HOTPATCH_POST_CHECK_START,
                                     roll_back_tasks)
        elif current_upgrade_state < RollBackStates.ROLLBACK_PRODUCT_START:
            self.del_ignore_hosts()
            self.del_upg_pkg()
            self.run_roll_back_tasks(RollBackStates.RESTORE_PRODUCT_REPO_DATA_START,
                                     roll_back_tasks)
        else:
            roll_back_tasks[-1] = "rollback_pkg_retry"
            self.run_roll_back_tasks(RollBackStates.ROLLBACK_PRODUCT_START,
                                     roll_back_tasks)

    def get_master_fsm_node(self):
        return UpgradeHotPatchOperate.get_master_node_client(
            [self.master_node, self.slaver_node])

    def get_post_check_status(self):
        logger.info("check upgrade status.")
        time_out = 600
        while time_out:
            ret_result, ret_data = self.opr.get_upgrade_status()
            if ret_result["code"] != '0':
                err_msg = "get upgrade status failed, " \
                          "Detail:[result:%s, data:%s]" \
                          % (ret_result, ret_data)
                logger.error(err_msg)
                raise Exception(err_msg)
            elif ret_data["currentPhase"] == "postcheck_success":
                break
            else:
                logger.info("current is %s" % ret_data["currentPhase"])
            time.sleep(20)
            time_out -= 20
        if time_out <= 0:
            err_msg = "Query post check status timeout"
            logger.error(err_msg)
            raise Exception(err_msg)

    def rollback_pkg_procedure(self):
        """
        标准升级回滚
        """
        UpgradePkg(self.project_id, self.pod_id, self.fs_args).procedure()

    def rollback_pkg_retry(self):
        """
        标准升级回滚重试
        """
        UpgradePkg(self.project_id, self.pod_id, self.fs_args).retry()

    def del_ignore_hosts(self):
        params = {
            "agentIps": self.node_list
        }
        ret_result, ret_data = self.opr.handle_ignore_host(params, "delete")
        if ret_result.get("code") != "0":
            err_msg = "Failed to del the node to ignore, Details: %s." % ret_data
            logger.error(err_msg)

    def get_host_list(self):
        """
        获取存储节点列表
        """
        osd_nodes_list, vbs_nodes_list = self.rest_opr.get_server_list()
        self.node_list.extend(osd_nodes_list)
        self.node_list.extend(vbs_nodes_list)
        self.node_list.append(self.master_node.ip)
        self.node_list.append(self.slaver_node.ip)
        # 去重
        self.node_list = list(set(self.node_list))

    def del_upg_pkg(self):
        """
        删除已上传的升级包
        """
        logger.info("Query upgrade package id.")
        ret_result, ret_data = self.opr.query_upg_pkg()
        if ret_result.get("code") != "0":
            err_msg = "Failed to query upgrade package id, Details: %s." % ret_data
            logger.error(err_msg)
            raise Exception(err_msg)
        if not ret_data:
            logger.info("There is no upgrade package info.")
            return
        package_id = ret_data.get("id")
        ret_result, ret_data = self.opr.del_upg_pkg(package_id)
        if ret_result.get("code") != "0":
            err_msg = "Failed to del upgrade package[%s], Details: %s." % (package_id, ret_data)
            logger.error(err_msg)

    @retry_adorn(stop=stop_after_attempt(5), wait=wait_fixed(120), reraise=True)
    def roll_back_post_check(self):
        """
        回退确认前检查
        """
        PostUpgradeCheck(self.project_id, self.pod_id, self.fs_args).procedure()

    def roll_back_commit(self):
        """
        回退确认
        """
        UpgradeConfirm(self.project_id, self.pod_id, self.fs_args).procedure()

    def restore_product_db(self):
        UpgradeHotPatchOperate.restore_db_data(self.remote_path, self.master_client, step="product")

    def restore_product_repo(self):
        UpgradeHotPatchOperate.restore_repo_data(self.remote_path, self.master_client)

    def upload_script_to_master_fsm(self, master_node):
        UpgradeHotPatchOperate.upload_script_to_fsm(self.master_client, master_node, self.remote_path)

    def update_rollback_state(self, state=0):
        """
        更新记录当前升级步骤
        0：补丁回滚初级状态
        1：补丁回滚完成完成
        2：补丁回滚确认完成
        3：恢复大包数据库完成
        4：恢复大包软件仓完成
        5：回滚大包完成
        """
        param_key = UpgradeStateKey.ROLLBACK_STATE.format(self.fs_args.get("suit_id"))
        ParamUtil().set_service_cloud_param(
            self.project_id, SERVICE_NAME, self.region_id, param_key, state)

    def get_rollback_state(self):
        param_key = UpgradeStateKey.ROLLBACK_STATE.format(self.fs_args.get("suit_id"))
        state = ParamUtil().get_value_from_cloud_param(
            self.project_id, SERVICE_NAME, param_key, self.region_id)
        return int(state)

    def get_product_upgrade_state(self):
        param_key = UpgradeStateKey.PRODUCT_UPGRADE_STATE.format(self.fs_args.get("suit_id"))
        return UpgradeHotPatchOperate.get_operate_state(self.project_id, self.region_id, param_key)

    def get_hotpatch_upgrade_state(self):
        try:
            param_key = UpgradeStateKey.HOTPATCH_UPGRADE_STATE.format(self.fs_args.get("suit_id"))
            return UpgradeHotPatchOperate.get_operate_state(self.project_id, self.region_id, param_key)
        except Exception as e:
            logger.info("Hotpatch upgrade is not start, details:%s" % e)
            return -1

