from __future__ import annotations

import os
import re
import threading
from dataclasses import dataclass
from functools import lru_cache
from typing import Any, Iterable, List, Optional, Sequence

from fastembed import TextEmbedding
from langchain_core.embeddings import Embeddings

# ============================
# Configuration (env-friendly)
# ============================

DEFAULT_EMBED_MODEL = os.getenv("ECDA_EMBED_MODEL", "BAAI/bge-small-en-v1.5")

MODEL_PREFERENCE = [
    "BAAI/bge-small-en-v1.5",
    "bge-small",
    "bge-small-en",
    "e5-small",
    "gte-small",
    "all-minilm",
    "minilm",
    "nomic-embed-text-v1.5",
    "nomic-ai/nomic-embed-text-v1.5",
]

# Safer defaults for large corpora (tune via env vars)
DEFAULT_BATCH_SIZE = int(os.getenv("ECDA_EMBED_BATCH_SIZE", "128"))
DEFAULT_THREADS = int(os.getenv("ECDA_EMBED_THREADS", "0"))          # 0 = library default
DEFAULT_MAX_LENGTH = int(os.getenv("ECDA_EMBED_MAX_LENGTH", "384"))
DEFAULT_CACHE_DIR = os.getenv("ECDA_EMBED_CACHE_DIR", "")
DEFAULT_CLEAN_TEXT = os.getenv("ECDA_EMBED_CLEAN_TEXT", "1") != "0"
DEFAULT_HARD_CHAR_LIMIT = int(os.getenv("ECDA_EMBED_HARD_CHAR_LIMIT", "8000"))

# ============================
# Text helpers
# ============================

_ws_re = re.compile(r"[ \t]+")
_many_nl_re = re.compile(r"\n{3,}")


def _clean_text(s: str) -> str:
    s = s.replace("\x00", " ")
    s = _ws_re.sub(" ", s)
    s = _many_nl_re.sub("\n\n", s)
    return s.strip()


def _truncate_hard(s: str, limit: int) -> str:
    if limit > 0 and len(s) > limit:
        return s[:limit]
    return s


def _batched(items: Sequence[str], batch_size: int) -> Iterable[List[str]]:
    bs = max(1, int(batch_size))
    for i in range(0, len(items), bs):
        yield list(items[i : i + bs])


# ============================
# FastEmbed model list robustness
# ============================

def _normalize_supported_models(raw: Any) -> List[str]:
    if raw is None:
        return []
    out: List[str] = []

    if isinstance(raw, (list, tuple)):
        for item in raw:
            if isinstance(item, str):
                name = item.strip()
                if name:
                    out.append(name)
                continue

            if isinstance(item, dict):
                for key in ("model", "model_name", "name", "id"):
                    val = item.get(key)
                    if isinstance(val, str) and val.strip():
                        out.append(val.strip())
                        break
                continue

            for attr in ("model", "model_name", "name", "id"):
                val = getattr(item, attr, None)
                if isinstance(val, str) and val.strip():
                    out.append(val.strip())
                    break

    seen = set()
    deduped: List[str] = []
    for m in out:
        if m in seen:
            continue
        seen.add(m)
        deduped.append(m)
    return deduped


@lru_cache(maxsize=1)
def _supported_models() -> List[str]:
    try:
        raw = TextEmbedding.list_supported_models()
        return _normalize_supported_models(raw)
    except Exception:
        return []


def _pick_model(preferred: Optional[str]) -> str:
    supported = _supported_models()

    if preferred:
        p = preferred.strip()
        if not supported:
            return p

        p_low = p.lower()
        for m in supported:
            if m.lower() == p_low:
                return m
        for m in supported:
            if p_low in m.lower():
                return m

    if supported:
        supported_low = [(m, m.lower()) for m in supported]
        for key in MODEL_PREFERENCE:
            k = key.lower()
            for m, ml in supported_low:
                if k == ml or k in ml:
                    return m
        return supported[0]

    return preferred.strip() if preferred else DEFAULT_EMBED_MODEL


def _build_text_embedding(model_name: str, *, cache_dir: str, threads: int) -> TextEmbedding:
    kwargs = {}
    if cache_dir:
        kwargs["cache_dir"] = cache_dir
    if threads and threads > 0:
        kwargs["threads"] = threads
    return TextEmbedding(model_name=model_name, **kwargs)


# ============================
# High-performance embeddings
# ============================

@dataclass(frozen=True)
class EmbedConfig:
    model_name: str
    batch_size: int = DEFAULT_BATCH_SIZE
    max_length: int = DEFAULT_MAX_LENGTH
    threads: int = DEFAULT_THREADS
    cache_dir: str = DEFAULT_CACHE_DIR
    clean_text: bool = DEFAULT_CLEAN_TEXT
    hard_char_limit: int = DEFAULT_HARD_CHAR_LIMIT


class FastEmbedHPEmbeddings(Embeddings):
    """
    - One model instance per process
    - Explicit batching
    - Optional cleanup + hard truncation
    - Thread-safe calls
    """

    def __init__(self, config: EmbedConfig):
        self.config = config
        self._model = _build_text_embedding(
            config.model_name,
            cache_dir=config.cache_dir,
            threads=config.threads,
        )
        self._lock = threading.Lock()

    def _prep_many(self, texts: Sequence[str]) -> List[str]:
        out: List[str] = []
        clean = self.config.clean_text
        hard_limit = self.config.hard_char_limit
        for t in texts:
            if clean:
                t = _clean_text(t)
            t = _truncate_hard(t, hard_limit)
            out.append(t)
        return out

    def _prep_one(self, text: str) -> str:
        if self.config.clean_text:
            text = _clean_text(text)
        return _truncate_hard(text, self.config.hard_char_limit)

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        texts = self._prep_many(texts)
        vectors: List[List[float]] = []

        with self._lock:
            for batch in _batched(texts, self.config.batch_size):
                try:
                    it = self._model.embed(
                        batch,
                        batch_size=self.config.batch_size,
                        max_length=self.config.max_length,
                    )
                except TypeError:
                    try:
                        it = self._model.embed(batch, batch_size=self.config.batch_size)
                    except TypeError:
                        it = self._model.embed(batch)

                vectors.extend([list(map(float, v)) for v in it])

        return vectors

    def embed_query(self, text: str) -> List[float]:
        text = self._prep_one(text)

        with self._lock:
            if hasattr(self._model, "query_embed"):
                try:
                    it = self._model.query_embed([text], max_length=self.config.max_length)
                except TypeError:
                    it = self._model.query_embed([text])
            else:
                try:
                    it = self._model.embed([text], batch_size=1, max_length=self.config.max_length)
                except TypeError:
                    it = self._model.embed([text])

            vec = next(iter(it))

        return list(map(float, vec))


@lru_cache(maxsize=16)
def make_embeddings(
    *,
    embedding_model: str | None = None,
    batch_size: int = DEFAULT_BATCH_SIZE,
    max_length: int = DEFAULT_MAX_LENGTH,
    threads: int = DEFAULT_THREADS,
    cache_dir: str = DEFAULT_CACHE_DIR,
    clean_text: bool = DEFAULT_CLEAN_TEXT,
    hard_char_limit: int = DEFAULT_HARD_CHAR_LIMIT,
) -> Embeddings:
    chosen = _pick_model(embedding_model or DEFAULT_EMBED_MODEL)
    cfg = EmbedConfig(
        model_name=chosen,
        batch_size=batch_size,
        max_length=max_length,
        threads=threads,
        cache_dir=cache_dir,
        clean_text=clean_text,
        hard_char_limit=hard_char_limit,
    )
    return FastEmbedHPEmbeddings(cfg)