python/src/server/services/ollama/model_discovery_service.py

"""
Ollama Model Discovery Service

Provides comprehensive model discovery, validation, and capability detection for Ollama instances.
Supports multi-instance configurations with automatic dimension detection and health monitoring.
"""

import asyncio
import time
from typing import Any

from src.server.config.logfire_config import get_logger

from .discovery.manifest_parser import OllamaManifestParser
from .discovery.models import InstanceHealthStatus, ModelCapabilities, OllamaModel
from .discovery.network_scanner import OllamaNetworkScanner

logger = get_logger(__name__)


class ModelDiscoveryService:
    """Service for discovering and validating Ollama models across multiple instances."""

    def __init__(self):
        self.model_cache: dict[str, list[OllamaModel]] = {}
        self.capability_cache: dict[str, ModelCapabilities] = {}
        self.capability_lock = asyncio.Lock()  # Restored for concurrency control
        self.health_cache: dict[str, InstanceHealthStatus] = {}
        self.cache_ttl = 300  # 5 minutes TTL
        self.scanner = OllamaNetworkScanner(timeout=30)
        self.parser = OllamaManifestParser()

    def _get_cached_models(self, instance_url: str) -> list[OllamaModel] | None:
        """Get cached models if not expired."""
        cache_key = f"models_{instance_url}"
        cached_data = self.model_cache.get(cache_key)
        if cached_data:
            # Check if any model in cache is still valid (simple TTL check)
            first_model = cached_data[0] if cached_data else None
            if first_model and first_model.last_updated:
                cache_time = float(first_model.last_updated)
                if time.time() - cache_time < self.cache_ttl:
                    logger.debug(f"Using cached models for {instance_url}")
                    return cached_data
                else:
                    # Expired, remove from cache
                    del self.model_cache[cache_key]
        return None

    def _cache_models(self, instance_url: str, models: list[OllamaModel]) -> None:
        """Cache models with current timestamp."""
        cache_key = f"models_{instance_url}"
        # Set timestamp for cache expiry
        current_time = str(time.time())
        for model in models:
            model.last_updated = current_time
        self.model_cache[cache_key] = models
        logger.debug(f"Cached {len(models)} models for {instance_url}")

    async def discover_models(self, instance_url: str, fetch_details: bool = False) -> list[OllamaModel]:
        """
        Discover all available models from an Ollama instance.

        Args:
            instance_url: Base URL of the Ollama instance
            fetch_details: If True, fetch comprehensive model details via /api/show

        Returns:
            List of OllamaModel objects with discovered capabilities
        """
        # Check cache first (but skip if we need detailed info)
        if not fetch_details:
            cached_models = self._get_cached_models(instance_url)
            if cached_models:
                return cached_models

        try:
            logger.info(f"Discovering models from Ollama instance: {instance_url}")

            # Use network scanner to fetch raw tags
            data = await self.scanner.fetch_tags(instance_url)

            # Use manifest parser to parse the tags
            models = self.parser.parse_tags_response(data, instance_url)

            logger.info(f"Discovered {len(models)} models from {instance_url}")

            # Enrich models with capability information
            enriched_models = await self._enrich_model_capabilities(models, instance_url, fetch_details=fetch_details)

            # Cache the results
            self._cache_models(instance_url, enriched_models)

            return enriched_models

        except Exception as e:
            logger.error(f"Error discovering models from {instance_url}: {e}")
            # The network_scanner raises standard exceptions which we propagate
            raise Exception(f"Failed to discover models: {str(e)}") from e

    async def _enrich_model_capabilities(
        self, models: list[OllamaModel], instance_url: str, fetch_details: bool = False
    ) -> list[OllamaModel]:
        """Pattern-match and enrich model capabilities. Delegates to capabilities submodule."""
        from .discovery.capabilities import enrich_model_capabilities_logic

        return await enrich_model_capabilities_logic(self, models, instance_url, fetch_details)

    async def _detect_model_capabilities_optimized(self, model_name: str, instance_url: str) -> ModelCapabilities:
        """Optimized capability detection. Delegates to capabilities submodule."""
        from .discovery.capabilities import detect_model_capabilities_logic

        return await detect_model_capabilities_logic(self, model_name, instance_url, optimized=True)

    async def _detect_model_capabilities(self, model_name: str, instance_url: str) -> ModelCapabilities:
        """Comprehensive capability detection. Delegates to capabilities submodule."""
        from .discovery.capabilities import detect_model_capabilities_logic

        return await detect_model_capabilities_logic(self, model_name, instance_url, optimized=False)

    # --- Private Facades for backward compatibility and testing ---
    async def _get_model_details(self, model_name: str, instance_url: str) -> dict[str, Any] | None:
        from .discovery.capabilities import get_model_details_logic

        return await get_model_details_logic(model_name, instance_url)

    async def _test_chat_capability(self, model_name: str, instance_url: str) -> bool:
        from .discovery.capabilities import test_chat_capability_logic

        return await test_chat_capability_logic(model_name, instance_url)

    async def _test_embedding_capability(self, model_name: str, instance_url: str) -> int | None:
        from .discovery.capabilities import test_embedding_capability_logic

        return await test_embedding_capability_logic(model_name, instance_url)

    async def _test_function_calling_capability(self, model_name: str, instance_url: str) -> bool:
        from .discovery.capabilities import test_function_calling_capability_logic

        return await test_function_calling_capability_logic(model_name, instance_url)

    async def _test_structured_output_capability(self, model_name: str, instance_url: str) -> bool:
        from .discovery.capabilities import test_structured_output_capability_logic

        return await test_structured_output_capability_logic(model_name, instance_url)

    async def _test_embedding_capability_fast(self, model_name: str, instance_url: str) -> int | None:
        from .discovery.capabilities import test_embedding_capability_fast_logic

        return await test_embedding_capability_fast_logic(model_name, instance_url)

    async def _test_chat_capability_fast(self, model_name: str, instance_url: str) -> bool:
        from .discovery.capabilities import test_chat_capability_fast_logic

        return await test_chat_capability_fast_logic(model_name, instance_url)

    async def _test_structured_output_capability_fast(self, model_name: str, instance_url: str) -> bool:
        from .discovery.capabilities import test_structured_output_capability_fast_logic

        return await test_structured_output_capability_fast_logic(model_name, instance_url)

    async def validate_model_capabilities(self, model_name: str, instance_url: str, required_capability: str) -> bool:
        """
        Validate that a model supports a required capability.

        Args:
            model_name: Name of the model to validate
            instance_url: Ollama instance URL
            required_capability: 'chat' or 'embedding'

        Returns:
            True if model supports the capability, False otherwise
        """
        try:
            capabilities = await self._detect_model_capabilities(model_name, instance_url)

            if required_capability == "chat":
                return capabilities.supports_chat
            elif required_capability == "embedding":
                return capabilities.supports_embedding
            elif required_capability == "function_calling":
                return capabilities.supports_function_calling
            elif required_capability == "structured_output":
                return capabilities.supports_structured_output
            else:
                logger.warning(f"Unknown capability requirement: {required_capability}")
                return False

        except Exception as e:
            logger.error(f"Error validating model {model_name} for {required_capability}: {e}")
            return False

    async def get_model_info(self, model_name: str, instance_url: str) -> OllamaModel | None:
        """
        Get comprehensive information about a specific model.

        Args:
            model_name: Name of the model
            instance_url: Ollama instance URL

        Returns:
            OllamaModel object with complete information or None if not found
        """
        try:
            models = await self.discover_models(instance_url)

            for model in models:
                if model.name == model_name:
                    return model

            logger.warning(f"Model {model_name} not found on instance {instance_url}")
            return None

        except Exception as e:
            logger.error(f"Error getting model info for {model_name}: {e}")
            return None

    async def check_instance_health(self, instance_url: str) -> InstanceHealthStatus:
        """
        Check the health status of an Ollama instance.

        Args:
            instance_url: Base URL of the Ollama instance

        Returns:
            InstanceHealthStatus with current health information
        """
        # Check cache first (shorter TTL for health checks)
        cache_key = f"health_{instance_url}"
        if cache_key in self.health_cache:
            cached_health = self.health_cache[cache_key]
            if cached_health.last_checked:
                cache_time = float(cached_health.last_checked)
                # Use shorter cache for health (30 seconds)
                if time.time() - cache_time < 30:
                    return cached_health

        status = InstanceHealthStatus(is_healthy=False)
        is_healthy, response_time_ms, models_count, error_message = await self.scanner.ping_health(instance_url)

        if is_healthy:
            status.is_healthy = True
            status.response_time_ms = response_time_ms
            status.models_available = models_count
            status.last_checked = str(time.time())
            logger.debug(f"Instance {instance_url} is healthy: {models_count} models, {status.response_time_ms:.0f}ms")
        else:
            status.error_message = error_message
            logger.warning(f"Health check failed for {instance_url}: {error_message}")

        # Cache the result
        self.health_cache[cache_key] = status

        return status

    async def discover_models_from_multiple_instances(
        self, instance_urls: list[str], fetch_details: bool = False
    ) -> dict[str, Any]:
        """
        Discover models from multiple Ollama instances concurrently.

        Args:
            instance_urls: List of Ollama instance URLs
            fetch_details: If True, fetch comprehensive model details via /api/show

        Returns:
            Dictionary with discovery results and aggregated information
        """
        if not instance_urls:
            return {
                "total_models": 0,
                "chat_models": [],
                "embedding_models": [],
                "host_status": {},
                "discovery_errors": [],
            }

        logger.info(f"Discovering models from {len(instance_urls)} Ollama instances with fetch_details={fetch_details}")

        # Discover models from all instances concurrently
        tasks = [self.discover_models(url, fetch_details=fetch_details) for url in instance_urls]
        results = await asyncio.gather(*tasks, return_exceptions=True)

        # Delegate aggregation to the parser
        discovery_result = self.parser.aggregate_discovery_results(instance_urls, results)

        logger.info(
            f"Discovery complete: {discovery_result['total_models']} total models, "
            f"{len(discovery_result['chat_models'])} chat, {len(discovery_result['embedding_models'])} embedding"
        )

        return discovery_result


# Global service instance
model_discovery_service = ModelDiscoveryService()