import os
import platform
import sys

from loguru import logger
from engines.engine_factory import (
    EngineFactory,
    OnnxEngineFactory,
    CoreMLEngineFactory,
    TensorRTEngineFactory,
    JetsonTensorRTEngineFactory,
)


def _check_tensor_gpu_index():
    try:
        import pynvml
        pynvml.nvmlInit()
        device_count = pynvml.nvmlDeviceGetCount()
        if device_count == 0:
            return -1
        for i in range(device_count):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
            if major > 6 or (major == 6 and minor >= 1):
                return i
        return -1
    except Exception:
        return -1
    finally:
        try:
            import pynvml
            pynvml.nvmlShutdown()
        except Exception:
            pass


def _is_apple_silicon():
    if sys.platform != "darwin" or platform.machine() != "arm64":
        return False
    try:
        import coremltools
        return True
    except ImportError:
        return False


def _is_jetson():
    return (
        platform.machine() == "aarch64"
        and tensor_gpu_index > -1
        and os.path.isfile("/etc/nv_tegra_release")
    )


tensor_gpu_index = _check_tensor_gpu_index()


def _create_engine_factory() -> EngineFactory:
    if _is_jetson():
        logger.info("Engine factory: JetsonTensorRTEngineFactory")
        return JetsonTensorRTEngineFactory()
    if tensor_gpu_index > -1:
        logger.info("Engine factory: TensorRTEngineFactory")
        return TensorRTEngineFactory()
    if _is_apple_silicon():
        logger.info("Engine factory: CoreMLEngineFactory")
        return CoreMLEngineFactory()
    logger.info("Engine factory: OnnxEngineFactory")
    return OnnxEngineFactory()


engine_factory = _create_engine_factory()
