"""
RBSP Embedding Pipeline.

Provider-agnostic embedding abstraction with document chunking.  Supports
pluggable backends (OpenAI, Cohere, Ollama, ONNX) with optional dependencies
guarded by lazy imports.  Ships a zero-dependency ``DummyProvider`` for testing.

Architecture:
    - ``EmbeddingConfig``    Configuration dataclass (follows HNSWConfig pattern)
    - ``EmbeddingProvider``  Abstract base class for all embedding backends
    - ``DummyProvider``      Deterministic random vectors (zero deps, for testing)
    - ``OpenAIProvider``     OpenAI text-embedding-ada-002 / text-embedding-3-*
    - ``CohereProvider``     Cohere embed-v3 models
    - ``OllamaProvider``     Local Ollama server embeddings
    - ``EmbeddingPipeline``  Combines chunker + provider into a single vectorizer

Document chunking is provided by ``rbsp.core.chunker.DocumentChunker``.

Thread-safe.  Zero required external dependencies.
"""

from __future__ import annotations

import hashlib
import math
import threading
from abc import ABC, abstractmethod
from typing import Dict, List, Optional, Sequence, Tuple

from rbsp.core.chunker import Chunk, ChunkResult, DocumentChunker


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


class EmbeddingConfig:
    """Configuration for the embedding pipeline.

    Follows the HNSWConfig pattern: __slots__, keyword-only __init__,
    validation in constructor, __repr__.

    Args:
        provider:       Provider name -- ``"dummy"``, ``"openai"``, ``"cohere"``,
                        or ``"ollama"`` (default ``"dummy"``).
        model:          Model identifier (provider-specific).  Defaults vary
                        per provider.
        dim:            Embedding dimensionality.  If ``None``, inferred from
                        the first call to ``embed()``.
        batch_size:     Maximum texts per ``embed_batch()`` call before
                        internal batching (default 64).
        chunk_size:     Target tokens per chunk for ``DocumentChunker``
                        (default 384).
        overlap:        Token overlap between consecutive chunks (default 64).
        strategy:       Chunking strategy -- ``"fixed"``, ``"sentence"``,
                        or ``"paragraph"`` (default ``"sentence"``).
        min_chunk_size: Minimum tokens for a trailing chunk before merging
                        into the previous chunk (default 32).
        api_key:        API key for cloud providers (optional).
        base_url:       Base URL override (e.g. for Ollama, default
                        ``"http://localhost:11434"``).
        normalize:      Whether to L2-normalize output vectors (default True).
        seed:           RNG seed for DummyProvider reproducibility (default 42).
    """

    __slots__ = (
        "provider",
        "model",
        "dim",
        "batch_size",
        "chunk_size",
        "overlap",
        "strategy",
        "min_chunk_size",
        "api_key",
        "base_url",
        "normalize",
        "seed",
    )

    _VALID_PROVIDERS = ("dummy", "openai", "cohere", "ollama")
    _VALID_STRATEGIES = ("fixed", "sentence", "paragraph")

    def __init__(
        self,
        *,
        provider: str = "dummy",
        model: Optional[str] = None,
        dim: Optional[int] = None,
        batch_size: int = 64,
        chunk_size: int = 384,
        overlap: int = 64,
        strategy: str = "sentence",
        min_chunk_size: int = 32,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        normalize: bool = True,
        seed: int = 42,
    ) -> None:
        if provider not in self._VALID_PROVIDERS:
            raise ValueError(
                f"provider must be 'dummy', 'openai', 'cohere', or 'ollama', "
                f"got {provider!r}"
            )
        if batch_size < 1:
            raise ValueError("batch_size must be >= 1")
        if chunk_size < 1:
            raise ValueError("chunk_size must be >= 1")
        if overlap < 0:
            raise ValueError("overlap must be >= 0")
        if overlap >= chunk_size:
            raise ValueError("overlap must be < chunk_size")
        if strategy not in self._VALID_STRATEGIES:
            raise ValueError(
                f"strategy must be 'fixed', 'sentence', or 'paragraph', "
                f"got {strategy!r}"
            )
        if min_chunk_size < 1:
            raise ValueError("min_chunk_size must be >= 1")
        self.provider = provider
        self.model = model
        self.dim = dim
        self.batch_size = batch_size
        self.chunk_size = chunk_size
        self.overlap = overlap
        self.strategy = strategy
        self.min_chunk_size = min_chunk_size
        self.api_key = api_key
        self.base_url = base_url
        self.normalize = normalize
        self.seed = seed

    def __repr__(self) -> str:
        return (
            f"EmbeddingConfig(provider={self.provider!r}, model={self.model!r}, "
            f"dim={self.dim}, chunk_size={self.chunk_size}, "
            f"strategy={self.strategy!r})"
        )


# ---------------------------------------------------------------------------
# Abstract provider
# ---------------------------------------------------------------------------


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

    Subclasses MUST implement ``embed()`` and the ``dim`` property.
    ``embed_batch()`` has a default implementation that loops over ``embed()``.

    Thread-safety: implementations SHOULD be thread-safe.
    """

    @abstractmethod
    def embed(self, text: str) -> List[float]:
        """Embed a single text string.

        Args:
            text: Input text to embed.

        Returns:
            Float vector of length ``self.dim``.
        """

    def embed_batch(self, texts: Sequence[str]) -> List[List[float]]:
        """Embed multiple texts.

        Default implementation calls ``embed()`` in a loop.
        Providers with native batch APIs SHOULD override this for efficiency.

        Args:
            texts: Sequence of input texts.

        Returns:
            List of float vectors, one per input text.
        """
        return [self.embed(t) for t in texts]

    @property
    @abstractmethod
    def dim(self) -> int:
        """Dimensionality of the embedding vectors."""


# ---------------------------------------------------------------------------
# DummyProvider -- zero external dependencies, deterministic
# ---------------------------------------------------------------------------


class DummyProvider(EmbeddingProvider):
    """Deterministic pseudo-random embedding provider for testing.

    Generates reproducible vectors by hashing the input text with BLAKE2b
    (matching RBSP's ``rbsp_hash`` convention) and using the digest bytes
    to seed a deterministic float sequence.  Zero external dependencies.

    Args:
        dim:       Vector dimensionality (default 128).
        seed:      Base seed mixed into the hash (default 42).
        normalize: L2-normalize output vectors (default True).
    """

    def __init__(
        self, dim: int = 128, *, seed: int = 42, normalize: bool = True
    ) -> None:
        self._dim = dim
        self._seed = seed
        self._normalize = normalize

    def embed(self, text: str) -> List[float]:
        """Generate a deterministic pseudo-random vector from *text*.

        Same text always produces the same vector (given same seed + dim).

        Args:
            text: Input text to embed.

        Returns:
            Float vector of length ``self.dim``.
        """
        # Hash text with seed to get deterministic bytes
        key = f"{self._seed}:{text}".encode("utf-8")
        # Generate enough bytes: 4 bytes per float
        needed = self._dim * 4
        chunks: List[bytes] = []
        total = 0
        i = 0
        while total < needed:
            h = hashlib.blake2b(key + i.to_bytes(4, "little"), digest_size=64)
            chunks.append(h.digest())
            total += 64
            i += 1
        raw = b"".join(chunks)

        # Convert bytes to floats in [-1, 1]
        vec: List[float] = []
        for j in range(self._dim):
            # Take 4 bytes, interpret as uint32, map to [-1, 1]
            b = int.from_bytes(raw[j * 4 : j * 4 + 4], "little")
            vec.append(b / 0xFFFFFFFF * 2.0 - 1.0)

        if self._normalize:
            vec = _l2_normalize(vec)
        return vec

    @property
    def dim(self) -> int:
        """Dimensionality of the embedding vectors."""
        return self._dim

    def __repr__(self) -> str:
        return f"DummyProvider(dim={self._dim}, seed={self._seed})"


# ---------------------------------------------------------------------------
# OpenAIProvider -- optional dependency: ``openai``
# ---------------------------------------------------------------------------


class OpenAIProvider(EmbeddingProvider):
    """OpenAI embedding provider.

    Requires the ``openai`` package (``pip install openai``).
    Import is guarded -- the class can be *defined* without the dependency
    installed; it raises ``ImportError`` at construction time if missing.

    Args:
        api_key:   OpenAI API key (or set ``OPENAI_API_KEY`` env var).
        model:     Model name (default ``"text-embedding-3-small"``).
        dim:       Output dimensionality (for models that support it).
                   ``None`` uses the model's native dimension.
        normalize: L2-normalize vectors (default True).
    """

    def __init__(
        self,
        *,
        api_key: Optional[str] = None,
        model: str = "text-embedding-3-small",
        dim: Optional[int] = None,
        normalize: bool = True,
    ) -> None:
        try:
            import openai  # noqa: F401
        except ImportError:
            raise ImportError(
                "OpenAIProvider requires the 'openai' package. "
                "Install it with: pip install openai"
            )
        self._model = model
        self._dim_override = dim
        self._normalize = normalize
        self._client = openai.OpenAI(api_key=api_key)
        # Resolve actual dimension from first call if not overridden
        self._resolved_dim: Optional[int] = dim

    def embed(self, text: str) -> List[float]:
        """Embed a single text via the OpenAI API.

        Args:
            text: Input text to embed.

        Returns:
            Float vector.
        """
        kwargs: Dict = {"input": [text], "model": self._model}
        if self._dim_override is not None:
            kwargs["dimensions"] = self._dim_override
        response = self._client.embeddings.create(**kwargs)
        vec = list(response.data[0].embedding)
        if self._resolved_dim is None:
            self._resolved_dim = len(vec)
        if self._normalize:
            vec = _l2_normalize(vec)
        return vec

    def embed_batch(self, texts: Sequence[str]) -> List[List[float]]:
        """Embed multiple texts in a single OpenAI API call.

        Args:
            texts: Sequence of input texts.

        Returns:
            List of float vectors, one per input text.
        """
        if not texts:
            return []
        kwargs: Dict = {"input": list(texts), "model": self._model}
        if self._dim_override is not None:
            kwargs["dimensions"] = self._dim_override
        response = self._client.embeddings.create(**kwargs)
        results: List[List[float]] = []
        for item in sorted(response.data, key=lambda x: x.index):
            vec = list(item.embedding)
            if self._resolved_dim is None:
                self._resolved_dim = len(vec)
            if self._normalize:
                vec = _l2_normalize(vec)
            results.append(vec)
        return results

    @property
    def dim(self) -> int:
        """Dimensionality of the embedding vectors."""
        if self._resolved_dim is not None:
            return self._resolved_dim
        # Default dimensions for known models
        defaults = {
            "text-embedding-3-small": 1536,
            "text-embedding-3-large": 3072,
            "text-embedding-ada-002": 1536,
        }
        return defaults.get(self._model, 1536)

    def __repr__(self) -> str:
        return f"OpenAIProvider(model={self._model!r}, dim={self.dim})"


# ---------------------------------------------------------------------------
# CohereProvider -- optional dependency: ``cohere``
# ---------------------------------------------------------------------------


class CohereProvider(EmbeddingProvider):
    """Cohere embedding provider.

    Requires the ``cohere`` package (``pip install cohere``).
    Import guarded at construction time.

    Args:
        api_key:    Cohere API key (or set ``CO_API_KEY`` env var).
        model:      Model name (default ``"embed-english-v3.0"``).
        input_type: Cohere input type -- ``"search_document"`` or
                    ``"search_query"`` (default ``"search_document"``).
        normalize:  L2-normalize vectors (default True).
    """

    def __init__(
        self,
        *,
        api_key: Optional[str] = None,
        model: str = "embed-english-v3.0",
        input_type: str = "search_document",
        normalize: bool = True,
    ) -> None:
        try:
            import cohere  # noqa: F401
        except ImportError:
            raise ImportError(
                "CohereProvider requires the 'cohere' package. "
                "Install it with: pip install cohere"
            )
        self._model = model
        self._input_type = input_type
        self._normalize = normalize
        self._client = cohere.Client(api_key=api_key)
        self._resolved_dim: Optional[int] = None

    def embed(self, text: str) -> List[float]:
        """Embed a single text via the Cohere API.

        Args:
            text: Input text to embed.

        Returns:
            Float vector.
        """
        response = self._client.embed(
            texts=[text],
            model=self._model,
            input_type=self._input_type,
        )
        vec = list(response.embeddings[0])
        if self._resolved_dim is None:
            self._resolved_dim = len(vec)
        if self._normalize:
            vec = _l2_normalize(vec)
        return vec

    def embed_batch(self, texts: Sequence[str]) -> List[List[float]]:
        """Embed multiple texts in a single Cohere API call.

        Args:
            texts: Sequence of input texts.

        Returns:
            List of float vectors, one per input text.
        """
        if not texts:
            return []
        response = self._client.embed(
            texts=list(texts),
            model=self._model,
            input_type=self._input_type,
        )
        results: List[List[float]] = []
        for emb in response.embeddings:
            vec = list(emb)
            if self._resolved_dim is None:
                self._resolved_dim = len(vec)
            if self._normalize:
                vec = _l2_normalize(vec)
            results.append(vec)
        return results

    @property
    def dim(self) -> int:
        """Dimensionality of the embedding vectors."""
        if self._resolved_dim is not None:
            return self._resolved_dim
        defaults = {
            "embed-english-v3.0": 1024,
            "embed-multilingual-v3.0": 1024,
            "embed-english-light-v3.0": 384,
            "embed-multilingual-light-v3.0": 384,
        }
        return defaults.get(self._model, 1024)

    def __repr__(self) -> str:
        return f"CohereProvider(model={self._model!r}, dim={self.dim})"


# ---------------------------------------------------------------------------
# OllamaProvider -- optional dependency: ``urllib.request`` (stdlib only!)
# ---------------------------------------------------------------------------


class OllamaProvider(EmbeddingProvider):
    """Ollama local embedding provider.

    Communicates with a local Ollama server via HTTP.  Uses only stdlib
    ``urllib.request`` + ``json`` -- NO external dependencies.

    Args:
        model:    Ollama model name (default ``"nomic-embed-text"``).
        base_url: Ollama server URL (default ``"http://localhost:11434"``).
        dim:      Expected dimensionality (inferred on first call if None).
        normalize: L2-normalize vectors (default True).
    """

    def __init__(
        self,
        *,
        model: str = "nomic-embed-text",
        base_url: str = "http://localhost:11434",
        dim: Optional[int] = None,
        normalize: bool = True,
    ) -> None:
        self._model = model
        self._base_url = base_url.rstrip("/")
        self._normalize = normalize
        self._resolved_dim: Optional[int] = dim

    def embed(self, text: str) -> List[float]:
        """Embed a single text via the Ollama HTTP API.

        Args:
            text: Input text to embed.

        Returns:
            Float vector.
        """
        import json
        import urllib.request

        payload = json.dumps({
            "model": self._model,
            "prompt": text,
        }).encode("utf-8")

        req = urllib.request.Request(
            f"{self._base_url}/api/embeddings",
            data=payload,
            headers={"Content-Type": "application/json"},
        )
        with urllib.request.urlopen(req) as resp:
            data = json.loads(resp.read().decode("utf-8"))

        vec = list(data["embedding"])
        if self._resolved_dim is None:
            self._resolved_dim = len(vec)
        if self._normalize:
            vec = _l2_normalize(vec)
        return vec

    @property
    def dim(self) -> int:
        """Dimensionality of the embedding vectors."""
        if self._resolved_dim is not None:
            return self._resolved_dim
        defaults = {
            "nomic-embed-text": 768,
            "all-minilm": 384,
            "mxbai-embed-large": 1024,
        }
        return defaults.get(self._model, 768)

    def __repr__(self) -> str:
        return (
            f"OllamaProvider(model={self._model!r}, "
            f"base_url={self._base_url!r})"
        )


# ---------------------------------------------------------------------------
# EmbeddingPipeline
# ---------------------------------------------------------------------------


class EmbeddingPipeline:
    """Combines a DocumentChunker with an EmbeddingProvider to vectorize documents.

    Thread-safe via internal lock (matches VectorStore/HNSWIndex pattern).

    Args:
        provider: An ``EmbeddingProvider`` instance.
        chunker:  A ``DocumentChunker`` instance.  If ``None``, a default
                  chunker is created with default settings.
    """

    def __init__(
        self,
        provider: EmbeddingProvider,
        chunker: Optional[DocumentChunker] = None,
    ) -> None:
        self._provider = provider
        self._chunker = chunker or DocumentChunker()
        self._lock = threading.Lock()

    @property
    def provider(self) -> EmbeddingProvider:
        """The underlying embedding provider."""
        return self._provider

    @property
    def dim(self) -> int:
        """Embedding dimensionality (delegated to provider)."""
        return self._provider.dim

    def embed_text(self, text: str) -> List[float]:
        """Embed a single text string (no chunking).

        Args:
            text: Input text.

        Returns:
            Float vector of length ``self.dim``.
        """
        with self._lock:
            return self._provider.embed(text)

    def embed_texts(self, texts: Sequence[str]) -> List[List[float]]:
        """Embed multiple texts (no chunking).

        Args:
            texts: Sequence of input texts.

        Returns:
            List of float vectors.
        """
        with self._lock:
            return self._provider.embed_batch(texts)

    def vectorize_document(
        self, text: str, doc_id: str = ""
    ) -> List[Tuple[Chunk, List[float]]]:
        """Chunk a document and embed each chunk.

        Args:
            text:   Full document text.
            doc_id: Document identifier passed to the chunker.

        Returns:
            List of ``(Chunk, vector)`` tuples.  Each Chunk carries
            full metadata (offsets, token count, index).
        """
        result = self._chunker.chunk(text, doc_id=doc_id)
        with self._lock:
            vectors = self._provider.embed_batch(
                [c.text for c in result.chunks]
            )
        return list(zip(result.chunks, vectors))

    def __repr__(self) -> str:
        return (
            f"EmbeddingPipeline(provider={self._provider!r}, "
            f"chunker={self._chunker!r})"
        )


# ---------------------------------------------------------------------------
# Factory
# ---------------------------------------------------------------------------


def create_provider(config: EmbeddingConfig) -> EmbeddingProvider:
    """Create an embedding provider from configuration.

    Args:
        config: An ``EmbeddingConfig`` instance.

    Returns:
        An ``EmbeddingProvider`` subclass instance.

    Raises:
        ValueError: If the provider name is not recognized.
    """
    if config.provider == "dummy":
        return DummyProvider(
            dim=config.dim or 128,
            seed=config.seed,
            normalize=config.normalize,
        )
    if config.provider == "openai":
        return OpenAIProvider(
            api_key=config.api_key,
            model=config.model or "text-embedding-3-small",
            dim=config.dim,
            normalize=config.normalize,
        )
    if config.provider == "cohere":
        return CohereProvider(
            api_key=config.api_key,
            model=config.model or "embed-english-v3.0",
            normalize=config.normalize,
        )
    if config.provider == "ollama":
        return OllamaProvider(
            model=config.model or "nomic-embed-text",
            base_url=config.base_url or "http://localhost:11434",
            dim=config.dim,
            normalize=config.normalize,
        )
    raise ValueError(f"Unknown provider: {config.provider!r}")


def create_pipeline(config: Optional[EmbeddingConfig] = None) -> EmbeddingPipeline:
    """Create a complete embedding pipeline from configuration.

    Args:
        config: An ``EmbeddingConfig`` instance.  If ``None``, uses
                defaults (DummyProvider, dim=128).

    Returns:
        An ``EmbeddingPipeline`` instance ready to use.
    """
    if config is None:
        config = EmbeddingConfig()
    provider = create_provider(config)
    chunker = DocumentChunker(
        chunk_size=config.chunk_size,
        overlap=config.overlap,
        strategy=config.strategy,
        min_chunk_size=config.min_chunk_size,
    )
    return EmbeddingPipeline(provider=provider, chunker=chunker)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _l2_normalize(vec: List[float]) -> List[float]:
    """L2-normalize a vector.  Returns zeros unchanged."""
    norm = math.sqrt(sum(x * x for x in vec))
    if norm == 0.0:
        return vec
    inv = 1.0 / norm
    return [x * inv for x in vec]
