"""Embedders for RAG.

Two implementations:
- `HashEmbedder`: offline, deterministic bag-of-words + hash projection. Used for dev/test.
- `OpenAIEmbedder`: calls an OpenAI-compatible embeddings API at runtime.

The active embedder is chosen via `settings.embeddings.provider`.
"""

from __future__ import annotations

import hashlib
import math
import re
from collections import Counter
from typing import Protocol, runtime_checkable

import httpx

from app.core.logging import get_logger

_logger = get_logger(__name__)

_WORD_RE = re.compile(r"\w+", re.UNICODE)


def _tokenize(text: str) -> list[str]:
    return [w.lower() for w in _WORD_RE.findall(text)]


@runtime_checkable
class Embedder(Protocol):
    async def embed(self, texts: list[str]) -> list[list[float]]: ...

    @property
    def dimension(self) -> int: ...


class HashEmbedder:
    """Offline bag-of-words embedder with hash projection.

    Not semantically meaningful, but deterministic and fast — sufficient for
    integration tests and local dev. Cosine similarity is non-zero only when
    texts share tokens.
    """

    def __init__(self, dimension: int = 256):
        if dimension <= 0:
            raise ValueError("dimension must be positive")
        self._dim = dimension

    @property
    def dimension(self) -> int:
        return self._dim

    async def embed(self, texts: list[str]) -> list[list[float]]:
        out: list[list[float]] = []
        for text in texts:
            out.append(self._hash_project(text))
        return out

    def _hash_project(self, text: str) -> list[float]:
        vec = [0.0] * self._dim
        tokens = _tokenize(text)
        if not tokens:
            return vec
        counts = Counter(tokens)
        for token, count in counts.items():
            h = hashlib.md5(token.encode("utf-8")).digest()
            # Use first 4 bytes for index, next 4 bytes for sign
            idx = int.from_bytes(h[:4], "little") % self._dim
            sign = 1.0 if (h[4] & 1) == 0 else -1.0
            vec[idx] += sign * math.sqrt(count)
        # L2 normalize
        norm = math.sqrt(sum(v * v for v in vec))
        if norm > 0:
            vec = [v / norm for v in vec]
        return vec


class OpenAIEmbedder:
    """OpenAI-compatible embeddings API client."""

    def __init__(
        self,
        api_url: str,
        api_key: str,
        model: str,
        dimension: int,
        timeout: float = 30.0,
        batch_size: int = 32,
    ):
        self._api_url = api_url.rstrip("/")
        self._api_key = api_key
        self._model = model
        self._dim = dimension
        self._timeout = timeout
        self._batch_size = batch_size

    @property
    def dimension(self) -> int:
        return self._dim

    async def embed(self, texts: list[str]) -> list[list[float]]:
        if not texts:
            return []
        out: list[list[float]] = []
        async with httpx.AsyncClient(timeout=self._timeout) as client:
            for i in range(0, len(texts), self._batch_size):
                batch = texts[i : i + self._batch_size]
                resp = await client.post(
                    f"{self._api_url}/embeddings",
                    headers={"Authorization": f"Bearer {self._api_key}"},
                    json={"model": self._model, "input": batch},
                )
                resp.raise_for_status()
                data = resp.json()
                # Sort by index to preserve order
                sorted_data = sorted(data["data"], key=lambda x: x["index"])
                out.extend(d["embedding"] for d in sorted_data)
        return out

    async def probe_dimension(self, sample_text: str = "hello world") -> int:
        """Make a single embedding call and return the dimension of the result.

        Useful for the "auto-probe dimension" admin button.
        """
        result = await self.embed([sample_text])
        if not result:
            raise RuntimeError("Empty embeddings response")
        return len(result[0])


def build_hash_embedder(dimension: int) -> HashEmbedder:
    return HashEmbedder(dimension=dimension)


def build_openai_embedder(
    api_url: str,
    api_key: str,
    model: str,
    dimension: int,
    timeout: float = 30.0,
    batch_size: int = 32,
) -> OpenAIEmbedder:
    return OpenAIEmbedder(
        api_url=api_url,
        api_key=api_key,
        model=model,
        dimension=dimension,
        timeout=timeout,
        batch_size=batch_size,
    )