from urllib.parse import urlparse

import aiohttp

from ...config.logfire_config import get_logger
from ..models import (
    EMBEDDING_DIMENSIONS,
    EMBEDDING_MODEL_PATTERNS,
    MODEL_CONTEXT_WINDOWS,
    VISION_MODEL_PATTERNS,
    ModelSpec,
)

logger = get_logger(__name__)


async def discover_ollama_models(base_urls: list[str], session: aiohttp.ClientSession, test_tool_fn) -> list[ModelSpec]:
    """1:1 Physical Parity Implementation from ProviderDiscoveryService."""
    all_models = []
    for base_url in base_urls:
        try:
            parsed = urlparse(base_url)
            api_url = base_url.replace("/v1", "") if parsed.path.endswith("/v1") else base_url

            async with session.get(f"{api_url}/api/tags") as response:
                if response.status == 200:
                    data = await response.json()
                    models = []
                    for model_info in data.get("models", []):
                        full_name = model_info.get("name", "")
                        model_name = full_name.split(":")[0]

                        supports_tools = await test_tool_fn(model_name, api_url)
                        supports_vision = any(p in model_name.lower() for p in VISION_MODEL_PATTERNS)
                        supports_embeddings = any(p in model_name.lower() for p in EMBEDDING_MODEL_PATTERNS)

                        context_window = 4096
                        for family, window_size in MODEL_CONTEXT_WINDOWS.items():
                            if family in model_name.lower():
                                context_window = window_size
                                break

                        embedding_dims = next(
                            (dims for pattern, dims in EMBEDDING_DIMENSIONS.items() if pattern in model_name.lower()),
                            None,
                        )

                        models.append(
                            ModelSpec(
                                name=full_name,
                                provider="ollama",
                                context_window=context_window,
                                supports_tools=supports_tools,
                                supports_vision=supports_vision,
                                supports_embeddings=supports_embeddings,
                                embedding_dimensions=embedding_dims,
                                description=f"Ollama model on {base_url}",
                                aliases=[model_name] if ":" in full_name else [],
                            )
                        )
                    all_models.extend(models)
                    logger.info(f"Discovered {len(models)} Ollama models from {base_url}")
        except Exception as e:
            logger.error(f"Error discovering Ollama models from {base_url}: {e}")
    return all_models