"""Embedding providers for RAG with fallback chain.

Auto-detects the best available provider from the user's existing config.
No extra API keys required — reuses whatever provider keys are already set up.

Priority chain: Gemini (free tier) → OpenAI → FTS-only mode (no embeddings).
Anthropic and Groq do not offer embedding APIs.

Fallback distinguishes:
- Soft failures (missing API key, unsupported provider) → skip, try next
- Hard failures (network error, rate limit during embed()) → raise
"""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod

logger = logging.getLogger(__name__)


class EmbedderSoftError(Exception):
    """Non-fatal error: provider unavailable, try next in chain."""


class EmbedderHardError(Exception):
    """Fatal error: network failure, rate limit — stop trying."""


class Embedder(ABC):
    """Abstract base for embedding providers."""

    dimensions: int
    provider_name: str = "unknown"

    @abstractmethod
    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Embed a batch of texts into vectors."""
        ...

    async def embed_single(self, text: str) -> list[float]:
        """Embed a single text string."""
        return (await self.embed([text]))[0]


class FastEmbedEmbedder(Embedder):
    """CPU-based embedding via FastEmbed (ONNX runtime, no GPU needed).

    Uses nomic-embed-text-v1.5 (137MB, 768 dims).
    Production-ready: runs on CPU, no VRAM conflict with chat model.
    """

    provider_name = "fastembed"

    def __init__(self, model: str = "nomic-ai/nomic-embed-text-v1.5"):
        from fastembed import TextEmbedding

        self._model = TextEmbedding(model)
        self.dimensions = 768
        logger.info("FastEmbed initialized: %s (CPU, %d dims)", model, self.dimensions)

    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Embed texts on CPU via ONNX runtime."""
        import asyncio

        if not texts:
            return []

        def _sync_embed():
            return [emb.tolist() for emb in self._model.embed(texts)]

        return await asyncio.to_thread(_sync_embed)


_EMBED_BATCH_SIZE = 100

_OPENAI_MODEL_DIMS: dict[str, int] = {
    "nomic-embed-text": 768,
    "text-embedding-3-small": 1536,
    "text-embedding-3-large": 3072,
}


class GeminiEmbedder(Embedder):
    """Google Gemini embedding via OpenAI-compatible endpoint (free tier)."""

    provider_name = "gemini"

    def __init__(self, api_key: str, model: str = "gemini-embedding-001", base_url: str | None = None):
        import openai

        self.client = openai.AsyncOpenAI(
            api_key=api_key,
            base_url=base_url or "https://generativelanguage.googleapis.com/v1beta/openai/",
        )
        self.model = model
        self.dimensions = 3072  # gemini-embedding-001 returns 3072 dims

    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Embed texts using Gemini's OpenAI-compatible endpoint."""
        if not texts:
            return []
        all_embeddings: list[list[float]] = [None] * len(texts)  # type: ignore[list-item]
        for i in range(0, len(texts), _EMBED_BATCH_SIZE):
            batch = texts[i : i + _EMBED_BATCH_SIZE]
            response = await self.client.embeddings.create(
                input=batch,
                model=self.model,
            )
            # Some providers return d.index=None — use enumerate as fallback
            for j, d in enumerate(response.data):
                idx = d.index if d.index is not None else j
                all_embeddings[i + idx] = d.embedding
            logger.debug("Gemini embedded batch %d-%d (%d texts)", i, i + len(batch), len(batch))
        return all_embeddings


_EMBED_BATCH_SIZE = 100

_OPENAI_MODEL_DIMS: dict[str, int] = {
    "nomic-embed-text": 768,
    "text-embedding-3-small": 1536,
    "text-embedding-3-large": 3072,
}


class OpenAIEmbedder(Embedder):
    """OpenAI text-embedding-3-small (1536 dims, $0.02/MTok)."""

    provider_name = "openai"

    def __init__(self, api_key: str, model: str = "text-embedding-3-small", base_url: str | None = None):
        import openai

        self.client = openai.AsyncOpenAI(
            api_key=api_key,
            base_url=base_url or None,
        )
        self.model = model
        self.dimensions = _OPENAI_MODEL_DIMS.get(model, 768 if base_url and "11434" in base_url else 1536)

    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Embed texts in batches of _EMBED_BATCH_SIZE."""
        if not texts:
            return []
        all_embeddings: list[list[float]] = []
        for i in range(0, len(texts), _EMBED_BATCH_SIZE):
            batch = texts[i : i + _EMBED_BATCH_SIZE]
            response = await self.client.embeddings.create(
                input=batch,
                model=self.model,
            )
            all_embeddings.extend(d.embedding for d in response.data)
            logger.debug("OpenAI embedded batch %d-%d (%d texts)", i, i + len(batch), len(batch))
        return all_embeddings


def _collect_provider_keys(config) -> dict[str, dict]:
    """Extract all available provider API keys from config."""
    providers: dict[str, dict] = {}

    for pc in getattr(config, "providers", []):
        if pc.provider not in providers:
            providers[pc.provider] = {
                "api_key": pc.api_key,
                "base_url": getattr(pc, "base_url", ""),
            }

    provider_type = getattr(config, "provider", "")
    if provider_type and provider_type not in providers:
        api_key = getattr(config, "api_key", "")
        if api_key:
            providers[provider_type] = {
                "api_key": api_key,
                "base_url": getattr(config, "base_url", ""),
            }

    return providers


def create_embedder(config) -> Embedder | None:
    """Auto-detect best available embedder with fallback chain.

    Chain: FastEmbed (CPU, for Ollama) → Gemini (free) → OpenAI → None (FTS-only).

    Soft failures (missing key, unsupported provider) skip to next.
    Returns None if no compatible provider found — RAG falls back to FTS-only.
    """
    providers = _collect_provider_keys(config)
    errors: list[str] = []

    # Priority 0: FastEmbed (CPU) — best for Ollama/local setups (no VRAM conflict)
    is_ollama = any(
        "ollama" in info.get("api_key", "").lower() or "11434" in info.get("base_url", "")
        for info in providers.values()
    )
    if is_ollama:
        try:
            embedder = FastEmbedEmbedder()
            logger.info("RAG embedder: using FastEmbed CPU (no VRAM conflict with Ollama)")
            return embedder
        except ImportError:
            errors.append("fastembed: not installed (pip install fastembed)")
            logger.info("FastEmbed not installed — falling back to Ollama embedding")
        except Exception as e:
            errors.append(f"fastembed: {e}")
            logger.warning("FastEmbed init failed: %s — trying next", e)

    # Priority 1: Gemini (free embedding tier)
    if info := providers.get("gemini"):
        if info.get("api_key"):
            try:
                embedder = GeminiEmbedder(
                    api_key=info["api_key"],
                    base_url=info.get("base_url") or None,
                )
                logger.info("RAG embedder: using Gemini gemini-embedding-001 (free tier)")
                return embedder
            except Exception as e:
                errors.append(f"gemini: {e}")
                logger.warning("Gemini embedder init failed: %s — trying next", e)
        else:
            errors.append("gemini: empty API key")

    # Priority 2: OpenAI (or Ollama via OpenAI-compatible API)
    if info := providers.get("openai"):
        if info.get("api_key"):
            base_url = info.get("base_url", "")
            via_ollama = "ollama" in info.get("api_key", "").lower() or "11434" in base_url
            try:
                model = "nomic-embed-text" if via_ollama else "text-embedding-3-small"
                embedder = OpenAIEmbedder(
                    api_key=info["api_key"],
                    model=model,
                    base_url=base_url or None,
                )
                label = f"Ollama {model}" if via_ollama else f"OpenAI {model}"
                logger.info("RAG embedder: using %s", label)
                return embedder
            except Exception as e:
                errors.append(f"openai: {e}")
                logger.warning("OpenAI embedder init failed: %s — trying next", e)
        else:
            errors.append("openai: empty API key")

    # Fallback: FTS-only mode (no vector search)
    if errors:
        logger.warning(
            "RAG embedder: all providers failed (%s). Falling back to FTS-only mode.",
            "; ".join(errors),
        )
    else:
        logger.warning(
            "RAG embedder: no compatible provider found (available: %s). "
            "RAG will use FTS-only keyword search. "
            "Add a Gemini or OpenAI provider for vector search.",
            list(providers.keys()),
        )
    return None
