#!/usr/bin/python3
import configparser
import ctypes
import json
import logging
import os
import shutil
import stat
import sys
import zipfile
from enum import Enum
from functools import wraps

from install import CommonMethod


class Flag(Enum):
    init = 1
    upgrading = 2


conf = configparser.ConfigParser()
SYSTEM_WINDOWS = 'Windows'
PERMISSION_700 = stat.S_IRWXU
PERMISSION_500 = stat.S_IRUSR | stat.S_IXUSR
PERMISSION_600 = stat.S_IWUSR | stat.S_IRUSR
FLAGS = os.O_WRONLY | os.O_CREAT


class RollBackWrap(object):
    def __call__(self, func):
        @wraps(func)
        def wrapper(*args, **kw):
            if not func(*args, **kw):
                self._rollback_agent()
                return False
            return True

        return wrapper

    @staticmethod
    def _rollback_agent():
        logging.error("Failed to upgrade the agent assistant.")
        image_path = os.path.join(sys.argv[2], 'AgentAssist_Image')
        python_image_path = os.path.join(sys.argv[2], 'AgentAssistPython_Image')
        rm_result = CommonMethod.remove_path(image_path, python_image_path)
        if not rm_result:
            logging.error(f"Failed to upgrade the agent assistant and roll "
                          f"back the version. ")
            return False


class ConfigHandler(object):
    def __init__(self, path, encoding="UTF-8"):
        """类初始化时创建conf对象"""
        self._path = path
        self._encoding = encoding
        self._conf = configparser.ConfigParser()
        self._conf.read(self._path, self._encoding)

    def get_option(self, section, option):
        """获取指定section下的指定option，默认返回str"""
        value = None
        try:
            value = self._conf.get(section, option)
        except Exception as err:
            logging.error(f"Get {section} {option} failed, err:{err}.")
        return value


    def get_int_option(self, section, option):
        """获取指定section下的指定option，并强制转化为int"""
        value = None
        try:
            value = self._conf.getint(section, option)
        except Exception as err:
            logging.error(f"Get {section} {option} failed, err:{err}.")
        return value

    def set_option(self, section, option, value=None):
        """设置section下某个option的值，如果不存在option会直接创建，存在即更新"""
        if not self._conf.has_section(section):
            self._conf.add_section(section)
        try:
            self._conf.set(section, option, value)
        except Exception as err:
            logging.error(f"Write {section} {option}:{value} to config failed, err:{err}.")
            return False
        with os.fdopen(os.open(self._path, FLAGS, PERMISSION_600), 'w') as fd_open:
            self._conf.write(fd_open)
        return True


class AgentAssistUpgrade(CommonMethod):
    def __init__(self, upgrade_path, install_path):
        super(AgentAssistUpgrade, self).__init__(install_path)
        self._upgrade_pkg_path = os.path.realpath(upgrade_path)
        self._agent_assist_upgrade_log = "upgrade.log"
        self._log_path = os.path.join(self._upgrade_pkg_path,
                                      self._agent_assist_upgrade_log)
        self._upgrade_flag = Flag.init

    def _log_init(self, cur_path):
        pkg_path = cur_path.split("AgentAssistImage")[0]
        self._log_path = os.path.join(pkg_path, self._agent_assist_upgrade_log)
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
            datefmt='%a, %d %b %Y %H:%M:%S',
            filename=self._log_path,
            filemode='a')

    @RollBackWrap()
    def _unzip_upgrade_pkg(self):
        pkg_name = self._find_pkg(self._upgrade_pkg_path)
        os.chdir(self._upgrade_pkg_path)
        if not os.path.isfile(pkg_name):
            logging.error(f"pkg to be extract not exists, pkg_name:{pkg_name}")
            return False
        logging.info(f'The path of the upgrade package is: '
                     f'{self._upgrade_pkg_path}. pkg_name: {pkg_name}')
        pkg_path = os.path.join(self._upgrade_pkg_path, pkg_name)
        max_file_num = 2000
        max_file_size = 1024 * 1024 * 100 * 20
        total_size = 0
        try:
            arc_file = zipfile.ZipFile(pkg_path)
            if len(arc_file.namelist()) > max_file_num:
                logging.error("compressed file nums exceed maximum")
                return False
            for file in arc_file.filelist:
                total_size += file.file_size
                if total_size > max_file_size:
                    logging.error("compressed package has an exception")
                    return False
            arc_file.extractall(
                os.path.join(self._install_path, 'AgentAssist_Image'))
            arc_file.close()
            return True
        except Exception as err:
            logging.error(f"Failed to decompress the installation package."
                          f" {err}")
            return False

    @RollBackWrap()
    def check_path_space(self):
        if not os.path.exists(self._install_path):
            print("The installation directory does not exist.")
            return False
        if self._cur_system == SYSTEM_WINDOWS:
            free_bytes = ctypes.c_ulonglong(0)
            ctypes.windll.kernel32.GetDiskFreeSpaceExW(
                ctypes.c_wchar_p(self._install_path), None, None,
                ctypes.pointer(free_bytes))
            free_space = free_bytes.value / 1024 / 1024 / 1024
        else:
            st = os.statvfs(self._install_path)
            free_space = st.f_bavail * st.f_frsize / 1024 / 1024 / 1024
        if free_space < 5:
            logging.error("The space of the installation directory is less than"
                          " 5 GB. ")
            print("The space of the installation directory is less than 5 GB")
            return False
        return True

    # delete the configuration file  AgentAssist.conf
    def _delete_system_conf(self, name):
        file_path = os.path.join(self._install_path,
                                 'AgentAssist_Image/conf/AgentAssist.conf')
        conf.read(file_path)
        conf.remove_option("system", name)
        with os.fdopen(os.open(file_path, FLAGS, PERMISSION_600), 'w') as f:
            conf.write(f)

    def _do_copy_file(self, srcdir, desdir):
        file_list = ['VersionDetails', 'AgentAssist.conf']
        for file in os.listdir(srcdir):
            source_file = os.path.join(srcdir, file)
            target_file = os.path.join(desdir, file)
            if os.path.isfile(source_file):
                if file not in file_list:
                    shutil.copyfile(source_file, target_file)
            if os.path.isdir(source_file):
                folder = os.path.exists(target_file)
                if not folder:
                    os.makedirs(target_file)
                self._do_copy_file(source_file, target_file)

    def _update_conf_info(self):
        old_conf_file = os.path.join(self._install_path, 'AgentAssist/conf/AgentAssist.conf')
        new_conf_file = os.path.join(self._install_path, 'AgentAssist_Image/conf/AgentAssist.conf')
        if not os.path.isfile(old_conf_file) or not os.path.isfile(new_conf_file):
            logging.error("The AgentAssist conf file not existed.")
            return False

        params_list = AgentAssistUpgrade.get_conf_info(old_conf_file)
        if not AgentAssistUpgrade.write_new_conf(new_conf_file, params_list):
            logging.error("Write new conf failed.")
            return False

        return True

    @staticmethod
    def get_conf_info(config_file):
        old_conf = ConfigHandler(config_file)

        # system
        memory_limit = old_conf.get_option('system', 'memory_limit')
        cpu = old_conf.get_option('system', 'cpu')
        if cpu is None:
            cpu = old_conf.get_option('system', 'cpu_limit')
        time = old_conf.get_option('system', 'time')
        if time is None:
            time = old_conf.get_option('system', 'check_time_limit')

        # agent_log
        per_compression_cycle = old_conf.get_option('agent_log', 'per_compression_cycle')
        if per_compression_cycle is None:
            per_compression_cycle = old_conf.get_option('agent_log', 'per_compression_cycle_h')
        max_log_size = old_conf.get_option('agent_log', 'max_log_size')
        if max_log_size is None:
            max_log_size = old_conf.get_option('agent_log', 'max_log_size_mb')
        max_log_num = old_conf.get_option('agent_log', 'max_log_num')
        archive_time = old_conf.get_option('agent_log', 'archive_time')
        if archive_time is None:
            archive_time = old_conf.get_option('agent_log', 'archive_time_d')
        retention_time = old_conf.get_option('agent_log', 'retention_time')
        if retention_time is None:
            retention_time = old_conf.get_option('agent_log', 'retention_time_d')

        info_list = {'memory_limit': memory_limit, 'cpu': cpu, 'time': time,
                     'per_compression_cycle': per_compression_cycle, 'max_log_size': max_log_size,
                     'max_log_num': max_log_num, 'archive_time': archive_time, 'retention_time': retention_time}
        return info_list

    @staticmethod
    def write_new_conf(config_file, params):
        new_conf = ConfigHandler(config_file)
        new_conf.set_option("system", "memory_limit", params.get('memory_limit'))
        new_conf.set_option("system", "cpu_limit", params.get("cpu"))
        new_conf.set_option("system", "check_time_limit", params.get("time"))
        new_conf.set_option("agent_log", "max_log_size_mb", params.get("max_log_size"))
        new_conf.set_option("agent_log", "max_log_num", params.get("max_log_num"))
        new_conf.set_option("agent_log", "per_compression_cycle_h", params.get("per_compression_cycle"))
        new_conf.set_option("agent_log", "archive_time_d", params.get("archive_time"))
        new_conf.set_option("agent_log", "retention_time_d", params.get("retention_time"))
        logging.info('The AgentAssist.conf is updated successfully.')
        return True

    def _copy_conf_file(self):
        old_path = os.path.join(self._install_path, 'AgentAssist/conf')
        dir_path = os.path.join(self._install_path, 'AgentAssist_Image/conf')
        try:
            self._do_copy_file(old_path, dir_path)
        except Exception as err:
            logging.error(f'Failed to copy conf file to Image folder. {err}')
            return False
        logging.info("The conf file is successfully moved.")
        return True

    @RollBackWrap()
    def _update_version_num(self):
        conf_path = os.path.join(self._install_path, 'AgentAssist_Image/conf')
        json_path = os.path.realpath(os.path.join(conf_path, 'agent_list.json'))
        version_path = os.path.join(conf_path, 'VersionDetails')
        upgrade_version_num = self._read_conf_without_section(version_path,
                                                              'Package_Version')
        if not self._check_conf_parameter(upgrade_version_num):
            logging.error('upgrade assist version parameter illegal')
            return False
        if not os.path.isfile(json_path):
            return False
        try:
            with open(json_path, "r", encoding='utf8') as fr:
                info_dict = json.load(fr)
                info_dict["agent_list"][0]["file_name"][
                    "agent_version"] = upgrade_version_num
            with os.fdopen(os.open(json_path, FLAGS, PERMISSION_600), 'w') as fw:
                json_str = json.dumps(info_dict)
                fw.write(json_str)
        except Exception as err:
            logging.error(f"Failed to mod agent_list.json. {err}")
            return False
        return True

    def _copy_log_file(self):
        log_path = os.path.join(self._install_path, 'AgentAssist/log')
        image_log_path = os.path.join(self._install_path,
                                      'AgentAssist_Image/log')

        try:
            self._do_copy_file(log_path, image_log_path)
        except Exception as err:
            logging.error(f"Failed to copy logs to the images folder. {err}")

        logging.info("success to copy logs to the images folder.")

        log_bak_path = os.path.join(self._install_path, 'AgentAssist_Image',
                                    'logbak/log_bak_log')
        try:
            os.makedirs(log_bak_path, mode=0o700, exist_ok=True)
        except Exception as err:
            logging.error(f"Failed to create [{log_bak_path}] folder. {err}")

        logbak_path = os.path.join(self._install_path, 'AgentAssist/logbak')
        image_logbak_path = os.path.join(self._install_path,
                                         'AgentAssist_Image/logbak')
        try:
            self._do_copy_file(logbak_path, image_logbak_path)
        except Exception as err:
            logging.info(
                f"failed to copy logbak file to the images folder. {err}")

        logging.info(f"success to copy logbak file to the images folder. ")

    @RollBackWrap()
    def _upgrade_create_config(self):
        return self._create_config(is_upgrade=True)

    def upgrade(self):
        if not self._install_path:
            return False

        self._log_init(self._upgrade_pkg_path)

        if not self.check_path_space():
            return False

        # Decompressing the Upgrade Package
        if not self._unzip_upgrade_pkg():
            return False

        # Modify the configuration file
        if not self._copy_conf_file():
            return False

        # 同步更新配置信息
        if not self._update_conf_info():
            return False

        if not self._update_version_num():
            return False
        self._copy_log_file()

        if not self._upgrade_create_config():
            return False
        return True


if __name__ == '__main__':
    if len(sys.argv) < 3:
        print("The input parameters are incorrect.")
        logging.error('The input parameters are incorrect.')
        sys.exit(1)
    upgrade = AgentAssistUpgrade(sys.argv[1], sys.argv[2])
    if not upgrade.upgrade():
        print("failed")
        sys.exit(1)
    print("success")
    sys.exit(0)
