from pathlib import Path

from llama_index.retrievers.bm25 import BM25Retriever



# from main import extract_court_decision_text
# INDEX_NAME = "bm25_retriever_long"
# TEST_CD_URL = "https://reyestr.court.gov.ua/Review/118766467"
# TEST_CD_URL = "https://reyestr.court.gov.ua/Review/118763429"
# query = extract_court_decision_text(TEST_CD_URL)

INDEX_NAME = "bm25_retriever_short"
query = (
    # "Викрадення майна, злочини, пов'язані з порушенням законодавчих норм щодо обігу та використання "
    # "документів, печаток, штампів, бланків, а також спеціальних технічних засобів "
    # "для отримання інформації та комунікаційних ліній."
    "Проникнення в житло, викрадення майна, використання підробного паспорта."
)

PERSIST_PATH = Path("Save_Index_Local")
PRINT_QUERY = True


retriever = BM25Retriever.from_persist_dir(str(PERSIST_PATH / INDEX_NAME))

if PRINT_QUERY:
    print(query)

nodes_with_score = retriever.retrieve(query)
for index, node_with_score in enumerate(nodes_with_score, start=1):
    source_title = node_with_score.node.metadata.get("title", "Невідомий заголовок")
    print(index, f"{node_with_score.score:.4f}", source_title, "\n", sep="\t")