# /*
#  * 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/>.
#  */

import os
import logging
import subprocess
from SystemUpdater.modules.i18n import _
from SystemUpdater.Core.errors import *
from SystemUpdater.Core.enums import *
from SystemUpdater.checks.registry import register_check

@register_check("check-download-space", roles=[ROLE_DOWNLOAD])
def check_download_space(trans):
    size = int(trans.update_metadata.get("download-size", 0))
    logging.info(f"Check: Need download size {size} Bytes...")
    check_free_space(size, 0)

@register_check("check-deploy-space", roles=[ROLE_DEPLOY])
def check_deploy_space(trans):
    install_size = int(trans.update_metadata.get("install-size", 0))
    logging.info(f"Check: Need deploy size {install_size} Bytes...")

    ovl_size = 0
    ovl_path = f"/sysroot/ostree/pkgs/ovl-{trans.from_revision}.{trans.booted_serial}"
    if os.path.exists(ovl_path):
        ovl_size = get_disk_usage(ovl_path)
        logging.info(f"Check: Need ovl({ovl_path}) size {ovl_size} Bytes...")
    else:
        logging.warning(f"Check: OVL path '{ovl_path}' not exist, skip check...")

    if trans.ovl_size > ovl_size:
        logging.info(f"Use ostree ovl_size{trans.ovl_size} Bytes and default size:{ovl_size}...")
        ovl_size = trans.ovl_size

    check_free_space(ovl_size, install_size)

def get_disk_usage(path: str) -> int:
    try:
        result = subprocess.run(
            ['du', '-sb', path],
            capture_output=True, text=True, check=False, timeout=20
        )
        if result.stdout:
            size_str = result.stdout.split(maxsplit=1)[0]
            return int(size_str)
        else:
            logging.error(f"No output from 'du' command for path {path}")
        return 0
    except Exception as e:
        logging.error(f"Failed to get disk usage for {path}: {e}")
        return 0

class diskManager():
    def __init__(self) -> None:
        self.mounted_list = None
        self.mounted_point_free_space = None
        self.init_finish = False
        self._start()

    def _start(self):
        try:
            self.mounted_list = self._get_mount_points()
            self.mounted_point_free_space = self._get_mounted_point_free_space()
            # 无效打印 先进行屏蔽
            # for mount, free_mb in self.mounted_point_free_space.items():
            #     print(f"{mount}: {free_mb} MB free")
            self.init_finish = True
        except Exception as e:
            logging.error(e)

    def _get_mount_points(self):
        # 获得系统的挂载点列表
        # get the mount points list from /proc/mounts
        mounted = []
        try:
            with open("/proc/mounts") as mounts:
                for line in mounts:
                    try:
                        what, where, fs, options, *_  = line.split()
                    except ValueError as e:
                        logging.error("line '%s' in /proc/mounts not understood (%s)" % (line, e))
                        continue
                    if os.path.exists(where):
                        mounted.append(where)
        except Exception as e:
            logging.error("Failed to read /proc/mounts: %s", e)
            return []
        return sorted(mounted, key=len, reverse=True)
    
    def _get_mounted_point_free_space(self):
        # 获取系统每个挂载点的剩余空间
        fs_free = {}
        for mount_point in self.mounted_list:
            try:
                statvfs = os.statvfs(mount_point)
                free_bytes = statvfs.f_frsize * statvfs.f_bavail
                fs_free[mount_point] = free_bytes
            except Exception as e:
                logging.warning("Failed to get free space for %s: %s", mount_point, e)
        return fs_free
    
    def _get_mounted_point_by_dir(self, path):
        """
        给定一个路径，返回其所在的挂载点，使用最长前缀法匹配。
        """
        path = os.path.realpath(path)
        best_match = "/"
        best_len = -1

        for mount_point in self.mounted_list:
            mount_real = os.path.realpath(mount_point.rstrip("/"))
            if path == mount_real or path.startswith(mount_real + "/"):
                if len(mount_real) > best_len:
                    best_match = mount_point
                    best_len = len(mount_real)

        return best_match

    # External interface
    def check_space_size_by_size(self, download_size, additional_size):
        '''
        download_size: 下载包的总大小
        additional_size: 除下载空间外还需要额外的空间大小
        通过输入的需求安装列表，计算挂载点大小
        安装下载的地址通过 archivedir = apt_pkg.config.find_dir("Dir::Cache::archives") 自动获取
        '''
        if self.init_finish is True:
            result = True
            install_req_size = {}
            notfree_mounted_pointed = []
            try:
                # 下载路径
                # archivedir = apt_pkg.config.find_dir("Dir::Cache::archives")
                # V11上使用ostree，安装路径为/ostree/repo
                # 除安装空间外还预留一些安全空间缓冲区，设置为固定值
                for (dir, size) in [
                                    ("/", download_size + additional_size + 500*1024*1024),
                                    ("/boot", 200*1024*1024)
                                    ]:
                    mounted_point = self._get_mounted_point_by_dir(dir)
                    logging.info(f"Path '{dir}' mount on {mounted_point}, '{dir}' needs {size} Bytes size.")
                    if size < 0:
                        # 兼容出现负数的BUG
                        continue
                    
                    # 计算总大小
                    if mounted_point in install_req_size:
                        install_req_size[mounted_point] += size
                    else:

                        install_req_size[mounted_point] = size
                
                for mounted_point in install_req_size:
                    logging.info(f"Mount point '{mounted_point}' all needs {install_req_size[mounted_point]} Bytes size.")
                    logging.info(f"Mount point '{mounted_point}' free {self.mounted_point_free_space[mounted_point]} Bytes size.")
                    if self.mounted_point_free_space[mounted_point] > install_req_size[mounted_point]:
                        continue
                    else:
                        result = False
                        notfree_mounted_pointed.append(mounted_point)
                        logging.error(f"Path '{mounted_point}' needs {install_req_size[mounted_point]} Bytes, But only free {self.mounted_point_free_space[mounted_point]} Bytes")

                if result is False:
                    logging.error(f"Not enough free space:{notfree_mounted_pointed}")
                else:
                    logging.info("Check: Enough free space For system update...")
            except Exception as e:
                logging.error(e)
                result = False

        else:
            # 如果初始化失败，默认不要影响更新流程
            result = True
            logging.error("diskManager Initialization is not complete")
            return result

        return result
    
def check_free_space(download_size, addition_size=0):
    '''
    download_size和addition_size单位都是Bytes
    '''
    try:
        _diskManager = diskManager()
        result =_diskManager.check_space_size_by_size(download_size, addition_size)
        if result is False:
            raise UpdateBaseError(ERROR_NOT_DISK_SPACE)
    except Exception as e:
        logging.error(e)
        raise UpdateBaseError(ERROR_NOT_DISK_SPACE)

    return True