import sys
import os
from typing import Dict, Any

# Ensure nxva is in path
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../")))

from nxva.yolo.yolo import YOLO
from nxva.yolo.pipeline import load_config
from nxva.sort.reid_multibackend import ReIDDetectMultiBackend

class ModelRunner:
    _model_cache = {}

    def __init__(self, model_config: Dict[str, Any]):
        self.name = model_config['name']
        self.type = model_config.get('type', 'yolo')
        self.model = None

        # Determine a cache key based on the model's source
        cache_key = model_config.get('config_path') or model_config.get('weights_path')
        
        if cache_key in ModelRunner._model_cache:
            self.model = ModelRunner._model_cache[cache_key]
            # print(f"[DEBUG] Reusing cached model for {self.name}")
        else:
            if self.type == 'yolo':
                # 1. Load the base config from YAML
                config_params = load_config(model_config['config_path'])
                
                # 2. Override parameters with values from MODEL_TEMPLATES
                mapping = {'input_size': 'size'}
                for k, v in model_config.items():
                    if k in ['name', 'type', 'config_path']:
                        continue
                    target_key = mapping.get(k, k)
                    config_params[target_key] = v
                
                # 3. Initialize YOLO
                self.model = YOLO(config=config_params)
                
            elif self.type == 'embedding':
                self.model = ReIDDetectMultiBackend(
                    weights=model_config['weights_path'],
                    device=model_config.get('device', 'cpu'),
                    fp16=False
                )
            
            # Store in cache
            if cache_key:
                ModelRunner._model_cache[cache_key] = self.model

    def run_inference(self, input_data):
        return self.model(input_data)

    def close(self):
        # We don't close the model instance here because it might be shared in the cache.
        # The underlying resources will be kept alive for the next test iteration.
        pass

    @classmethod
    def clear_cache(cls):
        """Explicitly close and clear all cached models."""
        for cache_key, model in cls._model_cache.items():
            print(f">>> Closing cached model: {cache_key}")
            if hasattr(model, 'engine') and model.engine:
                try:
                    model.engine.close()
                except:
                    pass
            elif hasattr(model, 'backend') and model.backend:
                try:
                    model.backend.close()
                except:
                    pass
            # Specific check for ReID HailoInference
            elif hasattr(model, 'inference') and hasattr(model.inference, 'inference_instance'):
                try:
                    inst = model.inference.inference_instance
                    if hasattr(inst, 'close'):
                        inst.close()
                except:
                    pass
        cls._model_cache.clear()
