# get_embedding.py
from __future__ import annotations

import asyncio
import os
import yaml
import torch
from typing import List, Optional

from huggingface_hub import snapshot_download
from langchain_huggingface import HuggingFaceEmbeddings

# If you have an auth helper, we keep it
from login import HuggingFaceLogin


class EmbeddingFetcher:
    """
    Async-friendly wrapper around a HuggingFace embedding model.

    - Lazily initializes the model and downloads repo snapshot at app startup.
    - Runs blocking HF / Torch operations in a worker thread via asyncio.to_thread.
    - Thread-safe through an asyncio.Lock to prevent duplicate initialization.
    """

    def __init__(self, config_path: str = "config.yml") -> None:
        self._config_path = config_path
        self._ready = False
        self._init_lock = asyncio.Lock()

        self._model: Optional[HuggingFaceEmbeddings] = None
        self._local_model_path: Optional[str] = None

        # Authenticate early (likely blocking), but off main thread in ensure_ready()
        self._login = HuggingFaceLogin()

        # Config defaults (overridden by config.yml)
        self._repo_id: str = "SPAL0028/default-model"
        self._normalize_embeddings: bool = False

        # Device
        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # -----------------------
    # Public properties
    # -----------------------
    @property
    def model_id(self) -> str:
        return self._local_model_path or self._repo_id

    @property
    def device_str(self) -> str:
        return str(self._device)

    # -----------------------
    # Initialization
    # -----------------------
    async def ensure_ready(self) -> None:
        """
        Idempotent async initializer. Safe to call multiple times.
        """
        if self._ready:
            return

        async with self._init_lock:
            if self._ready:
                return

            # 1) Login (may prompt environment-based auth); keep this off the main loop
            await asyncio.to_thread(self._login.authenticate)

            # 2) Load config
            cfg = await asyncio.to_thread(self._load_config)
            self._repo_id = cfg.get("hugging_face_url", self._repo_id)
            self._normalize_embeddings = bool(cfg.get("normalize_embeddings", self._normalize_embeddings))

            # 3) Download snapshot if needed (blocking IO) off main thread
            self._local_model_path = await asyncio.to_thread(
                snapshot_download,
                self._repo_id
            )

            # 4) Build embeddings (blocking CPU-bound creation) off main thread
            def _build_embeddings():
                return HuggingFaceEmbeddings(
                    model_name=self._local_model_path,
                    model_kwargs={"device": self._device},
                    encode_kwargs={"normalize_embeddings": self._normalize_embeddings},
                )

            self._model = await asyncio.to_thread(_build_embeddings)

            self._ready = True

    # -----------------------
    # Core API
    # -----------------------
    async def embed(self, texts: List[str] | str) -> List[List[float]]:
        """
        Generate embeddings for a list of texts.
        Ensures the model is initialized and runs blocking ops using a thread.
        """
        await self.ensure_ready()

        if isinstance(texts, str):
            texts = [texts]

        if not texts:
            raise ValueError("No texts provided for embedding.")

        # Defensive strip to avoid empty items sneaking through
        sanitized = [t if isinstance(t, str) else str(t) for t in texts]
        sanitized = [t.strip() for t in sanitized if t and t.strip()]

        if not sanitized:
            raise ValueError("All provided texts are empty after sanitization.")

        # langchain_huggingface.HuggingFaceEmbeddings.embed_documents is blocking
        vectors = await asyncio.to_thread(self._model.embed_documents, sanitized)  # type: ignore[union-attr]
        return vectors

    # -----------------------
    # Helpers
    # -----------------------
    def _load_config(self) -> dict:
        if not os.path.exists(self._config_path):
            # Keep behavior predictable if config is missing
            return {}
        with open(self._config_path, "r", encoding="utf-8") as f:
            return yaml.safe_load(f) or {}