import re

import yaml
import constants
from azaion_api import Api, ApiCredentials
from cdn_manager import CDNManager, CDNCredentials
from inference.inference import Inference
from inference.onnx_engine import OnnxEngine
from inference.tensorrt_engine import TensorRTEngine
from utils import Dotdict

if __name__ == "__main__":
    # Inference(OnnxEngine('azaion-2025-03-10.onnx', batch_size=4),
    #            confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4')
    # detection for the first 200sec of video:
    # onnxInference: 81 sec, 6.3Gb VRAM
    # tensorrt: 54 sec, 3.7Gb VRAM

    # Inference(TensorRTEngine('azaion-2025-03-10_int8.engine', batch_size=16),
    #          confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4')
    # INT8 for 200sec: 54 sec 3.7Gb

    # Inference(TensorRTEngine('azaion-2025-03-10_batch8.engine', batch_size=8),
    #           confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4')

    with open(constants.CONFIG_FILE, "r") as f:
        config_dict = yaml.safe_load(f)
        d_config = Dotdict(config_dict)
        cdn_c = Dotdict(d_config.cdn)
        api_c = Dotdict(d_config.api)

    cdn_manager = CDNManager(CDNCredentials(cdn_c.host,
                                            cdn_c.downloader_access_key, cdn_c.downloader_access_secret,
                                            cdn_c.uploader_access_key, cdn_c.uploader_access_secret))
    cdn_manager.download(cdn_c.bucket, constants.AI_TENSOR_MODEL_FILE_BIG)
    cdn_manager.download(cdn_c.bucket, constants.AI_ONNX_MODEL_FILE_BIG)

    api_client = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))

    tensor_model_bytes = api_client.load_resource(constants.AI_TENSOR_MODEL_FILE_BIG, constants.AI_TENSOR_MODEL_FILE_SMALL)
    onxx_model_bytes = api_client.load_resource(constants.AI_ONNX_MODEL_FILE_BIG, constants.AI_ONNX_MODEL_FILE_SMALL)


    # Inference(OnnxEngine(onxx_model_bytes, batch_size=4),
    #           confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4')

    Inference(TensorRTEngine(tensor_model_bytes, batch_size=4),
               confidence_threshold=0.5, iou_threshold=0.3).process('tests/ForAI_test.mp4')