#!/usr/bin/env python3
import rospy
import numpy as np

from lasr_vector_databases_msgs.srv import (
    TxtQueryRequest,
    TxtQueryResponse,
    TxtQuery,
)
from lasr_vector_databases_faiss import (
    load_model,
    parse_txt_file,
    get_sentence_embeddings,
    load_vector_database,
    query_database,
)


class TxtQueryService:
    def __init__(self):
        rospy.init_node("txt_query_service")
        self._sentence_embedding_model = load_model()
        rospy.Service("lasr_faiss/txt_query", TxtQuery, self.execute_cb)
        rospy.loginfo("Text Query service started")

    def execute_cb(self, req: TxtQueryRequest) -> TxtQueryResponse:
        txt_fp: str = req.txt_path
        index_path: str = req.index_path
        query_sentence: str = req.query_sentence
        possible_matches: list[str] = parse_txt_file(txt_fp)
        query_embedding: np.ndarray = get_sentence_embeddings(
            [query_sentence], self._sentence_embedding_model  # requires list of strings
        )
        distances, indices = query_database(index_path, query_embedding, k=req.k)
        nearest_matches = [possible_matches[i] for i in indices[0]]

        return TxtQueryResponse(
            closest_sentences=nearest_matches,
            cosine_similarities=distances[0].tolist(),
        )


if __name__ == "__main__":
    TxtQueryService()
    rospy.spin()
