"""Nuitka compilation of the orisys package."""

from __future__ import annotations

import shutil
import subprocess
import sys
from importlib.machinery import EXTENSION_SUFFIXES
from pathlib import Path

from paths import repo_path

NUITKA_INCLUDE_PACKAGES = (
    "orisys",
    "orisys.finger",
    "orisys.unroll",
    "orisys.hhd",
    "orisys.hhd.solvers",
)

NUITKA_NOFOLLOW_IMPORTS = (
    # Keep as runtime deps (shipped via wheel install_requires / extras).
    "numpy",
    "scipy",
    "cv2",
    "numba",
    "orisys_blend",
    "cupy",
    "cupyx",
    "pyopencl",
    "open3d",
    "PyQt5",
    "pyqtgraph",
)


def install_nuitka() -> None:
    print("Checking Nuitka installation...")
    try:
        import nuitka  # noqa: F401
        print("Nuitka is installed.")
    except ImportError:
        print("Installing Nuitka...")
        subprocess.check_call(
            [sys.executable, "-m", "pip", "install", "nuitka", "ordered-set", "zstandard"]
        )


def find_compiled_extension(output_dir: Path, module_name: str) -> Path:
    candidates = []
    for suffix in EXTENSION_SUFFIXES:
        candidates.extend(output_dir.glob(f"{module_name}*{suffix}"))
    if not candidates:
        all_files = ", ".join(p.name for p in output_dir.iterdir()) if output_dir.exists() else "(none)"
        raise RuntimeError(
            f"Nuitka did not produce a compiled extension for {module_name}. Found: {all_files}"
        )
    return sorted(candidates, key=lambda p: len(p.name))[0]


def compile_orisys(build_root: Path, env: dict[str, str]) -> Path:
    module_output_dir = build_root / "orisys"
    if module_output_dir.exists():
        shutil.rmtree(module_output_dir)
    module_output_dir.mkdir(parents=True)

    cmd = [
        sys.executable,
        "-m",
        "nuitka",
        "--module",
        "--output-dir=" + str(module_output_dir),
        "--remove-output",
        "--no-pyi-file",
        repo_path("src/orisys").as_posix(),
    ]
    for package_name in NUITKA_INCLUDE_PACKAGES:
        cmd.append(f"--include-package={package_name}")
    for import_name in NUITKA_NOFOLLOW_IMPORTS:
        cmd.append(f"--nofollow-import-to={import_name}")

    print(f"Running: {' '.join(cmd)}")
    subprocess.check_call(cmd, env=env)

    ext_file = find_compiled_extension(module_output_dir, "orisys")
    print(f"Generated compiled extension for orisys: {ext_file}")
    return ext_file


def copy_compiled_orisys(pkg_dir: Path, ext_file: Path) -> str:
    ext_suffix = "".join(ext_file.suffixes)
    dest_ext = pkg_dir / f"__init__{ext_suffix}"
    shutil.copy(ext_file, dest_ext)
    print(f"Copied compiled module orisys to {dest_ext}")
    return ext_suffix