from setuptools import setup, Extension
from Cython.Build import cythonize
import numpy as np
import os

SRC = "src"
np_inc = [np.get_include(), SRC]

extensions = [
    Extension('constants_inf', [f'{SRC}/constants_inf.pyx'], include_dirs=[SRC]),
    Extension('ai_availability_status', [f'{SRC}/ai_availability_status.pyx'], include_dirs=[SRC]),
    Extension('annotation', [f'{SRC}/annotation.pyx'], include_dirs=[SRC]),
    Extension('ai_config', [f'{SRC}/ai_config.pyx'], include_dirs=[SRC]),
    Extension('loader_http_client', [f'{SRC}/loader_http_client.pyx'], include_dirs=[SRC]),
    Extension('engines.inference_engine', [f'{SRC}/engines/inference_engine.pyx'], include_dirs=np_inc),
    Extension('engines.engine_factory', [f'{SRC}/engines/engine_factory.pyx'], include_dirs=[SRC]),
    Extension('engines.onnx_engine', [f'{SRC}/engines/onnx_engine.pyx'], include_dirs=np_inc),
    Extension('engines.coreml_engine', [f'{SRC}/engines/coreml_engine.pyx'], include_dirs=np_inc),
    Extension('inference', [f'{SRC}/inference.pyx'], include_dirs=np_inc),
]

build_tensorrt = os.environ.get("BUILD_TENSORRT_EXTENSIONS", "").lower() in ("1", "true", "yes")

if not build_tensorrt:
    try:
        import tensorrt  # pyright: ignore[reportMissingImports]
        build_tensorrt = True
    except ImportError:
        build_tensorrt = False

if build_tensorrt:
    extensions.append(
        Extension('engines.tensorrt_engine', [f'{SRC}/engines/tensorrt_engine.pyx'], include_dirs=np_inc)
    )
    extensions.append(
        Extension('engines.jetson_tensorrt_engine', [f'{SRC}/engines/jetson_tensorrt_engine.pyx'], include_dirs=np_inc)
    )

setup(
    name="azaion.detections",
    package_dir={"": SRC},
    packages=["engines"],
    ext_modules=cythonize(
        extensions,
        include_path=[SRC],
        compiler_directives={
            "language_level": 3,
            "emit_code_comments": False,
            "binding": True,
            'boundscheck': False,
            'wraparound': False,
        }
    ),
    zip_safe=False
)