#!/usr/bin/python3
import sys
import os
import subprocess
import time
import shlex
import re
import logging
import argparse
import hashlib
import json
import shutil
import fcntl
from contextlib import contextmanager
logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
STATE_FILE_DIR = '/var/run/alinas'
BACKUP_DIR = '/var/log/aliyun/alinas'
STATE_SIGN = 'sign'
UNAS_LOCK_STATEFILE_TIMEOUT_SECONDS = 20

###############################################################
# State file operations
###############################################################


def log_and_raise(message, error_type=ValueError):
    logging.error(message)
    raise error_type(message)

# 计算所有参数的 hash 校验值


def sign_state(state):
    keys = sorted(state.keys())
    md5 = hashlib.md5()

    for key in keys:
        val = state[key]
        if type(val) is list:
            for x in val:
                md5.update(str(x).encode('utf-8'))
        else:
            md5.update(str(val).encode('utf-8'))
    return md5.hexdigest()

# 校验 hash 值正确性，并将校验码从 state 中移除


def is_integral(state):
    saved_sign = state.pop(STATE_SIGN, '')
    computed_sign = sign_state(state)
    return saved_sign == computed_sign

# 读取 state 文件并解析，返回 hash 校验码和 dict(state)


def load_state_file(state_file, state_file_dir=STATE_FILE_DIR):
    state_file_path = os.path.join(state_file_dir, state_file)
    try:
        with open(state_file_path, 'r', encoding='utf-8') as f:
            try:
                state = json.load(f)
                result_sign = state.get(STATE_SIGN, '')
                if is_integral(state):
                    result_state = state
                    return result_sign, result_state
                else:
                    logging.error(
                        f"State file for {state_file_path} is modified or corrupted")
            except ValueError:
                log_and_raise(f"Unable to parse json in {state_file_path}")
    except IOError as e:
        log_and_raise(
            f"Fail to read state file {state_file_path}, err msg: {e}", IOError)

# 锁定 state 文件


@contextmanager
def lock_state_file(lock_type, state_file_dir=STATE_FILE_DIR, timeout=UNAS_LOCK_STATEFILE_TIMEOUT_SECONDS):
    if not os.path.exists(state_file_dir):
       os.makedirs(state_file_dir)
    path = os.path.join(state_file_dir, lock_type)
    if not os.path.isfile(path):
        os.mknod(path)
    if timeout == 0:
        try:
            fd = os.open(path, os.O_CREAT | os.O_RDWR)
            fcntl.lockf(fd, fcntl.LOCK_EX)
            yield
        finally:
            os.close(fd)
    else:
        locked = False
        start = time.time()
        while not locked:
            passed = time.time() - start
            if passed > timeout:
                raise Exception("lock_state_file timeout, path:%s", path)
            try:
                fd = os.open(path, os.O_CREAT | os.O_RDWR)
                fcntl.lockf(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
                locked = True
                yield
            except OSError as e:
                logging.error("Fail to open and lock file, sleep 0.3s")
                time.sleep(0.3)
            finally:
                os.close(fd)

# 备份 state 文件


def backup_state_file(state_file, sign, state_file_dir=STATE_FILE_DIR, backup_dir=BACKUP_DIR):
    state_file_path = os.path.join(state_file_dir, state_file)
    if not os.path.isdir(state_file_dir) or not os.path.isfile(state_file_path):
        log_and_raise("State file directory or state file is not exist")

    backup_file_dir = os.path.join(
        backup_dir, state_file.replace("eac-", "efc-"))
    if not os.path.isdir(backup_file_dir):
        logging.warning(
            f"Backup state file directory is not exist: {backup_file_dir}")
        try:
            os.makedirs(backup_file_dir)
            logging.info(
                f"Create backup state file directory completed: {backup_file_dir}")
        except Exception as e:
            log_and_raise(
                f"Fail to create backup state file directory, err msg:{str(e)}", IOError)

    timestamp_us = int(round(time.time() * 1000 * 1000))
    backup_file_path = os.path.join(
        backup_file_dir, f"{state_file}-{timestamp_us}-{sign}.backup")
    try:
        shutil.copy(state_file_path, backup_file_path)
    except Exception as e:
        log_and_raise(f"Fail to backup state file, err msg:{str(e)}", IOError)
    return backup_file_path

# 使用新的 state 覆盖 state 文件, 在覆盖失败时回退 state 文件，并在临时文件未正确删除时清除


def rewrite_state_file(state, state_file, backup_file_path, state_file_dir=STATE_FILE_DIR):
    tmp_state_file = os.path.join(state_file_dir, '~%s' % state_file)
    state_file_path = os.path.join(state_file_dir, state_file)
    try:
        with open(tmp_state_file, 'w') as f:
            signed_state = dict(state)
            signed_state[STATE_SIGN] = sign_state(state)
            json.dump(signed_state, f)
        os.rename(tmp_state_file, state_file_path)
        # logging.info(f"Content of updated {state_file_path}: {signed_state}")
    except Exception as e:
        logging.info(
            f"Rewrite state file failed (err msg:{str(e)}), try to roll back state file")
        try:
            shutil.copy(backup_file_path, state_file_path)
            logging.info(
                f"Success roll back state file with backup file: {backup_file_path}")
        except Exception as e:
            logging.error(
                f"Fail to roll back state file, please handle roll back manually, err msg: {str(e)}, backup file: {backup_file_path}")
        raise
    finally:
        try:
            if os.path.exists(tmp_state_file):
                os.unlink(tmp_state_file)
        except Exception as e:
            logging.error(
                f"Fail to clear tmp state file ({tmp_state_file}) in rewrite, err msg:{str(e)}")

# 使用备份文件覆盖 state 文件


def recover_state_file(state, state_file_path, backup_file_path):
    uuid_lock_name = state.get('mountkey')
    with lock_state_file(uuid_lock_name):
        if not os.path.exists(backup_file_path):
            log_and_raise(
                f"Backup state file is not exist: {backup_file_path}")
        try:
            shutil.copy(backup_file_path, state_file_path)
            logging.info(
                f"Success roll back {state_file_path} with backup file: {backup_file_path}")
        except Exception as e:
            log_and_raise(
                f"Fail to rewrite state file, manually roll back state file, err msg:{str(e)}")

# 只解析挂载参数，其他挂载选项不改变


def parse_mountcmd(mountcmd):
    if not isinstance(mountcmd, str):
        log_and_raise(f"Invalid input type of mountcmd")

    parts = mountcmd.split()
    vec1 = []       # 挂载命令非 设置参数 部分内容
    dict2 = {}      # 设置参数

    for part in parts:
        if part.startswith('--'):
            key_value = part[2:]
            if '=' in key_value:
                key, value = key_value.split('=', 1)
                dict2[key] = value
            else:
                log_and_raise(
                    f"Invalid format for -- option: missing '=' for {key_value}")
        elif part.startswith('-tier') or part.startswith('tier'):
            log_and_raise(f"Invalid format, should be --tier_xxx")
        else:
            vec1.append(part)
    str1 = ' '.join(vec1)
    return str1, dict2

# 解析更新参数


def parse_updatecmd(updatecmd):
    if not isinstance(updatecmd, str):
        log_and_raise(f"Invalid input type of updatecmd")
    res = {}
    key_values = updatecmd.rstrip(',')
    for key_value in key_values.split(','):
        if (not key_value.startswith('g_')) or ('=' not in key_value):
            log_and_raise(
                f"Invalid format for -o option: '{key_value}' should start with 'g_' and has '='")
        key_value = key_value[2:]
        key, value = key_value.split('=')
        res[key] = value
    if updatecmd.count('=') != len(res):
        logging.warning(
            f"Invalid format, some params may not be parsed, args: {res}")
    return res

# 构造新的 mount 命令


def construct_mountcmd(rest, param_dict):
    if not isinstance(rest, str) or not isinstance(param_dict, dict):
        log_and_raise(f"Invalid input types of cmd or params")

    param_dict_items = []
    for key, value in param_dict.items():
        if value is None:
            log_and_raise(
                f"Invalid format for -- option: missing value for key {key}")
        param_dict_items.append(f"--{key}={value}")
    param_dict_str = " ".join(param_dict_items)
    mountcmd = f"{rest} {param_dict_str}"
    return mountcmd.strip()


# 检查新参数是否合法
valid_keys = {
    "",
}


def check_tier_params(param_dict, update_params_dict):
    if not isinstance(param_dict, dict) or not isinstance(update_params_dict, dict):
        log_and_raise(f"Invalid input type of param_dict")

    invalid_keys = [key for key in update_params_dict.keys()
                    if key not in valid_keys]
    if invalid_keys:
        log_and_raise(
            f"Invalid keys found: {invalid_keys}. Please update valid keys in this script if new tier params are added to EFC")

    if 'tier_DadiDiskCacheCapacityMB' in param_dict and int(param_dict.get('tier_DadiDiskCacheCapacityMB')) > 0:
        if 'tier_DadiDiskCachePath' not in param_dict:
            log_and_raise(
                "tier_DadiDiskCachePath is required when tier_DadiDiskCacheCapacityMB is greater than zero")
        dadi_disk_paths = [path.strip() for path in param_dict.get(
            'tier_DadiDiskCachePath').split(':') if path.strip()]
        for path in dadi_disk_paths:
            if not os.path.isdir(path):
                log_and_raise(f"Directory {path} does not exist")

# 获取需要更新的参数


def get_update_params(params, new_params):
    res_params = {}
    for key, value in new_params.items():
        if key not in params or params[key] != value:
            res_params[key] = value
    return res_params

# 更新 state 文件


def update_state_file(state_file, update_cmd, state_file_dir=STATE_FILE_DIR, backup_dir=BACKUP_DIR):
    if not os.path.isdir(state_file_dir) or not os.path.isfile(os.path.join(state_file_dir, state_file)):
        log_and_raise("state_file_dir or state_file is not exist")

    try:
        sign, state = load_state_file(state_file, state_file_dir)
        uuid_lock_name = state.get('mountkey')
        with lock_state_file(uuid_lock_name):
            sign, state = load_state_file(state_file, state_file_dir)

            mountcmd = state.get('mountcmd')
            if mountcmd is None:
                log_and_raise("mountcmd is not exist in state file")

            try:
                rest, params = parse_mountcmd(mountcmd)
                logging.info(f"Parse original mount command completed")
            except Exception as e:
                raise

            try:
                new_params = parse_updatecmd(update_cmd)
                logging.info(f"Parse new params completed")
            except Exception as e:
                raise

            update_params = get_update_params(params, new_params)
            if update_params:
                params.update(update_params)
                logging.info(f"New params to be updated: {update_params}")
            else:
                logging.info(f"No params need to be updated")
                return

            # not check params valid here
            # try:
            #     check_tier_params(params, update_params)
            #     logging.info(f"Check new params, all params are valid")
            # except Exception as e:
            #     log_and_raise(f"Check new params, some params are invalid")

            try:
                state['mountcmd'] = construct_mountcmd(rest, params)
                logging.info(f"Construct new mount command completed")
            except Exception as e:
                raise

            try:
                backup_file_path = backup_state_file(
                    state_file, sign, state_file_dir, backup_dir)
                logging.info(
                    f"Backup state file completed: {backup_file_path}")
            except Exception as e:
                raise

            try:
                rewrite_state_file(state, state_file,
                                   backup_file_path, state_file_dir)
                logging.info(
                    f"Update {os.path.join(state_file_dir, state_file)} completed")
            except Exception as e:
                raise

    except Exception as e:
        log_and_raise(f"Fail to update state file")

# 回滚 state 文件


def rollback_state_file(state_file, state_file_dir=STATE_FILE_DIR, backup_dir=BACKUP_DIR, backup_state_file_path=None):
    state_file_path = os.path.join(state_file_dir, state_file)
    if backup_state_file_path is not None:
        if not os.path.isfile(backup_state_file_path):
            log_and_raise("Backup state file is not exist")
        try:
            # use the provided backup file to roll back the state file
            file_dir, file_name = os.path.split(backup_state_file_path)
            sign, state = load_state_file(file_name, file_dir)
            recover_state_file(state, state_file_path, backup_state_file_path)
            logging.info(f"Roll back state file sucess")
            return
        except Exception as e:
            log_and_raise(
                f"Check backup file({backup_state_file_path}) failed, err msg: {str(e)}")

    # use the default backup file to roll back the state file
    backup_file_dir = os.path.join(
        backup_dir, state_file.replace("eac-", "efc-"))
    if not os.path.isdir(backup_file_dir):
        log_and_raise(
            f"Backup state file directory is not exist: {backup_file_dir}")

    pattern = rf"{state_file}-(\d+)-([a-fA-F0-9]+)\.backup"
    files = []
    for filename in os.listdir(backup_file_dir):
        match = re.match(pattern, filename)
        if match:
            timestamp = int(match.group(1))
            hash_sign = match.group(2)
            files.append(
                (os.path.join(backup_file_dir, filename), timestamp, hash_sign))
    if not files:
        log_and_raise(f"Backup state file is not exist: {backup_file_dir}")

    latest_backup_file = max(files, key=lambda x: x[1])
    backup_file_path, timestamp, hash_sign = latest_backup_file
    try:
        file_dir, file_name = os.path.split(backup_file_path)
        sign, state = load_state_file(file_name, file_dir)
        if sign != hash_sign:
            log_and_raise(
                f"Backup state file is corrupted: {backup_file_path}")
        recover_state_file(state, state_file_path, backup_file_path)
        logging.info(f"Roll back state file sucess")
    except Exception as e:
        log_and_raise(
            f"Check backup file({backup_file_path}) failed, err msg: {str(e)}")

###############################################################
# System operations
###############################################################
# 检查 watchdog 进程是否存活


def check_watchdog():
    try:
        ps_result = subprocess.run(
            ["ps", "aux"], stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    except (OSError, subprocess.CalledProcessError) as e:
        log_and_raise(f"Failed to run command: {e}")
    try:
        grep_result = subprocess.run(["grep", "aliyun-alinas-mount-watchdog"], input=ps_result.stdout,
                                     stdout=subprocess.PIPE, stderr=subprocess.PIPE, universal_newlines=True)
    except (OSError, subprocess.CalledProcessError) as e:
        log_and_raise(f"Failed to run command: {e}")
    if grep_result.returncode == 0 and grep_result.stdout.strip():
        logging.info("Watchdog is alive")
    else:
        log_and_raise("Watchdog is not alive, stop")

# 通过挂载路径获取挂载点 uuid


def get_mountpoints_uuid(mount_str):
    # /proc/self/mounts 与 /var/run/alinas/nonfuse_mounts 格式一致，均按此解析
    mounts_files = ["/proc/self/mounts", "/var/run/alinas/nonfuse_mounts"]
    filtered_lines = []
    for mounts_file in mounts_files:
        try:
            with open(mounts_file, 'r', encoding='utf-8') as f:
                filtered_lines.extend(
                    line for line in f.read().splitlines() if "aliyuncs.com" in line)
        except FileNotFoundError:
            continue
        except IOError as e:
            log_and_raise(
                f"Failed to read {mounts_file}, err msg: {e}", IOError)

    mount_ids = [id.strip() for id in mount_str.split(",")]
    mountpoints = []
    for line in filtered_lines:
        fields = line.split()
        if len(fields) < 2:
            continue
        mount_path = fields[1]
        for mount_id in mount_ids:
            if mount_path == mount_id.rstrip('/'):
                mountpoint = fields[0].split(":", 1)[0].strip()
                if mountpoint not in mountpoints:
                    mountpoints.append(mountpoint)
                break
    if not any(mountpoints):
        log_and_raise(f"No filesystem mounts on the path(s): {mount_ids}")
    return mountpoints

# 传入 efc-mount_uuid 来查找进程获取 pid


def get_pid(keyword):
    efc_bin = "aliyun-alinas-efc"
    if not isinstance(keyword, str) or not keyword.strip():
        log_and_raise(f"Invalid input types of keyword")
    try:
        result = subprocess.run(["ps", "aux"], stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE, universal_newlines=True)
    except (OSError, subprocess.CalledProcessError) as e:
        log_and_raise(f"Failed to run command: {e}")
    process_ids = []
    for line in result.stdout.splitlines():
        parts = line.split()
        if efc_bin in line and keyword in line and parts[1] not in process_ids:
            process_ids.append(parts[1])
    if len(process_ids) == 1:
        # logging.info(f"Found one process id {process_ids[0]} for keyword:{keyword}")
        return process_ids[0]
    else:
        log_and_raise(
            f"Found {len(process_ids)} process ids for keyword:{keyword}")

# 通过 efc-uuid 找到对应的 efc 进程并关闭


def kill_process_by_keyword(keyword):
    try:
        pid = get_pid(keyword)
    except:
        log_and_raise(f"Failed to get pid for keyword:{keyword}")
    try:
        result = subprocess.run(["kill", "-9", pid], stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE, universal_newlines=True)
        return pid
    except (OSError, subprocess.CalledProcessError) as e:
        log_and_raise(f"Failed to run command: {e}")


def monitor_efc_restart(process_keyword, old_pid, mountpoint, wait_times=10):
    while wait_times > 0:
        wait_times -= 1
        time.sleep(1)
        try:
            new_pid = get_pid(process_keyword)  # 获取新进程的 PID
        except Exception as e:
            logging.info(f"Wait for EFC restart: {e}")
            continue

        if old_pid == new_pid:
            log_and_raise("Kill process failed")
        else:
            logging.info(
                f"Restart EFC with new params completed, process id: {new_pid}, mount uuid: {mountpoint}")
            break
    else:
        # 如果循环结束仍未检测到新进程，抛出异常
        log_and_raise("EFC restart monitoring timed out")

###############################################################
# Interface
###############################################################
# 检查挂载点状态并执行 EFC 更新操作


def update_efc(mount_points, is_update_efc, update_cmd, state_file_dir=STATE_FILE_DIR, backup_dir=BACKUP_DIR, backup_state_file_path=None):
    if not isinstance(mount_points, str) or not isinstance(update_cmd, str):
        log_and_raise(f"Invalid input types of mount info or update cmd")

    try:
        check_watchdog()
    except Exception as e:
        log_and_raise(f"Check watchdog failed, err msg: {str(e)}")

    try:
        mountpoints = get_mountpoints_uuid(mount_points)
        logging.info(
            f"Update state file start, the mountpoint(s) to be updated: {mountpoints}")
    except Exception as e:
        log_and_raise(f"Get mount uuid failed, err msg: {str(e)}")

    for mountpoint in mountpoints:
        state_file_name = f"eac-{mountpoint}"
        process_keyword = f"efc-{mountpoint}"
        try:
            if is_update_efc == True:
                update_state_file(state_file_name, update_cmd,
                                  state_file_dir, backup_dir)
                logging.info(
                    f"Update state file done, mount uuid: {mountpoint}")
            else:
                rollback_state_file(
                    state_file_name, state_file_dir, backup_dir, backup_state_file_path)
                logging.info(
                    f"Roll back state file done, mount uuid: {mountpoint}")
        except Exception as e:
            log_and_raise(
                f"Update/Roll back state file failed, err msg:{str(e)}")

        try:
            old_pid = kill_process_by_keyword(process_keyword)
            logging.info(
                f"Kill process {old_pid} done, mount uuid: {mountpoint}")
        except Exception as e:
            log_and_raise(f"Kill process {old_pid} failed, err msg:{str(e)}")

        try:
            monitor_efc_restart(process_keyword, old_pid, mountpoint)
        except Exception as e:
            log_and_raise(
                f"Restart EFC with new command failed, err msg:{str(e)}")


# 主函数
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Update EFC with new params or Roll back EFC with old state file")
    subparsers = parser.add_subparsers(dest="command", help="subcommand help")
    # 添加 update 子命令
    update_parser = subparsers.add_parser(
        "update", help="update EFC with new params, usage: python3 efc_runtime_ops.py update <mount_path> <update_cmd>")
    update_parser.add_argument(
        "mount_path", type=str, help="mount path(s), e.g. /mnt or /mnt1,/mnt2")
    update_parser.add_argument("-o", action="append", dest="options",
                               help="params to update, e.g. -o g_key1=value1,g_key2=value2 -o g_key3=value3")
    update_parser.set_defaults(func=lambda args: update_efc(
        args.mount_path, True, ",".join(args.options) if args.options else ""))
    # 添加 rollback 子命令
    rollback_parser = subparsers.add_parser(
        "rollback", help="roll back EFC with backup state file, usage: python3 script.py rollback <mount_path> [backup_file]")
    rollback_parser.add_argument(
        "mount_path", type=str, help="mount path(s), e.g. /mnt or /mnt1,/mnt2")
    rollback_parser.add_argument("backup_file", type=str, nargs='?', default=None,
                                 help="path for the backup state file, EFC will search default backup file if this argument is not provide")
    rollback_parser.set_defaults(func=lambda args: update_efc(
        args.mount_path, False, "", STATE_FILE_DIR, BACKUP_DIR, args.backup_file))

    args = parser.parse_args()
    if len(sys.argv) < 3:
        parser.print_help()
        sys.exit(1)
    if hasattr(args, 'func'):
        args.func(args)
    else:
        parser.print_help()
        sys.exit(1)