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


def install_nuitka():
    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 get_version() -> str:
    init_file = Path("src/orisys/__init__.py")
    if not init_file.exists():
        return "0.3.2"
    for line in init_file.read_text(encoding="utf-8").splitlines():
        if line.startswith("__version__"):
            return line.split("=", 1)[1].strip().strip('"').strip("'")
    return "0.3.2"


def find_compiled_extension(output_dir: Path) -> Path:
    candidates = []
    for suffix in EXTENSION_SUFFIXES:
        candidates.extend(output_dir.glob(f"*{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. Found: {all_files}")
    # Prefer the shortest filename (usually the package extension itself)
    return sorted(candidates, key=lambda p: len(p.name))[0]


def build():
    install_nuitka()

    version = get_version()
    src_dir = Path("src/orisys")
    output_dir = Path("build_nuitka")
    dist_dir = Path("dist_nuitka")

    if output_dir.exists():
        shutil.rmtree(output_dir)
    if dist_dir.exists():
        shutil.rmtree(dist_dir)
    output_dir.mkdir()
    dist_dir.mkdir()

    print("\n=== Compiling with Nuitka ===")
    env = os.environ.copy()
    env["PYTHONPATH"] = "src" + os.pathsep + env.get("PYTHONPATH", "")

    cmd = [
        sys.executable, "-m", "nuitka",
        "--module",
        "--include-package=orisys",
        "--output-dir=" + str(output_dir),
        "--remove-output",
        "--no-pyi-file",
        "src/orisys",
    ]
    print(f"Running: {' '.join(cmd)}")
    subprocess.check_call(cmd, env=env)

    ext_file = find_compiled_extension(output_dir)
    ext_suffix = "".join(ext_file.suffixes)
    print(f"Generated compiled extension: {ext_file}")

    print("\n=== Preparing Wheel Structure ===")
    wheel_root = dist_dir / "wheel_root"
    pkg_dir = wheel_root / "orisys"
    pkg_dir.mkdir(parents=True)

    dest_ext = pkg_dir / f"__init__{ext_suffix}"
    shutil.copy(ext_file, dest_ext)
    print(f"Copied compiled package to {dest_ext}")

    configs_src = src_dir / "configs"
    configs_dst = pkg_dir / "configs"
    configs_dst.mkdir()
    json_count = 0
    for json_file in configs_src.glob("*.json"):
        shutil.copy(json_file, configs_dst)
        json_count += 1
    print(f"Copied {json_count} JSON config files")

    examples_dst = pkg_dir / "examples"
    examples_dst.mkdir(exist_ok=True)
    (examples_dst / "__init__.py").write_text("# Examples for orisys package\n", encoding="utf-8")
    example_files = [Path("example1_basic_pipeline.py"), Path("example2_two_sensor_pipeline.py")]
    example_count = 0
    for example_file in example_files:
        if example_file.exists():
            shutil.copy(example_file, examples_dst)
            example_count += 1
            print(f"Copied example: {example_file.name}")
    print(f"Copied {example_count} example files to orisys/examples/")

    setup_py_path = dist_dir / "setup.py"
    setup_content = f"""
from setuptools import setup, Distribution

class BinaryDistribution(Distribution):
    def has_ext_modules(self):
        return True

setup(
    name="orisys",
    version="{version}",
    description="Orisys - Optical Tactile Sensor SDK (Nuitka Compiled)",
    author="Orisys",
    packages=['orisys', 'orisys.examples'],
    package_dir={{'': 'wheel_root'}},
    package_data={{
        'orisys': ['*{ext_suffix}', 'configs/*.json'],
        'orisys.examples': ['*.py'],
    }},
    include_package_data=True,
    install_requires=[
        "numpy>=1.24.4",
        "opencv-python>=4.10.0.84",
        "scipy>=1.10.1",
    ],
    extras_require={{
        "cuda12": ["cupy-cuda12x>=12.3.0"],
        "cuda13": ["cupy-cuda13x>=13.0.0"],
        "cuda": ["cupy-cuda12x>=12.3.0"],
    }},
    distclass=BinaryDistribution,
    zip_safe=False,
)
"""
    setup_py_path.write_text(setup_content, encoding="utf-8")

    print("\n=== Building Wheel ===")
    subprocess.check_call([sys.executable, "setup.py", "bdist_wheel"], cwd=dist_dir)

    wheels = list((dist_dir / "dist").glob("*.whl"))
    if wheels:
        print(f"\n[SUCCESS] Wheel created: {wheels[0]}")
    else:
        print("\n[ERROR] Wheel not found.")


if __name__ == "__main__":
    build()