from typing import Any, Dict, List

from haystack.schema import Document
from fastrag.rankers import QuantizedBiEncoderRanker


class EndpointHandler:
    def __init__(self, path=""):
        model_id = "Intel/bge-large-en-v1.5-rag-int8-static"
        self.ranker = QuantizedBiEncoderRanker(model_name_or_path=model_id)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        query = data.get("query", None)
        queries = data.get("queries", None)
        documents = data.get("documents", None)
        batch_size = data.get("batch_size", None)
        top_k = data.get("top_k", None)

        if query is not None:
            assert isinstance(query, str), "Expected query to be a string"
            assert isinstance(documents, list), "Expected documents to be a list"
            assert all(
                isinstance(d, dict) for d in documents
            ), "Expected each document in documents to be a dictionary"
            documents = [Document.from_dict(d) for d in documents]
            return self.ranker.predict(query=query, documents=documents, top_k=top_k)

        elif queries is not None:
            assert isinstance(queries, list), "Expected queries to be a list"
            assert all(
                isinstance(query, str) for query in queries
            ), "Expected each query in queries to be a string"
            assert isinstance(documents, list), "Expected documents to be a list"
            assert all(
                all(isinstance(d, dict) for d in doc) for doc in documents
            ), "Expected each document in list of documents to be a dictionary"
            documents = [Document.from_dict(d) for d in documents]
            return self.ranker.predict_batch(
                queries=queries, documents=documents, batch_size=batch_size, top_k=top_k
            )

        else:
            raise ValueError("Expected either query or queries")