import os
import sys
import subprocess
import shutil
from pathlib import Path

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

def build():
    install_nuitka()
    
    version = "0.3.0"
    src_dir = Path("src/orisys")
    output_dir = Path("build_nuitka")
    dist_dir = Path("dist_nuitka")
    
    # Clean previous builds
    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 ===")
    # Set PYTHONPATH to include src so nuitka finds the package
    env = os.environ.copy()
    env["PYTHONPATH"] = "src" + os.pathsep + env.get("PYTHONPATH", "")
    
    # Nuitka arguments
    # --module: compile as module/package
    # --include-package=orisys: include the whole package
    # --output-dir: where to put result
    # src/orisys: the target directory
    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)
    
    # Find the generated .so file
    # It should be something like orisys.cpython-310-x86_64-linux-gnu.so
    so_files = list(output_dir.glob("*.so"))
    if not so_files:
        raise RuntimeError("Nuitka did not produce a .so file!")
    
    so_file = so_files[0]
    print(f"Generated compiled extension: {so_file}")
    
    print("\n=== Preparing Wheel Structure ===")
    # We want the structure:
    # orisys/
    #   __init__.so  <-- renamed from orisys.so
    #   configs/
    #     *.json
    
    wheel_root = dist_dir / "wheel_root"
    pkg_dir = wheel_root / "orisys"
    pkg_dir.mkdir(parents=True)
    
    # Move and rename so
    dest_so = pkg_dir / "__init__.so"
    shutil.copy(so_file, dest_so)
    print(f"Copied to {dest_so}")
    
    # Copy JSON configs
    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")
    
    # Create setup.py for packaging
    setup_py_path = dist_dir / "setup.py"
    setup_content = f"""
from setuptools import setup, Distribution, find_packages

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'],
    package_dir={{'': 'wheel_root'}},
    package_data={{
        'orisys': ['*.so', 'configs/*.json'],
    }},
    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"],  # For backward compatibility, defaults to CUDA 12.x
    }},
    distclass=BinaryDistribution,
    zip_safe=False,
)
"""
    with open(setup_py_path, "w") as f:
        f.write(setup_content)
    
    print("\n=== Building Wheel ===")
    subprocess.check_call([sys.executable, "setup.py", "bdist_wheel"], cwd=dist_dir)
    
    # Find the wheel
    dist_output = dist_dir / "dist"
    wheels = list(dist_output.glob("*.whl"))
    if wheels:
        print(f"\n[SUCCESS] Wheel created: {wheels[0]}")
    else:
        print("\n[ERROR] Wheel not found.")

if __name__ == "__main__":
    build()