# ============================================================
# Project: EDGAR Research Database
# Description: SEC EDGAR Research Infrastructure
#
# Authors:
# - Manel Kammoun — manel.kammoun@uqo.ca
# - Djerrycharli Tandjambianda — djerrycharli.tandjambianda@uqo.ca
# - Simon-Pierre Boucher — simon-pierre.boucher@uqo.ca
#
# Institution:
# Université du Québec en Outaouais (UQO)
#
# License:
# Research / Academic Use
# ============================================================

from __future__ import annotations

import numpy as np

from src.utils.config import get_embeddings_config
from src.utils.logging import get_logger

logger = get_logger("embeddings.models")


class EmbeddingModelRegistry:

    def __init__(self) -> None:
        self._config = get_embeddings_config()
        self._models: dict[str, dict] = self._config.get("models", {})
        self._default_model: str = self._config.get("default_model", "all-MiniLM-L6-v2")
        self._loaded: dict = {}

    @property
    def default_model(self) -> str:
        return self._default_model

    def list_models(self) -> list[dict]:
        result: list[dict] = []
        for name, info in self._models.items():
            entry = {"name": name, **info}
            entry["is_default"] = name == self._default_model
            result.append(entry)
        return result

    def get_model_info(self, model_name: str) -> dict:
        if model_name not in self._models:
            available = ", ".join(self._models.keys())
            raise ValueError(
                f"Unknown model '{model_name}'. Available: {available}"
            )
        info = dict(self._models[model_name])
        info["name"] = model_name
        return info

    def load_model(self, model_name: str):
        if model_name in self._loaded:
            return self._loaded[model_name]

        info = self.get_model_info(model_name)
        provider = info.get("provider", "")

        if provider == "openai":
            raise NotImplementedError(
                f"OpenAI model '{model_name}' requires API key configuration. "
                "Set OPENAI_API_KEY environment variable and implement the OpenAI "
                "embedding client to use this model."
            )

        if provider != "sentence-transformers":
            raise ValueError(f"Unsupported provider '{provider}' for model '{model_name}'")

        from sentence_transformers import SentenceTransformer

        device = self._config.get("processing", {}).get("device", "cpu")
        logger.info("Loading sentence-transformers model '%s' on device '%s'", model_name, device)
        model = SentenceTransformer(model_name, device=device)
        self._loaded[model_name] = model
        logger.info("Model '%s' loaded successfully", model_name)
        return model

    def encode(
        self,
        texts: list[str],
        model_name: str | None = None,
        batch_size: int = 256,
        show_progress: bool = True,
    ) -> np.ndarray:
        model_name = model_name or self._default_model
        info = self.get_model_info(model_name)
        provider = info.get("provider", "")

        if provider == "openai":
            raise NotImplementedError(
                f"OpenAI model '{model_name}' requires API key configuration. "
                "Set OPENAI_API_KEY environment variable and implement the OpenAI "
                "embedding client to use this model."
            )

        model = self.load_model(model_name)
        normalize = self._config.get("processing", {}).get("normalize", True)

        logger.info(
            "Encoding %d texts with '%s' (batch_size=%d)",
            len(texts), model_name, batch_size,
        )
        vectors = model.encode(
            texts,
            batch_size=batch_size,
            show_progress_bar=show_progress,
            normalize_embeddings=normalize,
        )
        return np.asarray(vectors)