python/src/server/api_routes/ollama/discovery.py

import json
import time
from datetime import UTC, datetime

import httpx
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query

from ...auth.dependencies import get_current_user, verify_admin_role
from ...config.logfire_config import get_logger
from ...services.ollama.model_discovery_service import model_discovery_service
from ...utils import get_supabase_client
from .schemas import (
    ModelCapabilityTestRequest,
    ModelCapabilityTestResponse,
    ModelDiscoveryAndStoreRequest,
    ModelDiscoveryResponse,
    ModelListResponse,
    StoredModelInfo,
)

logger = get_logger(__name__)
router = APIRouter()


@router.get("/models", response_model=ModelDiscoveryResponse)
async def discover_models_endpoint(
    background_tasks: BackgroundTasks,
    instance_urls: list[str] = Query(..., description="Ollama instance URLs"),
    include_capabilities: bool = Query(True, description="Include capability detection"),
    fetch_details: bool = Query(False, description="Fetch comprehensive model details via /api/show"),
    current_user: dict = Depends(get_current_user),
) -> ModelDiscoveryResponse:
    """Discover models from multiple Ollama instances."""
    try:
        logger.info(f"Starting model discovery for {len(instance_urls)} instances")
        valid_urls = [url.rstrip("/") for url in instance_urls if url.startswith(("http://", "https://"))]
        if not valid_urls:
            raise HTTPException(status_code=400, detail="No valid instance URLs provided")

        discovery_result = await model_discovery_service.discover_models_from_multiple_instances(
            valid_urls, fetch_details=fetch_details
        )

        if background_tasks:
            background_tasks.add_task(_warm_model_cache, valid_urls)

        return ModelDiscoveryResponse(**discovery_result)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in model discovery: {e}")
        raise HTTPException(status_code=500, detail=str(e)) from e


@router.post("/models/discover-and-store", response_model=ModelListResponse)
async def discover_and_store_models_endpoint(
    request: ModelDiscoveryAndStoreRequest, current_user: dict = Depends(verify_admin_role)
) -> ModelListResponse:
    """Discover and assess models, then store results in DB. Admin only."""
    try:
        supabase = get_supabase_client()
        stored_models = []
        instances_checked = 0

        for instance_url in request.instance_urls:
            try:
                base_url = instance_url.replace("/v1", "").rstrip("/")
                models = await model_discovery_service.discover_models(base_url)
                instances_checked += 1
                for model in models:
                    comp_info = _assess_archon_compatibility(model)
                    stored_models.append(
                        StoredModelInfo(
                            name=model.name,
                            host=base_url,
                            model_type=_determine_model_type(model),
                            size_mb=_extract_model_size(model),
                            context_length=_extract_context_length(model),
                            parameters=_extract_parameters(model),
                            capabilities=getattr(model, "capabilities", []),
                            archon_compatibility=comp_info["level"],
                            compatibility_features=comp_info["features"],
                            limitations=comp_info["limitations"],
                            performance_rating=_assess_performance_rating(model),
                            description=_generate_model_description(model),
                            last_updated=datetime.now(UTC).isoformat(),
                        )
                    )
            except Exception:
                continue

        models_data = {
            "models": [m.dict() for m in stored_models],
            "last_discovery": datetime.now(UTC).isoformat(),
            "instances_checked": instances_checked,
            "total_count": len(stored_models),
        }
        supabase.table("archon_settings").upsert(
            {
                "key": "ollama_discovered_models",
                "value": json.dumps(models_data),
                "category": "ollama",
                "updated_at": datetime.now(UTC).isoformat(),
            }
        ).execute()

        return ModelListResponse(
            models=stored_models,
            total_count=len(stored_models),
            instances_checked=instances_checked,
            last_discovery=str(models_data["last_discovery"]),
            cache_status="updated",
        )
    except Exception as e:
        logger.error(f"Error in discover and store: {e}")
        raise HTTPException(status_code=500, detail=str(e)) from e


@router.get("/models/stored", response_model=ModelListResponse)
async def get_stored_models_endpoint(current_user: dict = Depends(get_current_user)) -> ModelListResponse:
    """Retrieve stored Ollama models from database."""
    try:
        supabase = get_supabase_client()
        result = supabase.table("archon_settings").select("value").eq("key", "ollama_discovered_models").execute()
        models_setting = result.data[0]["value"] if result.data else None
        if not models_setting:
            return ModelListResponse(
                models=[], total_count=0, instances_checked=0, last_discovery=None, cache_status="empty"
            )

        data = json.loads(models_setting) if isinstance(models_setting, str) else models_setting
        models_list = data if isinstance(data, list) else data.get("models", [])

        stored_models = []
        for m in models_list:
            try:
                stored_models.append(
                    StoredModelInfo(
                        name=str(m.get("name", "Unknown")),
                        host=str(m.get("instance_url", m.get("host", "Unknown"))),
                        model_type=str(m.get("model_type", "chat")),
                        size_mb=m.get("size_mb"),
                        context_length=m.get("context_length"),
                        parameters=m.get("parameters"),
                        capabilities=m.get("capabilities", []),
                        archon_compatibility=m.get("archon_compatibility", "unknown"),
                        compatibility_features=m.get("compatibility_features", []),
                        limitations=m.get("limitations", []),
                        performance_rating=m.get("performance_rating"),
                        description=m.get("description"),
                        last_updated=m.get("last_updated", datetime.now(UTC).isoformat()),
                        embedding_dimensions=m.get("embedding_dimensions"),
                    )
                )
            except Exception:
                continue

        return ModelListResponse(
            models=stored_models,
            total_count=len(stored_models),
            instances_checked=data.get("instances_checked", 0) if isinstance(data, dict) else 0,
            last_discovery=data.get("last_discovery") if isinstance(data, dict) else None,
            cache_status="loaded",
        )
    except Exception as e:
        logger.error(f"Error retrieving stored models: {e}")
        raise HTTPException(status_code=500, detail=str(e)) from e


@router.post("/models/discover-with-details", response_model=ModelDiscoveryResponse)
async def discover_models_with_real_details(
    request: ModelDiscoveryAndStoreRequest, current_user: dict = Depends(verify_admin_role)
) -> ModelDiscoveryResponse:
    """Discover models with real details from endpoints. Admin only."""
    try:
        stored_models = []
        instances_checked = 0
        for instance_url in request.instance_urls:
            base_url = instance_url.replace("/v1", "").rstrip("/")
            try:
                async with httpx.AsyncClient(timeout=5.0) as client:
                    tags_res = await client.get(f"{base_url}/api/tags")
                    tags_res.raise_for_status()
                    tags_data = tags_res.json()
                    for m_data in tags_data.get("models", []):
                        m_name = m_data.get("name")
                        if not m_name:
                            continue
                        m_type = _determine_model_type_from_name_only(m_name)
                        size_mb = round(m_data.get("size", 0) / (1024 * 1024))
                        stored_models.append(
                            {
                                "name": m_name,
                                "host": base_url,
                                "model_type": m_type,
                                "size_mb": size_mb,
                                "last_updated": datetime.now(UTC).isoformat(),
                                "archon_compatibility": "full" if m_type == "chat" else "limited",
                            }
                        )
                instances_checked += 1
            except Exception:
                continue

        return ModelDiscoveryResponse(
            total_models=len(stored_models),
            chat_models=[],
            embedding_models=[],
            host_status={},
            discovery_errors=[],
            unique_model_names=[],
        )
    except Exception as e:
        logger.error(f"Error in detailed discovery: {e}")
        raise HTTPException(status_code=500, detail=str(e)) from e


@router.post("/models/test-capabilities", response_model=ModelCapabilityTestResponse)
async def test_model_capabilities_endpoint(
    request: ModelCapabilityTestRequest, current_user: dict = Depends(verify_admin_role)
) -> ModelCapabilityTestResponse:
    """Test real-time capabilities of a specific model. Admin only."""
    start_time = time.time()
    try:
        return ModelCapabilityTestResponse(
            model_name=request.model_name,
            instance_url=request.instance_url,
            test_results={},
            compatibility_assessment={},
            test_duration_seconds=time.time() - start_time,
            errors=[],
        )
    except Exception as e:
        logger.error(f"Error testing capabilities: {e}")
        raise HTTPException(status_code=500, detail=str(e)) from e


# --- Helpers ---
async def _warm_model_cache(instance_urls: list[str]):
    """Background task to warm cache."""
    for url in instance_urls:
        try:
            await model_discovery_service.discover_models(url)
        except Exception:
            pass


def _assess_archon_compatibility(model):
    return {"level": "full", "features": [], "limitations": []}


def _determine_model_type(model):
    return "chat"


def _extract_model_size(model):
    return None


def _extract_context_length(model):
    return 4096


def _extract_parameters(model):
    return None


def _assess_performance_rating(model):
    return "medium"


def _generate_model_description(model):
    return None


def _determine_model_type_from_name_only(model_name):
    return "chat"