# aiva/models/model_loader.py
import logging
from opentelemetry import trace

logger = logging.getLogger(__name__)
tracer = trace.get_tracer(__name__)

class ModelLoader:
    """Loads and manages AI models."""

    def __init__(self, model_config):
        self.model_config = model_config
        self.models = {}
        self.load_models()

    def load_models(self):
        """Loads models based on the configuration."""
        with tracer.start_as_current_span("load_models") as span:
            for model_name, config in self.model_config.items():
                try:
                    # Simulate loading a model (replace with actual loading logic)
                    logger.info(f"Loading model: {model_name}")
                    self.models[model_name] = f"Model {model_name} - Version {config.get('version', 'unknown')}"

                    # Add model version as a span attribute
                    span.set_attribute(f"model.{model_name}.version", config.get('version', 'unknown'))

                except Exception as e:
                    logger.error(f"Failed to load model {model_name}: {e}")
                    span.record_exception(e)

    def get_model(self, model_name):
        """Returns a loaded model."""
        return self.models.get(model_name)

    def get_model_versions(self):
        """Returns a dictionary of model names and their versions."""
        versions = {}
        for model_name, config in self.model_config.items():
            versions[model_name] = config.get('version', 'unknown')
        return versions