"""
Dynamic model registry that collects models from all AI connectors.
Implements plugin-like architecture for connector discovery.
"""

import logging
import importlib
import os
from typing import Dict, List, Optional, Any
from modules.datamodels.datamodelAi import AiModel
from modules.aicore.aicoreBase import BaseConnectorAi

logger = logging.getLogger(__name__)


class ModelRegistry:
    """Dynamic registry for AI models from all connectors."""
    
    def __init__(self):
        self._models: Dict[str, AiModel] = {}
        self._connectors: Dict[str, BaseConnectorAi] = {}
        self._last_refresh: Optional[float] = None
        self._refresh_interval: float = 300.0  # 5 minutes
    
    def registerConnector(self, connector: BaseConnectorAi):
        """Register a connector and collect its models."""
        connector_type = connector.getConnectorType()
        self._connectors[connector_type] = connector
        
        # Collect models from this connector
        try:
            models = connector.getCachedModels()
            for model in models:
                self._models[model.name] = model
                logger.debug(f"Registered model: {model.name} from {connector_type}")
        except Exception as e:
            logger.error(f"Failed to register models from {connector_type}: {e}")
    
    def discoverConnectors(self) -> List[BaseConnectorAi]:
        """Auto-discover connectors by scanning aicorePlugin*.py files."""
        connectors = []
        connector_dir = os.path.dirname(__file__)
        
        # Scan for connector files
        for filename in os.listdir(connector_dir):
            if filename.startswith('aicorePlugin') and filename.endswith('.py'):
                module_name = filename[:-3]  # Remove .py extension
                
                try:
                    # Import the module
                    module = importlib.import_module(f'modules.connectors.{module_name}')
                    
                    # Find connector classes (classes that inherit from BaseConnectorAi)
                    for attr_name in dir(module):
                        attr = getattr(module, attr_name)
                        if (isinstance(attr, type) and 
                            issubclass(attr, BaseConnectorAi) and 
                            attr != BaseConnectorAi):
                            
                            # Instantiate the connector
                            connector = attr()
                            connectors.append(connector)
                            logger.info(f"Discovered connector: {connector.getConnectorType()}")
                            
                except Exception as e:
                    logger.warning(f"Failed to discover connector from {filename}: {e}")
        
        return connectors
    
    def refreshModels(self, force: bool = False):
        """Refresh models from all registered connectors."""
        import time
        
        current_time = time.time()
        
        # Check if refresh is needed
        if (not force and 
            self._last_refresh is not None and 
            current_time - self._last_refresh < self._refresh_interval):
            return
        
        logger.info("Refreshing model registry...")
        
        # Clear existing models
        self._models.clear()
        
        # Re-register all connectors
        for connector in self._connectors.values():
            try:
                connector.clearCache()  # Clear connector cache
                models = connector.getCachedModels()
                for model in models:
                    self._models[model.name] = model
            except Exception as e:
                logger.error(f"Failed to refresh models from {connector.getConnectorType()}: {e}")
        
        self._last_refresh = current_time
        logger.info(f"Model registry refreshed: {len(self._models)} models available")
    
    def getModel(self, name: str) -> Optional[AiModel]:
        """Get a specific model by name."""
        self.refreshModels()
        return self._models.get(name)
    
    def getModels(self) -> List[AiModel]:
        """Get all available models."""
        self.refreshModels()
        return list(self._models.values())
    
    def getModelsByConnector(self, connector_type: str) -> List[AiModel]:
        """Get models from a specific connector."""
        self.refreshModels()
        return [model for model in self._models.values() if model.connectorType == connector_type]
    
    def getModelsByCapability(self, capability: str) -> List[AiModel]:
        """Get models that support a specific capability."""
        self.refreshModels()
        return [model for model in self._models.values() if capability in model.capabilities]
    
    def getModelsByTag(self, tag: str) -> List[AiModel]:
        """Get models that have a specific tag."""
        self.refreshModels()
        return [model for model in self._models.values() if tag in model.tags]
    
    def getAvailableModels(self) -> List[AiModel]:
        """Get only available models."""
        self.refreshModels()
        return [model for model in self._models.values() if model.isAvailable]
    
    def getConnectorForModel(self, model_name: str) -> Optional[BaseConnectorAi]:
        """Get the connector instance for a specific model."""
        model = self.getModel(model_name)
        if model:
            return self._connectors.get(model.connectorType)
        return None
    
    def getModelStats(self) -> Dict[str, Any]:
        """Get statistics about the model registry."""
        self.refreshModels()
        
        stats = {
            "total_models": len(self._models),
            "available_models": len([m for m in self._models.values() if m.isAvailable]),
            "connectors": len(self._connectors),
            "by_connector": {},
            "by_capability": {},
            "by_tag": {}
        }
        
        # Count by connector
        for model in self._models.values():
            connector = model.connectorType
            if connector not in stats["by_connector"]:
                stats["by_connector"][connector] = 0
            stats["by_connector"][connector] += 1
        
        # Count by capability
        for model in self._models.values():
            for capability in model.capabilities:
                if capability not in stats["by_capability"]:
                    stats["by_capability"][capability] = 0
                stats["by_capability"][capability] += 1
        
        # Count by tag
        for model in self._models.values():
            for tag in model.tags:
                if tag not in stats["by_tag"]:
                    stats["by_tag"][tag] = 0
                stats["by_tag"][tag] += 1
        
        return stats


# Global registry instance
model_registry = ModelRegistry()