from langchain.embeddings import OpenAIEmbeddings
from langchain.schema import Document
from langchain.vectorstores.elasticsearch import ElasticsearchStore
from langchain.indexes import SQLRecordManager, index


collection_name = "test_index"

embedding = OpenAIEmbeddings()

vectorstore = ElasticsearchStore(
    es_url="http://localhost:9200",
    index_name="test_index",
    embedding=embedding)

namespace = f"elasticsearch/{collection_name}"
record_manager = SQLRecordManager(
    namespace, db_url="sqlite:///record_manager_cache.sql"
)
# record_manager.create_schema()

doc1 = Document(page_content="kitty", metadata={"source": "kitty.txt"})
doc2 = Document(page_content="doggy", metadata={"source": "doggy.txt"})
doc3 = Document(page_content="doggy1", metadata={"source": "doggy.txt"})


def _clear():
    """Hacky helper method to clear content. See the `full` mode section to to understand why it works."""
    index(
        [],
        record_manager,
        vectorstore,
        cleanup="full",
        source_id_key="source")

_clear()

res = index(
    [doc1, doc1, doc2],
    record_manager,
    vectorstore,
    cleanup=None,
    source_id_key="source",
)
print(res)

{'num_added': 2, 'num_updated': 0, 'num_skipped': 0, 'num_deleted': 0}

res = index(
    [doc1, doc3],
    record_manager,
    vectorstore,
    cleanup=None,
    source_id_key="source",
)
print(res)

{'num_added': 1, 'num_updated': 0, 'num_skipped': 1, 'num_deleted': 0}

res = index(
    [doc1, doc1, doc2],
    record_manager,
    vectorstore,
    cleanup="full",
    source_id_key="source",
)
print(res)

{'num_added': 2, 'num_updated': 0, 'num_skipped': 0, 'num_deleted': 0}

res = index(
    [doc1, doc3],
    record_manager,
    vectorstore,
    cleanup="full",
    source_id_key="source",
)
print(res)

{'num_added': 1, 'num_updated': 0, 'num_skipped': 1, 'num_deleted': 1}

_clear()

res = index(
    [doc1, doc1, doc2],
    record_manager,
    vectorstore,
    cleanup="incremental",
    source_id_key="source",
)
print(res)

res = index(
    [doc1, doc3],
    record_manager,
    vectorstore,
    cleanup="incremental",
    source_id_key="source",
)
print(res)

{'num_added': 2, 'num_updated': 0, 'num_skipped': 0, 'num_deleted': 0}
{'num_added': 1, 'num_updated': 0, 'num_skipped': 1, 'num_deleted': 1}

res = index(
    [doc1],
    record_manager,
    vectorstore,
    cleanup="incremental",
    source_id_key="source",
)
print(res)

{'num_added': 0, 'num_updated': 0, 'num_skipped': 1, 'num_deleted': 0}

{'num_added': 1, 'num_updated': 0, 'num_skipped':

def index(
    docs_source: Union[BaseLoader, Iterable[Document]],
    record_manager: RecordManager,
    vector_store: VectorStore,
    *,
    batch_size: int = 100,
    cleanup: Literal["incremental", "full", None] = None,
    source_id_key: Union[str, Callable[[Document], str], None] = None,
    cleanup_batch_size: int = 1_000,
) -> IndexingResult:
    ...

    if isinstance(docs_source, BaseLoader):
        try:
            doc_iterator = docs_source.lazy_load()
        except NotImplementedError:
            doc_iterator = iter(docs_source.load())
    else:
        doc_iterator = iter(docs_source)

    source_id_assigner = _get_source_id_assigner(source_id_key)

    # Mark when the update started.
    index_start_dt = record_manager.get_time()
    num_added = 0
    num_skipped = 0
    num_updated = 0
    num_deleted = 0
    
    
    for doc_batch in _batch(batch_size, doc_iterator):
        hashed_docs = list(
            _deduplicate_in_order(
                [_HashedDocument.from_document(doc) for doc in doc_batch]
            )
        )

        source_ids: Sequence[Optional[str]] = [
            source_id_assigner(doc) for doc in hashed_docs
        ]

        ....

        exists_batch = record_manager.exists([doc.uid for doc in hashed_docs])

        # Filter out documents that already exist in the record store.
        uids = []
        docs_to_index = []
        
        # 判斷哪些是要更新，哪些是要添加的
        for hashed_doc, doc_exists in zip(hashed_docs, exists_batch):
            if doc_exists:
                # Must be updated to refresh timestamp.
                record_manager.update([hashed_doc.uid], time_at_least=index_start_dt)
                num_skipped += 1
                continue
            uids.append(hashed_doc.uid)
            docs_to_index.append(hashed_doc.to_document())

      
        # 知識(shí)入向量庫(kù)
        if docs_to_index:
            vector_store.add_documents(docs_to_index, ids=uids)
            num_added += len(docs_to_index)

        # 更新數(shù)據(jù)庫(kù)記錄時(shí)間
        record_manager.update(
            [doc.uid for doc in hashed_docs],
            group_ids=source_ids,
            time_at_least=index_start_dt,
        )

        # 根據(jù)時(shí)間和source_ids 清理舊版本數(shù)據(jù)
        if cleanup == "incremental":
            ...

            uids_to_delete = record_manager.list_keys(
                group_ids=_source_ids, before=index_start_dt
            )
            if uids_to_delete:
                vector_store.delete(uids_to_delete)
                record_manager.delete_keys(uids_to_delete)
                num_deleted += len(uids_to_delete)

    if cleanup == "full":
        while uids_to_delete := record_manager.list_keys(
            before=index_start_dt, limit=cleanup_batch_size
        ):
            # First delete from record store.
            vector_store.delete(uids_to_delete)
            # Then delete from record manager.
            record_manager.delete_keys(uids_to_delete)
            num_deleted += len(uids_to_delete)

    return {
        "num_added": num_added,
        "num_updated": num_updated,
        "num_skipped": num_skipped,
        "num_deleted": num_deleted,
    }