from dataclasses import dataclass

import numpy as np
from sentence_transformers.util import cos_sim

from kensakun.encoder import Encoder


@dataclass
class Engine:
    encoder: Encoder
    top_n: int
    documents: list[str]
    document_embeddings: np.ndarray

    @classmethod
    def init(cls, encoder: Encoder, top_n: int) -> "Engine":
        return cls(encoder, top_n, [], np.ndarray([]))

    def add_documents(self, documents: list[str]):
        embs = self.encoder.encode(documents, is_query=False)
        if self.document_embeddings.shape == ():
            self.document_embeddings = embs
        else:
            self.document_embeddings = np.concatenate((self.document_embeddings, embs))
        self.documents += documents

    def search(self, query: str) -> tuple[list[str], list[int]]:
        emb = self.encoder.encode([query], is_query=True)
        sims = cos_sim(emb, self.document_embeddings).numpy().squeeze(0)
        indexes = sims.squeeze().argsort()[::-1][: self.top_n]
        sorted_docs = np.array(self.documents)[indexes].tolist()
        return sorted_docs, indexes.tolist()
