"""Model Routing - Capability-based model selection.

This module provides:
- get_model_for_capability(): Get the default model for a capability
- ModelCapability: Enum of supported capabilities

Configuration is loaded from config/models.yaml.

Usage:
    from stdlib.llm.routing import get_model_for_capability, ModelCapability

    # Get default model for a capability
    model = get_model_for_capability(ModelCapability.CODE)
    # Returns: "anthropic/claude-3.5-sonnet"

    # Or use string capability
    model = get_model_for_capability("reasoning")
    # Returns: "openai/gpt-4o"
"""

from __future__ import annotations

import logging
import os
from enum import Enum
from pathlib import Path
from typing import Any

import yaml

logger = logging.getLogger(__name__)


class ModelCapability(str, Enum):
    """Model capabilities for routing decisions.

    Use these instead of magic strings like "auto".
    """

    FAST = "fast"  # Quick, cheap: gpt-4o-mini
    REASONING = "reasoning"  # Complex thinking: gpt-4o
    CODE = "code"  # Code generation: claude-3.5-sonnet
    VISION = "vision"  # Image understanding: gpt-4o
    EMBEDDING = "embedding"  # Vector embeddings
    LOCAL = "local"  # Local/offline: llama3


# =============================================================================
# Configuration Loading
# =============================================================================

_config: dict[str, Any] | None = None


def _load_config() -> dict[str, Any]:
    """Load model configuration from YAML."""
    global _config

    if _config is not None:
        return _config

    # Find config file
    env_config = os.getenv("ALBUS_MODELS_CONFIG", "")
    config_paths = [
        Path(env_config) if env_config else None,
        Path.cwd() / "config" / "models.yaml",
        Path(__file__).parent.parent.parent.parent / "config" / "models.yaml",
    ]
    config_paths = [p for p in config_paths if p is not None]

    for config_path in config_paths:
        if config_path.exists():
            try:
                _config = yaml.safe_load(config_path.read_text())
                logger.debug("Loaded model config from %s", config_path)
                return _config
            except Exception as e:
                logger.warning("Failed to load model config: %s", e)

    # Fallback defaults
    _config = {
        "defaults": {
            "fast": "openai/gpt-4o-mini",
            "reasoning": "openai/gpt-4o",
            "code": "anthropic/claude-3.5-sonnet",
            "vision": "openai/gpt-4o",
            "embedding": "openai/text-embedding-3-small",
            "local": "llama3.1:8b",
        }
    }
    logger.debug("Using fallback model config")
    return _config


def reload_config() -> None:
    """Force reload of model configuration."""
    global _config
    _config = None
    _load_config()


# =============================================================================
# Model Selection
# =============================================================================


def get_model_for_capability(
    capability: str | ModelCapability,
    *,
    fallback: str | None = None,
) -> str:
    """Get the model for a capability.

    Resolution order:
    1. Runtime overrides (from albus.yaml ``models.routing`` or API)
    2. Defaults from ``config/models.yaml``
    3. Explicit fallback parameter
    4. Ultimate fallback (``openai/gpt-4o-mini``)

    Args:
        capability: The capability needed (e.g., ModelCapability.CODE or "code")
        fallback: Fallback model if capability not found

    Returns
    -------
        Model name in OpenRouter format (e.g., "openai/gpt-4o")
    """
    config = _load_config()
    defaults = config.get("defaults", {})

    # Normalize capability
    cap_str = capability.value if isinstance(capability, ModelCapability) else capability.lower()

    # 1. Check runtime overrides first (from albus.yaml or API)
    if cap_str in _runtime_routing:
        return _runtime_routing[cap_str]

    # 2. Look up in config/models.yaml defaults
    model = defaults.get(cap_str)
    if model:
        return model

    # 3. Explicit fallback
    if fallback:
        return fallback

    # 4. Ultimate fallback
    return _runtime_routing.get("fast", defaults.get("fast", "openai/gpt-4o-mini"))


def get_default_model() -> str:
    """Get the default model (fast capability)."""
    return get_model_for_capability(ModelCapability.FAST)


def list_capabilities() -> list[str]:
    """List all configured capabilities."""
    config = _load_config()
    return list(config.get("defaults", {}).keys())


def get_model_info(model: str) -> dict[str, Any] | None:
    """Get information about a specific model.

    Args:
        model: Model name (e.g., "openai/gpt-4o")

    Returns
    -------
        Model info dict or None if not found
    """
    config = _load_config()
    models = config.get("models", {})
    return models.get(model)


# =============================================================================
# Runtime Model Config (for dynamic updates)
# =============================================================================

_runtime_routing: dict[str, str] = {}


def set_runtime_model_config(routing: dict[str, str] | None = None) -> None:
    """Set runtime model routing overrides.

    Args:
        routing: Dict mapping capability to model name
    """
    global _runtime_routing
    if routing:
        _runtime_routing.update(routing)
        logger.info(f"Runtime model routing updated: {routing}")


def get_runtime_model_config() -> dict[str, Any]:
    """Get current runtime model configuration.

    Returns
    -------
        Dict with defaults and any runtime overrides
    """
    config = _load_config()
    return {
        "defaults": config.get("defaults", {}),
        "runtime_overrides": _runtime_routing,
        "effective": {
            **config.get("defaults", {}),
            **_runtime_routing,
        },
    }


__all__ = [
    # Enum
    "ModelCapability",
    "get_default_model",
    # Functions
    "get_model_for_capability",
    "get_model_info",
    "get_runtime_model_config",
    "list_capabilities",
    "reload_config",
    # Runtime config
    "set_runtime_model_config",
]
