from typing import Any, Dict, List

from haystack import Document

from fastrag.embedders import (
    IPEXSentenceTransformersDocumentEmbedder,
    IPEXSentenceTransformersTextEmbedder,
)


class EndpointHandler:
    def __init__(self):
        model_id = "Intel/bge-small-en-v1.5-rag-int8-static"

        self.query_embedder = IPEXSentenceTransformersTextEmbedder(model_id)
        self.document_embedder = IPEXSentenceTransformersDocumentEmbedder(model_id)

        self.query_embedder.warm_up()
        self.document_embedder.warm_up()

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        text = data.get("text", None)
        if text is not None:
            assert isinstance(text, list), "Expected text to be a string"
            return self.query_embedder.run(text)

        documents = data.get("documents", None)
        if documents is not None:
            assert isinstance(documents, list), "Expected documents to be a list"
            assert all(
                isinstance(document, dict) for document in documents
            ), "Expected each document in documents to be a dictionary"

            documents = [Document.from_dict(document) for document in documents]
            return self.document_embedder.run(documents)

        raise ValueError("Expected either text or documents")