import os
from dotenv import load_dotenv
from vectorstore import create_vectorstore, load_documents

def main():
    load_dotenv()  # загружаем TAVILY_API_KEY и другие переменные
    persist_dir = os.getenv("CHROMA_PERSIST_DIR", "./chroma_db")
    vectorstore = create_vectorstore(persist_directory=persist_dir)

    # Проверяем, пуста ли база данных
    try:
        if hasattr(vectorstore, "_collection"):
            doc_count = vectorstore._collection.count()
        else:
            # запасной способ: пробуем выполнить пустой поиск
            doc_count = len(vectorstore.similarity_search("test", k=1))
    except Exception:
        doc_count = 0

    if doc_count == 0:
        docs_dir = os.path.join(os.path.dirname(__file__), "documents")
        if not os.path.isdir(docs_dir):
            print(f"Каталог {docs_dir} не найден.")
            return
        print(f"Загрузка документов из {docs_dir} в ChromaDB ({persist_dir})...")
        load_documents(docs_dir, vectorstore)
        print("Загрузка завершена.")
    else:
        print(f"База данных в {persist_dir} уже содержит {doc_count} документов. Загрузка пропущена.")

if __name__ == "__main__":
    main()