"""LiteLLM integration -- velocity-aware model routing.

Usage (OpenRouter -- recommended, most models available):
    from wzrd.litellm import completion

    response = completion(
        messages=[{"role": "user", "content": "Write a quicksort"}],
        task="code",
    )

Usage (specific provider):
    response = completion(
        messages=[{"role": "user", "content": "Explain monads"}],
        task="reasoning",
        provider="together",  # model must be available on this provider
    )

Simple composition (no wrapper needed):
    import wzrd, litellm
    response = litellm.completion(model=f"openrouter/{wzrd.pick('code')}", messages=[...])

Note: WZRD tracks model velocity across HuggingFace, OpenRouter, GitHub, and
ArtificialAnalysis. The model names are HuggingFace-style (org/model). Most are
available on OpenRouter. If using a different provider, the model may not exist
there -- set fallback_model to handle that case.
"""

from __future__ import annotations

import logging
from typing import Any, Optional, Sequence

logger = logging.getLogger("wzrd.litellm")

# Default fallback if WZRD API is unreachable
_FALLBACK_MODEL = "groq/llama-3.3-70b-versatile"


def _infer_task(messages: Sequence[dict], task: Optional[str] = None) -> str:
    """Infer the task type from explicit parameter or message content."""
    if task:
        return task

    # Check the last user message for task hints
    for msg in reversed(messages):
        if msg.get("role") != "user":
            continue
        content = (msg.get("content") or "").lower()
        if any(w in content for w in ("code", "function", "implement", "debug", "refactor", "write a", "python", "javascript", "typescript", "rust ", "golang")):
            return "code"
        if any(w in content for w in ("reason", "logic", "prove", "derive", "math")):
            return "reasoning"
        break

    return "general"


def pick_model(
    task: str = "general",
    fallback: str = _FALLBACK_MODEL,
    provider: Optional[str] = None,
) -> str:
    """Pick the highest-velocity model for a task. Returns a LiteLLM-compatible model string.

    Args:
        task: Task type (code, chat, reasoning, vision, general).
        fallback: Model to use if WZRD is unreachable.
        provider: LiteLLM provider prefix (e.g. "groq", "openrouter", "together").
                  If not set, uses "openrouter" as default since most WZRD-tracked
                  models are available there.

    Falls back to `fallback` if WZRD is unreachable.
    """
    try:
        from . import pick
        model = pick(task)
        if model:
            # Already has a LiteLLM provider prefix? Use as-is.
            known_prefixes = (
                "openai/", "anthropic/", "groq/", "google/", "openrouter/",
                "together/", "fireworks/", "cerebras/", "deepseek/",
            )
            if any(model.startswith(p) for p in known_prefixes):
                return model

            # Apply user-specified or default provider prefix
            prefix = provider or "openrouter"
            return f"{prefix}/{model}"
    except Exception as exc:
        logger.debug("wzrd.pick failed, using fallback: %s", exc)

    return fallback


def completion(
    messages: list[dict[str, Any]],
    *,
    task: Optional[str] = None,
    provider: Optional[str] = None,
    fallback_model: str = _FALLBACK_MODEL,
    **kwargs: Any,
) -> Any:
    """LiteLLM completion with automatic WZRD velocity routing.

    Args:
        messages: Chat messages (same as litellm.completion).
        task: Task type for model selection (code, chat, reasoning, vision).
              If not provided, inferred from message content.
        provider: LiteLLM provider prefix ("groq", "openrouter", "together", etc).
        fallback_model: Model to use if WZRD is unreachable.
        **kwargs: Passed through to litellm.completion().

    Returns:
        LiteLLM completion response.
    """
    try:
        import litellm as _litellm
    except ImportError:
        raise ImportError(
            "litellm is required for wzrd.litellm.completion(). "
            "Install it: pip install litellm"
        )

    inferred_task = _infer_task(messages, task)
    model = pick_model(inferred_task, fallback=fallback_model, provider=provider)
    logger.info("wzrd routing: task=%s model=%s", inferred_task, model)

    return _litellm.completion(model=model, messages=messages, **kwargs)
