# /*
#  * Copyright (C) 2020, KylinSoft Co., Ltd.
#  *
#  * This program is free software: you can redistribute it and/or modify
#  * it under the terms of the GNU General Public License as published by
#  * the Free Software Foundation, either version 3 of the License, or
#  * (at your option) any later version.
#  *
#  * This program is distributed in the hope that it will be useful,
#  * but WITHOUT ANY WARRANTY; without even the implied warranty of
#  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  * GNU General Public License for more details.
#  *
#  * You should have received a copy of the GNU General Public License
#  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
#  */

# UpdateManager.py
# -*- Mode: Python; indent-tabs-mode: nil; tab-width: 4; coding: utf-8 -*-
import os
import pwd
import json
import shutil
import dbus,time,sys
import traceback
import logging
import apt_pkg
apt_pkg.init_config()
apt_pkg.init_system()
from .Core.errors import *
from .Core.enums import *
from gi.repository import GObject
from .Core.loop import mainloop
from SystemUpdater.modules.i18n import _
from dbus.mainloop.glib import DBusGMainLoop
DBusGMainLoop(set_as_default=True)
from .Core.Database import Sqlite3Server
from gi.repository import GLib
from SystemUpdater.UpdatesAvailable import Transaction
from SystemUpdater.modules.system import get_dist,system_shutdown
from .push_update import update_important,get_push_content
from SystemUpdater.Core.PlymouthManager import PlymouthManager
from SystemUpdater.checks.registry import run_checks
from SystemUpdater.modules.rate_limiter import get_download_limit
from SystemUpdater.configs import G_settings
from SystemUpdater.strategies import G_strategies
from SystemUpdater.Core.dbus_helper import get_bus
from SystemUpdater.update_hooks import scan_hooks
from .backend.BackendBackup import BackUpBackend 
from .backend.BackendOstreeNext import UpdateBackend,DownloadBackend,DeployBackend,CacheBackend
from .dataMigrator.datamigrator import MigrationBackend
from .dataMigrator.scanner import get_data_migration_dict, get_data_migration_osname
from .UpdateManagerDbus import UpdateManagerDbusController,UPDATER_DBUS_INTERFACE,UPDATER_DBUS_PATH,UPDATER_DBUS_SERVICE
from SystemUpdater.self_upgrade.hooks import execute_hooks
from SystemUpdater.constants import OUTPUT_JSON_PATH,KYLIN_OS_NAME,UPDATER_PACKAGE_NAME,HISTORY_LIST_LEN,CONTROL_CONTER_SENDER

class UpdateManager(GObject.GObject):
    __gsignals__ = {
                    "transaction-stage": (GObject.SignalFlags.RUN_FIRST,
                                         None,
                                         (GObject.TYPE_PYOBJECT,)),
                    "transaction-progress": (GObject.SIGNAL_RUN_FIRST,
                                GObject.TYPE_NONE,
                                (   GObject.TYPE_STRING,
                                    GObject.TYPE_INT
                                )),
                    }
    def __init__(self,options):
        try:
            GObject.GObject.__init__(self)
            self.cancellable = None
            self.options = options
            self.sqlite3_server = Sqlite3Server(self)

            self._check_exception_status()
            
            self.bus = dbus.SystemBus()
            self.dbus_send = self._setup_dbus()

            self.distro = get_dist()

            self._init_function()
            
        except Exception as e:
            logging.error(e)

    def _init_function(self):
        # 非麒麟系统不获取推送内容 
        if KYLIN_OS_NAME != self.distro:
            G_settings["enable_push"] = False

        cache = CacheBackend(self.options.sysroot,self.options.os)
        active_revision = cache.get_current_deployed().asdict()
        G_settings["active_revision"] = active_revision

    def run(self):
        """Start the daemon and listen for calls."""
        logging.info("Waiting for calls...")
        try:
            mainloop.run()
        except KeyboardInterrupt:
            self.dbus_send.Quit(None)

    def run_transaction(self, trans: Transaction):
        """Run the worker"""
        try:
            self.cancellable = trans._cancellable
            trans.download_limit = get_download_limit()
            trans.is_kylin = KYLIN_OS_NAME == self.distro
            G_strategies.reload()

            # 在执行各种事务之前检查状态
            run_checks(trans, G_settings)

            if trans.role == ROLE_UPDATE_CACHE:
                self.start_update(trans)
            if trans.role == ROLE_UPDATE_CACHE_NEXT:
                self.start_update_cache(trans)
            elif trans.role == ROLE_DOWNLOAD:
                self.start_download(trans)
            elif trans.role == ROLE_DEPLOY:
                self.start_upgrade(trans)

        except UpdateBaseError as excep:
            trans.success = False
            trans.error = excep.code
            self._failed_signal(trans,excep.code)
        except Exception as e:
            logging.error(e)
            tbk = traceback.format_exc()
            logging.error(str(tbk))
            self._failed_signal(trans,ERROR_PROGRAM_EXCEPTION)
        finally:
            pass

    def _failed_signal(self,trans,error_code):
        if trans.role in (ROLE_UPDATE_CACHE,ROLE_UPDATE_CACHE_NEXT):
            self.dbus_send.UpdateDetectFinishedNext(False,[''],error_code)
        elif trans.role == ROLE_DOWNLOAD:
            self.dbus_send.UpdateDownloadFinishedNext(False,[''],error_code)
        elif trans.role == ROLE_DEPLOY:
            self.dbus_send.DeployUpdatFinishedNext(False,[''],error_code)   

        self.emit("transaction-stage", trans)

    def start_cancel(self):
        self.cancellable.cancel()

    #进行更新的操作
    def start_update(self, trans: Transaction):
        if G_settings.get("enable_push",True):
            def _update_important_success(push_content):
                trans.origin_push = push_content
                trans.role = ROLE_UPDATE_CACHE_NEXT
                self.run_transaction(trans)
            
            def _update_important_failed(trans: Transaction,error_code):
                # 当更新推送列表失败时，仍然进行更新操作
                trans.origin_push = []
                trans.delay_error = error_code
                logging.warning("Update important list failed with error code: %s", error_code)
                trans.role = ROLE_UPDATE_CACHE_NEXT
                self.run_transaction(trans)

            update_important(trans,reply_handler=_update_important_success,error_handler=_update_important_failed)
            return

        trans.origin_push = get_push_content()
        self.start_update_cache(trans)

    #进行更新的操作
    def start_update_cache(self,trans: Transaction):
        def _download_info(parent,fetched,requested,downloaded,download_all,bytes_sec):
            self.dbus_send.UpdateDownloadInfo(fetched,requested,downloaded,download_all,bytes_sec)

        def _force_kill_self():
            logging.warning("Forcefully terminating current process using os._exit(0)")
            os._exit(0)

        def _update_finished(parent,trans: Transaction):
            if trans.error == PRIORITY_UPGRADE_SUCCCESSED:
                self.dbus_send.UpdateDetectFinishedNext(trans.success,[UPDATER_PACKAGE_NAME],trans.error,trans.error_message)
                GLib.timeout_add_seconds(3,_force_kill_self)
                self.dbus_send.Quit()
                return
            
            trans.hooks = scan_hooks()
            trans.hooks["post-check.d"] = execute_hooks(trans.hooks.get("post-check.d",[]))
            if trans.success:
                if trans.push_content:
                    self._make_meta_for_panel(trans.push_content,OUTPUT_JSON_PATH)
                    self.dbus_send.UpdateDetectFinishedNext(True,[p.update_type for p in trans.push_content])
                else:
                    logging.info("System is up-to-date...")
                    self.dbus_send.UpdateDetectFinishedNext(trans.success,[])
            else:
                if not trans.origin_push and not trans.delay_error and trans.is_kylin:
                    logging.info("Force to show system up-to-date...")
                    self.dbus_send.UpdateDetectFinishedNext(True,[])
                else:
                    self.dbus_send.UpdateDetectFinishedNext(trans.success,[""],trans.delay_error if trans.delay_error else trans.error,trans.error_message)

            if trans.sender == CONTROL_CONTER_SENDER or G_settings.get("replace",False):
                G_settings["current_update"] = trans.asdict()
            self.emit("transaction-stage",trans)

        update_backend = UpdateBackend(self.options.sysroot,self.options.os)
        update_backend.connect("transaction-done", _update_finished)
        update_backend.connect("download-info", _download_info)
        update_backend.run(trans)

    def start_download(self,trans: Transaction):
        def _download_info(parent,fetched,requested,downloaded,download_all,bytes_sec):
            self.dbus_send.UpdateDownloadInfo(fetched,requested,downloaded,download_all,bytes_sec)
            self.emit("transaction-progress",trans.role,self.calc_progress(downloaded,download_all))

        def _download_finished(parent,trans):
            self.dbus_send.UpdateDownloadFinishedNext(trans.success,[trans.selected.update_type],trans.error,trans.error_message)
            if trans.success:
                # 标记已下载在推送更新中
                for item in trans.push_content:
                    if item.revision == trans.selected.revision:
                        item.downloaded = True
                        break
                G_settings["current_update"] = trans.asdict()
            self.emit("transaction-stage",trans)

        download_backend = DownloadBackend(self.options.sysroot,self.options.os)
        download_backend.connect("transaction-done", _download_finished)
        download_backend.connect("download-info", _download_info)
        download_backend.run(trans)

    def start_upgrade(self,trans: Transaction):
        trans.status = STATUS_DEPLOYING
        trans.should_backup = G_strategies.get("autoupdate", {}).get("autobackup", True) \
                        and not self.options.prohibit_backup
        func_feedback = self.start_backup if trans.should_backup else self.start_deploy
        func_feedback(trans)

    def _upgrade_finished(self,trans: Transaction):
        trans.status = STATUS_FINISHED
        G_settings["current_update"] = {}
        G_settings["history"] = ([trans.asdict()] + G_settings.get("history", []))[:HISTORY_LIST_LEN]
        
        trans.hooks["post-deploy.d"] = execute_hooks(trans.hooks.get("post-deploy.d",[]))
        # 还原流程 通过post-backup-flag.py脚本进行判断，在目录scripts/post-backup-flag.py中

        self.emit("transaction-stage",trans)
        self.dbus_send.DeployUpdatFinishedNext(trans.success,[trans.selected.update_type],trans.error,trans.error_message)
        PlymouthManager.show_countdown(trans.success,trans.reboot_required)
        system_shutdown(trans.reboot_required)

    def start_backup(self,trans: Transaction):
        plymouth = PlymouthManager()
        def _progress_info(parent,progress):
            logging.info("Backup progress: %d%%", progress)
            plymouth.update(progress,_("The system is installing updates. Please do not turn off your computer."),1,10)

        def _backup_finished(parent,result):
            logging.info("Backup finished with result: %s", str(result))
            trans.backup_success = bool(result)
            if trans.backup_success:
                self.start_deploy(trans, plymouth)
            else:
                trans.success = False
                trans.error = ERROR_BACKUP_FAILED
                self._upgrade_finished(trans)
        
        backup_backend = BackUpBackend(self)
        backup_backend.connect("transaction-done", _backup_finished)
        backup_backend.connect("backup-progress", _progress_info)
        backup_backend.run(trans)

    def start_deploy(self,trans: Transaction, plymouth=None):
        plymouth = plymouth or PlymouthManager()
        def _deploy_progress(parent, progress):
            self.dbus_send.UpdateDloadAndInstStaChanged([trans.selected.update_type],progress,"","")
            plymouth.update(progress,_("The system is installing updates. Please do not turn off your computer."),10,90)
        
        def _deploy_finished(parent, trans):
            logging.info("Deploy finished with result: %s", str(trans.success))
            if trans.success:
                self.start_migration(trans, plymouth)
            else:
                self._upgrade_finished(trans)

        _deploy_progress(self,1)
        G_settings["current_update"] = trans.asdict() # 同步当前的状态给脚本读取使用
        trans.hooks["pre-deploy.d"] = execute_hooks(trans.hooks.get("pre-deploy.d",[]))
        deploy_backend = DeployBackend(self.options.sysroot,self.options.os)
        deploy_backend.connect("transaction-done", _deploy_finished)
        deploy_backend.connect("deploy-progress", _deploy_progress)
        deploy_backend.run(trans)

    def start_migration(self,trans: Transaction, plymouth=None):
        logging.info("----------Start migration----------")
        plymouth = plymouth or PlymouthManager()
        def _migration_progress(parent, status, pkg, percent):
            self.dbus_send.UpdateDloadAndInstStaChanged([trans.selected.update_type],percent,"","")
            plymouth.update(percent,_("The system is installing updates. Please do not turn off your computer."),90,100)

        def _migration_finished(parent, success, status):
            logging.info(f"Migration finished with result: {success} and status: {status}")
            trans.success = success
            trans.error = status if not success else ''
            self._upgrade_finished(trans)

        new_version = trans.new_revision
        # 获取OS名称并构建deploy_base_dir
        osname = get_data_migration_osname()
        deploy_base_dir = f"/sysroot/ostree/deploy/{osname}/deploy"
        logging.info(f"Using deploy_base_dir: {deploy_base_dir}")
        # 传递deploy_base_dir参数
        migration_backend = MigrationBackend(new_version, deploy_base_dir=deploy_base_dir)
        trans.data_migration_path_chroot = migration_backend.chroot_output_dir
        trans.data_migration_path = migration_backend.output_dir
        G_settings["current_update"] = trans.asdict()
        data_migration_dict, data_migration_dict_old = get_data_migration_dict()
        migration_backend.connect("migration-done", _migration_finished)
        migration_backend.connect("migration-progress", _migration_progress)
        migration_backend.run(data_migration_dict, data_migration_dict_old)

    def _make_meta_for_panel(self,push_content, output_path):
        for push in push_content:
            metadata = push.metadata
            try: 
                if not os.path.exists(output_path):
                    os.makedirs(output_path)
                    logging.info('making the ConfigPath(%s) is complete...',output_path)

                update_info = {}
                update_info.update(metadata)
                
                # 适配旧版本，在没有更新name
                if "update-name" not in update_info:
                    update_name = {"zh_CN":update_info.get("update-name_zh_CN","系统更新"),"en_US":update_info.get("update-name_en_US","System update")}
                    update_info.update({"update-name":update_name})

                # 适配旧版本以及新版本元数据的制定
                if isinstance(update_info.get('version',{}), str):
                    update_info.update({"version":update_info.get('version',"")})
                elif isinstance(update_info.get('version',{}), dict):
                    update_info.update({"version":update_info.get('version',{}).get("value","")})

                # 当description是字符串的时候 转化为字典类型
                if isinstance(update_info.setdefault('description',{"zh_CN":"","en_US":""}), str):
                    temp_dict = {"zh_CN":update_info.get("description",""),"en_US":update_info.get("description","")}
                    update_info.update({"description":temp_dict})

                if isinstance(update_info.setdefault('changelog',{"zh_CN":"","en_US":""}), str):
                    temp_dict = {"zh_CN":update_info.get("changelog",""),"en_US":update_info.get("changelog","")}
                    update_info.update({"changelog":temp_dict})
                
                update_info.update({"total_download_size":int(update_info.get("download-size",0))})
                update_info.update({"total_install_size":int(update_info.get("install-size",0))})

                update_info.update({"icon":update_info.setdefault("icon","/usr/share/upgrade/icon/kylin-update-desktop-system.png")})

                output_config_name = os.path.join(output_path, update_info.get("update-type") + '.json')
                with open(output_config_name, 'w', encoding='utf-8') as f:
                    json.dump(update_info, f, ensure_ascii=False, indent=4)                   
                logging.info("Generate Jsonfile(%s) to complete... ",output_config_name)
            except Exception as e:
                logging.error(e)

    def _check_exception_status(self):
        """检查异常状态，例如部署过程中异常关机。"""
        VAR_CONFIG_DIR = '/var/lib/kylin-system-updater'
        os.makedirs(VAR_CONFIG_DIR, exist_ok=True)

        current_update = G_settings.get("current_update")
        if not current_update:
            return

        trans = Transaction(**current_update)

        # 标识：部署过程中异常关机
        if trans.status == STATUS_DEPLOYING:
            trans.success = False
            trans.error = ERROR_INSTALL_POWER_OUTAGE
            self.emit("transaction-stage", trans)

            # 更新历史记录并清理当前事务
            history = G_settings.get("history", [])
            G_settings["history"] = [trans.asdict(), *history]
            G_settings["current_update"] = {}

            # 清理数据迁移相关缓存
            if os.path.exists(trans.data_migration_path):
                try:
                    shutil.rmtree(trans.data_migration_path)
                except Exception as e:
                    logging.error(f"Clean data migration cache {trans.data_migration_path} failed:")
                    logging.error(e)
            if os.path.exists(trans.data_migration_path_chroot):
                try:
                    shutil.rmtree(trans.data_migration_path_chroot)
                except Exception as e:
                    logging.error(f"Clean data migration cache {trans.data_migration_path_chroot} failed:")
                    logging.error(e)
            

    def _setup_dbus(self):
        # check if there is another g-a-i already and if not setup one
        # listening on dbus
        try:
            bus = get_bus()
            bus_name = dbus.service.BusName(UPDATER_DBUS_SERVICE,
                                            bus,do_not_queue=True)
            logging.info("Initiate dbus success ...")
            return  UpdateManagerDbusController(self,bus_name)
        except dbus.exceptions.NameExistsException:
            if not G_settings.get("replace",False):
                logging.critical("Another daemon is already running")
                sys.exit(1)
            logging.warning("Replacing already running daemon")
            
            retry_reboot_times = 0
            the_other_guy = bus.get_object(UPDATER_DBUS_SERVICE,
                                            UPDATER_DBUS_PATH)
            the_other_guy.Quit(dbus_interface=UPDATER_DBUS_INTERFACE,
                                timeout=300)
            time.sleep(1)
            while True:
                retry_reboot_times = retry_reboot_times + 1
                #当重试次数超过5次时退出程序
                if retry_reboot_times > 5:
                    logging.critical("Reboot backend is Failed...")
                    sys.exit(1)
                try:
                    bus_name = dbus.service.BusName(UPDATER_DBUS_SERVICE,
                                                    bus,
                                                    do_not_queue=True)
                    logging.warning("Replacing already running daemon to Success...")
                    return  UpdateManagerDbusController(self,bus_name)
                except dbus.exceptions.NameExistsException:
                    the_other_guy = bus.get_object(UPDATER_DBUS_SERVICE,
                                                    UPDATER_DBUS_PATH)
                    the_other_guy.Quit(dbus_interface=UPDATER_DBUS_INTERFACE,
                                        timeout=300)
                    logging.warning("Dbus has not withdrawn and retry reboot times:%d...",retry_reboot_times)
                    time.sleep(1)

    def calc_progress(self, downloaded: int, total: int) -> int:
        if total <= 0:
            return 1

        progress = (downloaded * 100) // total
        return min(progress, 100)

    # 获取当前deploy的系统hash
    def get_cur_sys_hash(self):
        update_backend = UpdateBackend(self.options.sysroot,self.options.os)
        return update_backend.cur_deployment.get_csum()
