from .embeddings_type import EmbeddingsType
from ..chunkers.chunk_strategy import ChunkStrategy


class Embeddings:
    def __init__(self, embeddings_type: EmbeddingsType = EmbeddingsType.OPEN_AI,
                 chunk_size: int = 256,
                 chunk_overlap: int = 128,
                 chunk_strategy: ChunkStrategy = ChunkStrategy.EXACT,
                 docker_image: str = None,
                 hugging_face_model_name: str = None):
        self.embeddings_type = embeddings_type
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.chunk_strategy = chunk_strategy
        self.docker_image = docker_image
        self.hugging_face_model_name = hugging_face_model_name

    def serialize(self):
        data = {
            'embeddings_type': self.embeddings_type.name if self.embeddings_type else None,
            'chunk_size': self.chunk_size,
            'chunk_overlap': self.chunk_overlap,
            'chunk_strategy': self.chunk_strategy.name if self.chunk_strategy else None,
            'docker_image': self.docker_image,
            'hugging_face_model_name': self.hugging_face_model_name
        }

        return {k: v for k, v in data.items() if v is not None}