from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import OpenAIEmbeddings
from pydantic import PrivateAttr

from typing import Any, ClassVar


class HuggingFaceEmbeddings(HuggingFaceEmbeddings):
    model_version: str

    def __init__(self, **kwargs: Any):
        """Initialize the sentence_transformer."""
        kwargs.setdefault("model_version", "1")
        super().__init__(**kwargs)


class OpenAIMockClient:
    """
    Mock client for OpenAI embeddings. Replicates the behavior of HuggingFaceEmbeddings _client property.
    """

    def __init__(self, dimensions: int):
        self.dimensions = dimensions

    def get_sentence_embedding_dimension(self):
        return self.dimensions


class OctostarOpenAIEmbeddings(OpenAIEmbeddings):
    # Define static variables as class variables
    DEFAULT_MODEL: ClassVar[str] = "text-embedding-3-small"
    DEFAULT_DIMENSIONS: ClassVar[int] = 1536

    model_version: str
    model_name: str
    _client: OpenAIMockClient = PrivateAttr()

    def __init__(self, **kwargs: Any):
        if not kwargs:
            kwargs = {}
        kwargs.setdefault("model_version", "1")
        kwargs.setdefault("model", self.DEFAULT_MODEL)
        kwargs.setdefault("model_name", self.DEFAULT_MODEL)
        kwargs.setdefault("dimensions", self.DEFAULT_DIMENSIONS)

        super().__init__(**kwargs)
        # To mock the HuggingFaceEmbeddings _client property that we somehow seem to be relying on
        self._client = OpenAIMockClient(
            kwargs.get("dimensions", self.DEFAULT_DIMENSIONS)
        )
