# update.py
# /*
#  * 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 stat
import json
import shutil
import logging
from typing import Dict,List
from pydantic import BaseModel

from .partial_checkout import checkout_files_partly
from .dpkg import simulate_apt_install_and_parse_versions,not_upgrade_pkgs_list
from .hooks import execute_hooks
from .config import VERSIONS_DIR, MANIFESTS_DIR, CURRENT_PROGRESS,VERSION_MANIFEST_FILE
from .manifest import (
    PostAction,
    ComponentManifest,
    parse_manifests,
    collect_installed_files_from_manifests,
    get_all_dpkg_list_files,
    get_version_policy,
    get_move_files_from_manifests,
    get_all_post_hooks,
    get_failure_rate,
    handle_move_operations,
    get_remapped_files
)

DPKG_STATUS_PATH = "/var/lib/dpkg/status"
VERSION_ID_FILE_NAME = "HEAD"

class VersionConfig(BaseModel):
    pull_ref: str = None               # 拉取来源分支
    commit_id: str = None              # 当前版本commit id
    packages: Dict[str, dict] = {}
    moved_files: Dict[str, str] = {}
    executed_hooks: List[PostAction] = []     # 所有更新组件自升级完成后 需要做的处理，例如重启服务
    failed_files: Dict[str, str] = {}          # 记录那些文件都拉取失败了
    manifests: Dict[str, ComponentManifest] = {}  # 每个组件的manifest描述

def load_current_version_info():
    """
    加载当前版本的 manifest 文件，返回 commit_id 和 moved_files 信息。

    Returns:
        Tuple[Optional[str], VersionConfig]: 包含 commit_id 和 moved_files 的元组。
    """
    version_manifest_path = os.path.join(CURRENT_PROGRESS, VERSION_MANIFEST_FILE)
    if not os.path.exists(version_manifest_path):
        return None, {}

    try:
        with open(version_manifest_path, 'r') as f:
            data = json.load(f)
    except Exception as e:
        logging.error(f"Failed to load version manifest: {e}")
        return None, {}

    commit_id = data.get('commit_id',None)
    moved_files = data.get('moved_files', {})

    return commit_id, moved_files

def _clean_old_versions(new_version_id: str, old_version_id: str):
    for version_name in os.listdir(VERSIONS_DIR):
        version_path = os.path.join(VERSIONS_DIR, version_name)
        if version_name not in {new_version_id, old_version_id} and os.path.isdir(version_path):
            logging.info("Deleting outdated version: %s", version_name)
            shutil.rmtree(version_path)

def _clean_old_config(new_moved_files: Dict[str, str], old_moved_files: Dict[str, str]):
    """
    清理旧版本遗留的配置文件。
    
    规则：
    - 如果旧版本迁移的源文件在新版本中已不存在，删除对应的目标文件。
    - 如果源文件存在，但目标路径变化，删除旧目标文件。
    - 如果源文件和目标路径均未变化，保留目标文件。
    """
    for src_file, old_dst_file in old_moved_files.items():
        new_dst_file = new_moved_files.get(src_file)

        if new_dst_file is None:
            # 新版本已无此迁移项，直接删除旧文件
            if os.path.exists(old_dst_file):
                logging.info("Deleting outdated config file: %s", old_dst_file)
                os.remove(old_dst_file)
        elif new_dst_file != old_dst_file:
            # 迁移路径发生变化，删除旧目标文件
            if os.path.exists(old_dst_file):
                logging.info("Deleting outdated config file due to path change: %s",old_dst_file)
                os.remove(old_dst_file)
        else:
            # 迁移路径未变化，保留文件
            logging.info("Keeping valid config file: %s", old_dst_file)

def _activate_new_version(version_path: str) -> None:
    """
    激活指定版本目录，通过替换 CURRENT_PROGRESS 链接指向该版本。

    参数:
        version_path (str): 新版本目录的路径，必须为有效目录。

    逻辑说明:
        - 校验版本路径有效性。
        - 清理当前 CURRENT_PROGRESS，支持符号链接、目录或普通文件情况。
        - 将当前版本目录重命名备份。
        - 将新版本目录重命名为 CURRENT_PROGRESS，实现版本切换。
    """
    if not os.path.isdir(version_path):
        logging.error("Invalid version path: %s", version_path)
        return

    # 清理当前链接（符号链接、目录或文件）
    if os.path.islink(CURRENT_PROGRESS):
        os.unlink(CURRENT_PROGRESS)  # 删除符号链接
    elif os.path.isdir(CURRENT_PROGRESS):
        version_file = os.path.join(CURRENT_PROGRESS, VERSION_ID_FILE_NAME)
        with open(version_file, 'r', encoding='utf-8') as f:
            current_commit_id = f.read().strip()
        old_move_post = os.path.join(VERSIONS_DIR, current_commit_id)
        os.rename(CURRENT_PROGRESS, old_move_post)

    try:
        # 重命名新版本目录为 CURRENT_PROGRESS，完成切换
        os.rename(version_path, CURRENT_PROGRESS)
        logging.info("Activated version: %s", CURRENT_PROGRESS)
    except OSError as e:
        logging.error("Failed to activate version '%s': %s", version_path, e)

def _should_upgrade(new_version_id: str) -> bool:
    """
    判断是否需要执行升级逻辑。

    :param new_version_id: 新版本的 commit ID
    :return: 如果需要升级，返回 True；否则返回 False
    """
    progress_path = CURRENT_PROGRESS
    version_file = os.path.join(progress_path, VERSION_ID_FILE_NAME)

    if not os.path.isdir(progress_path):
        logging.info("Upgrade needed: version path does not exist: %s", progress_path)
        return True

    if os.path.islink(progress_path):
        try:
            current_commit_id = os.path.basename(os.readlink(progress_path))
            if current_commit_id == new_version_id:
                logging.info("No upgrade needed: already running version %s", new_version_id)
                return False
        except OSError as e:
            logging.warning(f"Failed to read current symlink: {e}")
    else:
        try:
            with open(version_file, 'r', encoding='utf-8') as f:
                current_commit_id = f.read().strip()

            if current_commit_id == new_version_id:
                logging.info("No upgrade needed: already running version %s", new_version_id)
                return False

            logging.info("Upgrade needed: current version %s != new version %s", current_commit_id, new_version_id)
        except (OSError, FileNotFoundError) as e:
            logging.warning("Failed to read current version file '%s': %s", version_file, e)

    return True

def add_new_version(repo_obj, trans,version_root) -> VersionConfig:
    version_config = VersionConfig(
        pull_ref=trans.available_refs,
        commit_id=trans.new_revision,
    )

    logging.info("Checking out update manifests and status files...")
    failed_file = checkout_files_partly(repo_obj, trans, [MANIFESTS_DIR], [DPKG_STATUS_PATH], version_root)
    version_config.failed_files.update(failed_file)
    version_config.manifests = parse_manifests(version_root)
    if not version_config.manifests:
        logging.warning("No manifests found, aborting self upgrade add new version addition.")
        return None

    # 检查是否禁止自我升级
    if any(m.prohibit_upgrade for m in version_config.manifests.values()):
        logging.warning("Upgrade is prohibited by manifest rules, aborting self upgrade addition.")
        return None

    # 解析status文件，获取自更新包的版本以及依赖并拉取依赖        
    packages_list = [pkg for pkg, m in version_config.manifests.items() if m.use_dpkg_list]
    logging.info(f"Simulating apt install to get package depends {packages_list}...")
    packages_info = simulate_apt_install_and_parse_versions(version_root, packages_list)
    version_config.packages = packages_info
    extra_depends = list(set(packages_info.keys()) - set(version_config.manifests.keys()))
    logging.info(f"Extra dependencies from simulation: {extra_depends}")

    # 增加额外的依赖包到 manifests 中
    if extra_depends and all(m.append_depends for m in version_config.manifests.values()):
        version_config.manifests.update({pkg: ComponentManifest(package=pkg) for pkg in extra_depends})

    # 根据版本策略过滤不需要升级的包
    version_policy = get_version_policy(version_config.manifests)
    if version_policy:
        not_upgrade = not_upgrade_pkgs_list(packages_info, version_policy)
        version_config.manifests = {k: v for k, v in version_config.manifests.items() if k not in not_upgrade}
        if not version_config.manifests:
            logging.warning("No manifests found, aborting self upgrade add new version addition.")
            return None

    logging.info("Checking out dpkg list...")
    dpkg_files = get_all_dpkg_list_files(version_config.manifests)
    failed_file = checkout_files_partly(repo_obj, trans, [], dpkg_files, version_root)
    version_config.failed_files.update(failed_file)

    logging.info("Starting checkout package files:%s",packages_info)
    all_files = collect_installed_files_from_manifests(version_root,version_config.manifests)
    if not all_files:
        logging.info("No files need to pull and aborting self upgrade addition...")
        return version_config

    failed_file = checkout_files_partly(repo_obj, trans, [], all_files, version_root)
    version_config.failed_files.update(failed_file)

    # 检查是否超过一定比例的失败 超过时 认为自更新失败
    failed_limit = int(len(all_files) * get_failure_rate(version_config.manifests))
    if len(failed_file) > failed_limit:
        logging.warning(f"Failed files({len(failed_file)}) larger than {failed_limit} files checked out, aborting...",)
        return None

    return version_config

def _add_version_id_file(version_path: str) -> None:
    """
    在指定版本目录下创建一个包含 commit ID 的标识文件。
    
    :param version_path: 版本目录的路径，该目录名应为 commit ID。
    """
    os.makedirs(version_path, exist_ok=True)
    version_file = os.path.join(version_path, VERSION_ID_FILE_NAME)
    commit_id = os.path.basename(os.path.normpath(version_path))

    with open(version_file, 'w', encoding='utf-8') as f:
        f.write(commit_id)

def _grant_executable_permissions(hooks_dir: str) -> None:
    """
    递归扫描指定目录（base_dir/sub_path），为其中所有文件添加执行权限。

    Args:
        base_dir (str): 根目录路径，如 CURRENT_PROGRESS。
        sub_path (str): 相对路径，如 'hooks' 或 MANIFESTS_DIR 下的 hooks。

    Example:
        grant_executable_permissions(os.path.join(CURRENT_PROGRESS, MANIFESTS_DIR.lstrip('/')), 'hooks')
    """

    if not os.path.isdir(hooks_dir):
        logging.debug("No hooks directory found at: %s", hooks_dir)
        return

    updated = 0
    try:
        for root, _, files in os.walk(hooks_dir):
            for file in files:
                file_path = os.path.join(root, file)
                try:
                    mode = os.stat(file_path).st_mode
                    os.chmod(file_path, mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
                    updated += 1
                    logging.debug("Granted execute permission: %s", file_path)
                except Exception as e:
                    logging.warning("Failed to chmod %s: %s", file_path, e)
        logging.info("Granted executable permissions to %d files under %s", updated, hooks_dir)
    except Exception as e:
        logging.error("Error while processing hooks directory %s: %s", hooks_dir, e)

def _update_current_version_id(commit_id: str, version_path: str = CURRENT_PROGRESS) -> None:
    """
    在指定版本目录下创建或更新版本标识文件。
    如果文件已存在，会先读取旧内容并写入日志，再覆盖为新的 commit ID。
    
    :param version_path: 版本目录路径，通常为 commit ID 命名
    :param commit_id: 需要写入的 commit ID
    """
    os.makedirs(version_path, exist_ok=True)
    version_file = os.path.join(version_path, VERSION_ID_FILE_NAME)

    if os.path.exists(version_file):
        try:
            with open(version_file, 'r', encoding='utf-8') as f:
                old_content = f.read().strip()
            logging.info("Version ID file already exists: %s (content: %s)", version_file, old_content)
        except Exception as e:
            logging.warning("Failed to read existing version ID file %s: %s", version_file, e)
    else:
        logging.info("Version ID file does not exist, will create: %s", version_file)

    try:
        with open(version_file, 'w', encoding='utf-8') as f:
            f.write(commit_id)
        logging.info("Version ID file updated with commit ID: %s", commit_id)
    except Exception as e:
        logging.error("Failed to write version ID file %s: %s", version_file, e)

def run_priority_upgrade(repo_obj, trans):
    new_version_id = trans.new_revision
    if not _should_upgrade(new_version_id):
        return False,[]

    # 创建/opt/update-module/versions/下面根目录
    version_root = os.path.join(VERSIONS_DIR, new_version_id)
    shutil.rmtree(version_root, ignore_errors=True)
    os.makedirs(version_root, exist_ok=True)

    # 下载优先升级版本内容
    version_config = add_new_version(repo_obj, trans,version_root)
    if not version_config:
        # 优先升级失败或没有描述文件或者不需要升级时 仅打标完成优先升级 下次不再触发
        shutil.rmtree(version_root, ignore_errors=True)
        _update_current_version_id(new_version_id)
        return False,[]

    # 进行规则映射，将配置转移到/opt规则下面
    remapped_files = get_remapped_files(version_config.manifests)
    has_moved = handle_move_operations(remapped_files,version_root)
    version_config.moved_files.update(has_moved)
    
    # 处理需要复制文件到指定的位置
    move_files = get_move_files_from_manifests(version_config.manifests)
    has_moved = handle_move_operations(move_files,version_root)
    version_config.moved_files.update(has_moved)

    # 增加记录版本id的文件，方便进行解析
    _add_version_id_file(version_root)

    # 获取旧版本的基础信息
    old_version_id,old_move_files = load_current_version_info()

    # 最后进行激活目录，保证原子性，激活之后 所有的代码被移动到/opt/update-module/current目录下面
    _activate_new_version(version_root)

    # 删除旧版本与清除旧配置
    _clean_old_versions(new_version_id, old_version_id)
    _clean_old_config(version_config.moved_files, old_move_files)

    # 导出配置到本地，首次写入 来支持下面执行的脚本 需要读取这个配置
    manifest_path = os.path.join(CURRENT_PROGRESS, VERSION_MANIFEST_FILE)
    with open(manifest_path, "w") as f:
        json.dump(json.loads(version_config.json()), f, indent=4)

    # 为hooks目录下面的脚本加执行权限
    hooks_dir = os.path.join(CURRENT_PROGRESS, MANIFESTS_DIR.lstrip('/'), 'hooks')
    _grant_executable_permissions(hooks_dir)

    # 执行自更新后的服务
    post_hooks = get_all_post_hooks(version_config.manifests)
    actual_manifest = os.path.join(CURRENT_PROGRESS, MANIFESTS_DIR.lstrip('/'))
    for action in post_hooks:
        if isinstance(action.command, str) and MANIFESTS_DIR in action.command:
            action.command = action.command.replace(MANIFESTS_DIR, actual_manifest)

        if isinstance(action.value, str) and MANIFESTS_DIR in action.value:
            action.value = action.value.replace(MANIFESTS_DIR, actual_manifest)

    version_config.executed_hooks = execute_hooks(post_hooks)

    # 导出配置到本地
    with open(manifest_path, "w") as f:
        json.dump(json.loads(version_config.json()), f, indent=4)

    return True,list(version_config.manifests.keys())