from typing import Optional

from openkosmos_ai.common.model import RerankResult, EmbeddingResult, RelevanceResult
from openkosmos_ai.nodes.api import ApiConfig, ApiNode


class ModelGatewayConfig(ApiConfig):
    rerank_model: Optional[str] = "bge-reranker-v2-m3"
    embedding_model: Optional[str] = "bge-m3"
    api_prefix: str = "/v1"


class ModelGatewayNode(ApiNode[ModelGatewayConfig]):

    def __init__(self, config: ModelGatewayConfig):
        super().__init__(config)

    def rerank(self, query: str, documents: list[str], top_n=3) -> RerankResult:
        result = self.post("/rerank",
                           json_data={
                               "model": self.config().rerank_model,
                               "query": query,
                               "documents": documents,
                               "top_n": top_n
                           })

        return RerankResult(results=[RelevanceResult(index=r["index"],
                                                     relevance_score=r["relevance_score"])
                                     for r in result["results"]], metadata=result["meta"])

    def embeddings(self, input: str) -> EmbeddingResult:
        return EmbeddingResult.model_validate_json(self.post("/embeddings",
                                                             json_data={
                                                                 "model": self.config().embedding_model,
                                                                 "input": input,
                                                                 "encoding_format": "float"
                                                             }))
