import os
import tempfile
from engines.tensorrt_engine cimport TensorRTEngine
from loader_http_client cimport LoaderHttpClient, LoadResult


cdef class JetsonTensorRTEngine(TensorRTEngine):
    @staticmethod
    def convert_from_source(bytes onnx_model, LoaderHttpClient loader_client, str models_dir):
        cdef str calib_cache_path
        calib_cache_path = JetsonTensorRTEngine._download_calib_cache(loader_client, models_dir)
        try:
            return TensorRTEngine.convert_from_source(onnx_model, calib_cache_path)
        finally:
            if calib_cache_path is not None:
                try:
                    os.unlink(calib_cache_path)
                except Exception:
                    pass

    @staticmethod
    def _download_calib_cache(LoaderHttpClient loader_client, str models_dir):
        cdef LoadResult res
        import constants_inf
        try:
            res = loader_client.load_big_small_resource(
                constants_inf.INT8_CALIB_CACHE_FILE, models_dir
            )
            if res.err is not None:
                constants_inf.log(f"INT8 calibration cache not available: {res.err}")
                return None
            fd, path = tempfile.mkstemp(suffix=".cache")
            with os.fdopen(fd, "wb") as f:
                f.write(res.data)
            constants_inf.log("INT8 calibration cache downloaded")
            return path
        except Exception as e:
            constants_inf.log(f"INT8 calibration cache download failed: {str(e)}")
            return None