"""
cogniq.embedders — Auto-detecting, zero-to-GPU embedder stack.

Priority chain (auto mode):
  sentence-transformers → OpenAI → Ollama → TF-IDF (zero-dep fallback)
"""

from __future__ import annotations
import re, math, hashlib, logging
from abc import ABC, abstractmethod
from collections import Counter, defaultdict
from functools import lru_cache
from typing import List, Optional, Callable, Union
import numpy as np

log = logging.getLogger("cogniq")

# ─────────────────────────────────────────────
# BASE
# ─────────────────────────────────────────────

class BaseEmbedder(ABC):
    """All embedders must implement embed() and expose .dim"""

    @abstractmethod
    def embed(self, text: str) -> np.ndarray: ...

    def embed_batch(self, texts: List[str]) -> np.ndarray:
        """Batch embed — override for native batch support"""
        return np.vstack([self.embed(t) for t in texts])

    @property
    @abstractmethod
    def dim(self) -> int: ...

    def _normalize(self, vec: np.ndarray) -> np.ndarray:
        n = np.linalg.norm(vec)
        return vec / n if n > 1e-10 else vec

    def __repr__(self):
        return f"{self.__class__.__name__}(dim={self.dim})"


# ─────────────────────────────────────────────
# 1. TF-IDF (ZERO DEPS — always available)
# ─────────────────────────────────────────────

class TFIDFEmbedder(BaseEmbedder):
    """
    Feature-hashed TF-IDF with incremental IDF — works with numpy only.
    Fast, deterministic, surprisingly effective for domain-specific text.
    """

    def __init__(self, dim: int = 1024, ngram_max: int = 2):
        self._dim = dim
        self.ngram_max = ngram_max
        self._df: dict = defaultdict(int)
        self._idf: dict = {}
        self._doc_count = 0

    # ── tokenize ──────────────────────────────
    def _tokenize(self, text: str) -> List[str]:
        text = text.lower()
        text = re.sub(r"[^a-z0-9\s]", " ", text)
        tokens = text.split()
        result = list(tokens)
        for n in range(2, self.ngram_max + 1):
            result += ["_".join(tokens[i:i+n]) for i in range(len(tokens) - n + 1)]
        return result

    def _hash(self, token: str) -> int:
        return int(hashlib.blake2b(token.encode(), digest_size=4).hexdigest(), 16) % self._dim

    # ── fit / partial_fit ─────────────────────
    def fit(self, texts: List[str]) -> "TFIDFEmbedder":
        self._df = defaultdict(int)
        self._doc_count = len(texts)
        for text in texts:
            for tok in set(self._tokenize(text)):
                self._df[tok] += 1
        self._rebuild_idf()
        return self

    def partial_fit(self, text: str) -> "TFIDFEmbedder":
        """Incrementally update IDF — call when adding new docs"""
        self._doc_count += 1
        for tok in set(self._tokenize(text)):
            self._df[tok] += 1
        self._rebuild_idf()
        return self

    def _rebuild_idf(self):
        N = self._doc_count
        self._idf = {
            tok: math.log((N + 1) / (df + 1)) + 1.0
            for tok, df in self._df.items()
        }

    # ── embed ─────────────────────────────────
    def embed(self, text: str) -> np.ndarray:
        tokens = self._tokenize(text)
        if not tokens:
            return np.zeros(self._dim, dtype=np.float32)
        tf = Counter(tokens)
        vec = np.zeros(self._dim, dtype=np.float32)
        inv_len = 1.0 / len(tokens)
        for tok, cnt in tf.items():
            idx = self._hash(tok)
            idf = self._idf.get(tok, math.log(2) + 1.0)
            vec[idx] += cnt * inv_len * idf
        return self._normalize(vec)

    def embed_batch(self, texts: List[str]) -> np.ndarray:
        # Vectorized batch — all tokens hashed in one pass per doc
        return np.vstack([self.embed(t) for t in texts])

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


# ─────────────────────────────────────────────
# 2. SENTENCE-TRANSFORMER (optional)
# ─────────────────────────────────────────────

class SentenceEmbedder(BaseEmbedder):
    """Wraps sentence-transformers — install: pip install sentence-transformers"""

    def __init__(self, model: str = "all-MiniLM-L6-v2"):
        try:
            from sentence_transformers import SentenceTransformer
        except ImportError:
            raise ImportError("pip install sentence-transformers")
        self._model = SentenceTransformer(model)
        self._dim = self._model.get_sentence_embedding_dimension()
        log.info(f"SentenceEmbedder loaded: {model} (dim={self._dim})")

    def embed(self, text: str) -> np.ndarray:
        v = self._model.encode(text, normalize_embeddings=True, show_progress_bar=False)
        return np.array(v, dtype=np.float32)

    def embed_batch(self, texts: List[str]) -> np.ndarray:
        return np.array(
            self._model.encode(texts, normalize_embeddings=True,
                               batch_size=64, show_progress_bar=False),
            dtype=np.float32,
        )

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


# ─────────────────────────────────────────────
# 3. OPENAI (optional)
# ─────────────────────────────────────────────

class OpenAIEmbedder(BaseEmbedder):
    """OpenAI text-embedding-3-small/large"""

    def __init__(self, model: str = "text-embedding-3-small", api_key: Optional[str] = None):
        try:
            import openai
        except ImportError:
            raise ImportError("pip install openai")
        import os
        self._client = openai.OpenAI(api_key=api_key or os.environ["OPENAI_API_KEY"])
        self._model = model
        _dim_map = {"text-embedding-3-small": 1536, "text-embedding-3-large": 3072,
                    "text-embedding-ada-002": 1536}
        self._dim = _dim_map.get(model, 1536)

    def embed(self, text: str) -> np.ndarray:
        r = self._client.embeddings.create(input=[text], model=self._model)
        v = np.array(r.data[0].embedding, dtype=np.float32)
        return self._normalize(v)

    def embed_batch(self, texts: List[str]) -> np.ndarray:
        r = self._client.embeddings.create(input=texts, model=self._model)
        vecs = [np.array(d.embedding, dtype=np.float32) for d in r.data]
        return np.vstack([self._normalize(v) for v in vecs])

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


# ─────────────────────────────────────────────
# 4. OLLAMA (optional — local LLM)
# ─────────────────────────────────────────────

class OllamaEmbedder(BaseEmbedder):
    """Local Ollama embeddings — install: https://ollama.ai"""

    def __init__(self, model: str = "nomic-embed-text", base_url: str = "http://localhost:11434"):
        self._model = model
        self._base_url = base_url.rstrip("/")
        self._dim: Optional[int] = None
        # Probe dim
        try:
            v = self._call("hello")
            self._dim = len(v)
            log.info(f"OllamaEmbedder: {model} (dim={self._dim})")
        except Exception as e:
            raise ConnectionError(f"Ollama not reachable at {base_url}: {e}")

    def _call(self, text: str) -> List[float]:
        import urllib.request, json
        body = json.dumps({"model": self._model, "prompt": text}).encode()
        req = urllib.request.Request(
            f"{self._base_url}/api/embeddings",
            data=body, headers={"Content-Type": "application/json"}, method="POST"
        )
        with urllib.request.urlopen(req, timeout=30) as r:
            return json.loads(r.read())["embedding"]

    def embed(self, text: str) -> np.ndarray:
        v = np.array(self._call(text), dtype=np.float32)
        return self._normalize(v)

    @property
    def dim(self) -> int:
        return self._dim or 768


# ─────────────────────────────────────────────
# 5. CUSTOM (user-provided function)
# ─────────────────────────────────────────────

class CustomEmbedder(BaseEmbedder):
    """
    Wrap any callable as an embedder.
    
    Example:
        def my_embed(text: str) -> list:
            return my_model.encode(text).tolist()
        
        emb = CustomEmbedder(my_embed, dim=768)
    """

    def __init__(self, fn: Callable[[str], Union[list, np.ndarray]], dim: int):
        self._fn = fn
        self._dim = dim

    def embed(self, text: str) -> np.ndarray:
        v = np.array(self._fn(text), dtype=np.float32)
        return self._normalize(v)

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


# ─────────────────────────────────────────────
# 6. AUTO (smart fallback chain)
# ─────────────────────────────────────────────

def AutoEmbedder(
    prefer: str = "auto",
    model: Optional[str] = None,
    api_key: Optional[str] = None,
    ollama_url: str = "http://localhost:11434",
) -> BaseEmbedder:
    """
    Auto-selects the best available embedder.
    
    Priority: sentence-transformers → Ollama → OpenAI → TF-IDF
    
    Args:
        prefer: "auto" | "sentence" | "openai" | "ollama" | "tfidf"
        model: override model name
        api_key: OpenAI API key
        ollama_url: Ollama base URL
    """
    import os

    if prefer == "openai" or (prefer == "auto" and os.environ.get("OPENAI_API_KEY")):
        try:
            return OpenAIEmbedder(model or "text-embedding-3-small", api_key=api_key)
        except Exception as e:
            log.warning(f"OpenAI embedder failed: {e}")

    if prefer in ("sentence", "auto"):
        try:
            return SentenceEmbedder(model or "all-MiniLM-L6-v2")
        except Exception as e:
            log.debug(f"SentenceTransformer not available: {e}")

    if prefer in ("ollama", "auto"):
        try:
            return OllamaEmbedder(model or "nomic-embed-text", base_url=ollama_url)
        except Exception as e:
            log.debug(f"Ollama not available: {e}")

    log.info("Using built-in TF-IDF embedder (zero dependencies)")
    return TFIDFEmbedder(dim=1024)
