import duckdb
import lancedb

from src.setting import AVAILABLE_WORDS


class VectorDatabaseHandler:
    QUERY_TEMPLATE = "SELECT word, vector FROM {table_name} WHERE word = '{user_word}'"

    def __init__(self, db_path: str, table_name: str, metrics_cfg: dict):
        db = lancedb.connect(db_path)

        self.metrics_cfg = metrics_cfg
        self.embeddings_tbl = db.open_table(table_name)

    def __call__(self, guessed_word: str, supposed_word: str) -> dict:
        arrow_table = self.embeddings_tbl.to_arrow()
        word_embedding = self.get_word_vector(guessed_word, "arrow_table")

        df_emb = self.embeddings_tbl.search(word_embedding) \
            .metric(self.metrics_cfg.metric) \
            .limit(len(AVAILABLE_WORDS)) \
            .to_df()

        supposed_word_row = df_emb[df_emb['word'] == supposed_word].iloc[0]
        cosine_distance = supposed_word_row['_distance']

        words_between_count = len(df_emb[df_emb['_distance'] < cosine_distance])
        closest_word = df_emb[df_emb['word'] != guessed_word].iloc[0]['word'] if words_between_count else supposed_word

        return {
            "score": cosine_distance,
            "rating": words_between_count,
            "percentage": 100 - words_between_count / len(df_emb) * 100,
            "closest_word": closest_word
        }

    def get_word_vector(self, word: str, table_name: str):
        vector = duckdb.query(
            self.QUERY_TEMPLATE.format(table_name=table_name, user_word=word)
        ).to_df()["vector"].values[0]
        return vector