"""
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._lastRefresh: Optional[float] = None
        self._refreshInterval: float = 300.0  # 5 minutes
    
    def registerConnector(self, connector: BaseConnectorAi):
        """Register a connector and collect its models."""
        connectorType = connector.getConnectorType()
        self._connectors[connectorType] = 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 {connectorType}")
        except Exception as e:
            logger.error(f"Failed to register models from {connectorType}: {e}")
    
    def discoverConnectors(self) -> List[BaseConnectorAi]:
        """Auto-discover connectors by scanning aicorePlugin*.py files."""
        connectors = []
        connectorDir = os.path.dirname(__file__)
        
        # Scan for connector files
        for filename in os.listdir(connectorDir):
            if filename.startswith('aicorePlugin') and filename.endswith('.py'):
                moduleName = filename[:-3]  # Remove .py extension
                
                try:
                    # Import the module
                    module = importlib.import_module(f'modules.connectors.{moduleName}')
                    
                    # Find connector classes (classes that inherit from BaseConnectorAi)
                    for attrName in dir(module):
                        attr = getattr(module, attrName)
                        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
        
        currentTime = time.time()
        
        # Check if refresh is needed
        if (not force and 
            self._lastRefresh is not None and 
            currentTime - self._lastRefresh < self._refreshInterval):
            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._lastRefresh = currentTime
        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, connectorType: str) -> List[AiModel]:
        """Get models from a specific connector."""
        self.refreshModels()
        return [model for model in self._models.values() if model.connectorType == connectorType]
    
    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 getModelsByPriority(self, priority: str) -> List[AiModel]:
        """Get models that have a specific priority."""
        self.refreshModels()
        return [model for model in self._models.values() if model.priority == priority]
    
    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, modelName: str) -> Optional[BaseConnectorAi]:
        """Get the connector instance for a specific model."""
        model = self.getModel(modelName)
        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 = {
            "totalModels": len(self._models),
            "availableModels": len([m for m in self._models.values() if m.isAvailable]),
            "connectors": len(self._connectors),
            "byConnector": {},
            "byCapability": {},
            "byPriority": {}
        }
        
        # Count by connector
        for model in self._models.values():
            connector = model.connectorType
            if connector not in stats["byConnector"]:
                stats["byConnector"][connector] = 0
            stats["byConnector"][connector] += 1
        
        # Count by capability
        for model in self._models.values():
            for capability in model.capabilities:
                if capability not in stats["byCapability"]:
                    stats["byCapability"][capability] = 0
                stats["byCapability"][capability] += 1
        
        # Count by priority
        for model in self._models.values():
            priority = model.priority
            if priority not in stats["byPriority"]:
                stats["byPriority"][priority] = 0
            stats["byPriority"][priority] += 1
        
        return stats


# Global registry instance
model_registry = ModelRegistry()