#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import os
import sys
import subprocess
import argparse
import shutil
import stat
from datetime import datetime
from pathlib import Path

PACKAGE_NAME = "affinity-sched"
SCRIPT_DIR = Path(__file__).parent.resolve()


def _is_root_user() -> bool:
    try:
        return os.geteuid() == 0
    except AttributeError:
        return False


def _log(level: str, message: str, *, stream=None):
    ts = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    print(f"[msboost][{ts}][{level}]: {message}", file=stream or sys.stdout)


def log_info(message: str):
    _log("INFO", message)


def log_warning(message: str):
    _log("WARNING", message, stream=sys.stderr)


def log_error(message: str):
    _log("ERROR", message, stream=sys.stderr)


def _install_for_all_mode(mode: int) -> int:
    perm = f"{mode & 0o777:o}".zfill(3)[-3:]
    return int(perm[:2] + perm[1], 8)


def _apply_install_for_all(path: Path):
    if not path.exists():
        return

    if path.is_file() or path.is_symlink():
        path.chmod(_install_for_all_mode(path.stat().st_mode))
        return

    for root, dirs, files in os.walk(path):
        root_path = Path(root)
        for name in dirs:
            entry = root_path / name
            entry.chmod(_install_for_all_mode(entry.stat().st_mode))
        for name in files:
            entry = root_path / name
            entry.chmod(_install_for_all_mode(entry.stat().st_mode))
    path.chmod(_install_for_all_mode(path.stat().st_mode))


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--install-path", required=True, help="Specify the installation path")
    parser.add_argument(
        "--install-for-all",
        action="store_true",
        help="Grant group-equivalent permissions to other users",
    )
    args = parser.parse_args()

    log_info(f"Starting installation of {PACKAGE_NAME}")

    install_path = Path(args.install_path)

    whl_pattern = f"{PACKAGE_NAME.replace('-', '_')}-*.whl"
    whl_list = sorted(SCRIPT_DIR.glob(whl_pattern))

    if not whl_list:
        log_error(f"Installation package not found: {whl_pattern}")
        sys.exit(1)

    if len(whl_list) > 1:
        log_error(
            "Multiple installation packages found in the directory. "
            "Please clean them up and try again."
        )
        sys.exit(1)

    whl_file = whl_list[0]
    log_info(f"Installing package: {whl_file.name}")

    version_file = SCRIPT_DIR / "VERSION"
    if version_file.exists():
        version = version_file.read_text().strip()
    else:
        version = whl_file.name.split("-")[1]

    target_dir = install_path / "python" / "site-packages"
    install_for_all = args.install_for_all
    if _is_root_user() or install_for_all:
        old_umask = os.umask(0o022)
    else:
        old_umask = os.umask(0o027)

    try:
        target_dir.mkdir(parents=True, exist_ok=True)

        pip_cmd = [
            "python3", "-m", "pip", "install",
            "-qq",
            "--no-color",
            "--disable-pip-version-check",
            "--force-reinstall",
            "--no-index",
            "--no-deps",
            "--target", str(target_dir),
            str(whl_file),
        ]

        try:
            subprocess.run(
                pip_cmd,
                check=True,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
            )
        except subprocess.CalledProcessError as exc:
            if exc.stderr:
                pip_error = exc.stderr.decode(errors="replace").strip()
                if pip_error:
                    log_error(pip_error)
            log_error("Installation failed")
            sys.exit(1)

        info_dir = install_path / "share" / "info" / "msboost"
        info_dir.mkdir(parents=True, exist_ok=True)

        uninstall_sh_src = SCRIPT_DIR / "uninstall.sh"
        if uninstall_sh_src.exists():
            shutil.copy2(uninstall_sh_src, info_dir / "uninstall.sh")
            (info_dir / "uninstall.sh").chmod((info_dir / "uninstall.sh").stat().st_mode | 0o755)

        version_info = info_dir / "version.info"
        version_info.write_text(
            f"[PACKAGE]\nName=mindstudio-boost\nVersion={version}\n"
        )
    finally:
        os.umask(old_umask)

    cann_uninstall = install_path / "cann_uninstall.sh"
    entry = 'uninstall_package "share/info/msboost"'

    if cann_uninstall.exists():
        content = cann_uninstall.read_text()
        if entry not in content:
            lines = content.splitlines(keepends=True)
            new_lines = []
            for line in lines:
                if line.strip() == "exit ${TOTAL_RET}":
                    new_lines.append(entry + "\n")
                new_lines.append(line)

            original_mode = cann_uninstall.stat().st_mode
            need_restore = not os.access(cann_uninstall, os.W_OK)
            if need_restore:
                cann_uninstall.chmod(original_mode | stat.S_IWUSR)
            try:
                cann_uninstall.write_text("".join(new_lines))
            finally:
                if need_restore:
                    try:
                        cann_uninstall.chmod(original_mode)
                    except PermissionError:
                        pass

    if install_for_all:
        package_dir = target_dir / PACKAGE_NAME.replace("-", "_")
        for dist_info in sorted(target_dir.glob(f"{PACKAGE_NAME.replace('-', '_')}-*.dist-info")):
            _apply_install_for_all(dist_info)
        _apply_install_for_all(package_dir)
        _apply_install_for_all(info_dir)

    log_info("Installation completed successfully")


if __name__ == "__main__":
    main()