from __future__ import annotations

import os
import time
from typing import Iterable

from dotenv import load_dotenv
from huggingface_hub import InferenceClient
from langchain_core.embeddings import Embeddings
import numpy as np
from sklearn.feature_extraction.text import HashingVectorizer


DEFAULT_EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
DEFAULT_EMBED_PROVIDER = "local"
DEFAULT_EMBED_BATCH_SIZE = 16
DEFAULT_EMBED_RETRIES = 3

load_dotenv()


def get_hf_token() -> str | None:
    return os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACEHUB_API_TOKEN")


def to_single_vector(result) -> list[float]:
    array = np.asarray(result, dtype=np.float32)
    if array.ndim == 1:
        vector = array
    elif array.ndim == 2:
        vector = array.mean(axis=0)
    elif array.ndim == 3:
        vector = array[0].mean(axis=0)
    else:
        raise RuntimeError(f"Unexpected embedding shape: {array.shape}")

    norm = np.linalg.norm(vector)
    if norm:
        vector = vector / norm
    return vector.astype(float).tolist()


class HFTextEmbeddings(Embeddings):
    def __init__(self, model: str | None = None, provider: str | None = None) -> None:
        self.model = model or os.getenv("HF_EMBED_MODEL", DEFAULT_EMBED_MODEL)
        self.provider = provider or os.getenv("HF_EMBED_PROVIDER", DEFAULT_EMBED_PROVIDER)
        self.batch_size = int(os.getenv("HF_EMBED_BATCH_SIZE", str(DEFAULT_EMBED_BATCH_SIZE)))
        self.retries = int(os.getenv("HF_EMBED_RETRIES", str(DEFAULT_EMBED_RETRIES)))
        self.client = None
        self.local_vectorizer = None
        if self.provider == "local":
            # A stateless vectorizer keeps indexing and querying compatible without
            # downloading a model or consuming Hugging Face inference credits.
            self.local_vectorizer = HashingVectorizer(
                n_features=384,
                analyzer="char_wb",
                ngram_range=(3, 5),
                lowercase=True,
                alternate_sign=False,
                norm="l2",
            )
        else:
            self.client = InferenceClient(provider=self.provider, api_key=get_hf_token())

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        return self._embed(texts)

    def embed_query(self, text: str) -> list[float]:
        return self._embed([text])[0]

    def _embed(self, texts: Iterable[str]) -> list[list[float]]:
        inputs = list(texts)
        if not inputs:
            return []

        if self.local_vectorizer is not None:
            print(f"Local embeddings payload: inputs={len(inputs)} dimensions=384", flush=True)
            return self.local_vectorizer.transform(inputs).toarray().astype(float).tolist()

        vectors: list[list[float]] = []
        for start in range(0, len(inputs), self.batch_size):
            batch = inputs[start : start + self.batch_size]
            print(
                f"Outgoing HF embeddings payload: model={self.model} inputs={len(batch)} "
                f"offset={start}",
                flush=True,
            )
            for text in batch:
                vectors.append(self._embed_one(text))
        return vectors

    def _embed_one(self, text: str) -> list[float]:
        last_error: Exception | None = None
        for attempt in range(1, self.retries + 1):
            try:
                return to_single_vector(
                    self.client.feature_extraction(
                        text,
                        model=self.model,
                        normalize=True,
                        truncate=True,
                    )
                )
            except Exception as exc:
                last_error = exc
                if attempt >= self.retries:
                    break
                time.sleep(min(2 * attempt, 5))

        raise RuntimeError(f"HF embeddings failed after {self.retries} attempts: {last_error}") from last_error