import logging

from haystack import Document
from haystack.components.writers import DocumentWriter
from haystack.document_stores.errors import DuplicateDocumentError
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.document_stores.types import DuplicatePolicy


logging.getLogger("haystack").setLevel(logging.ERROR)

document_store = InMemoryDocumentStore()
document_id = "release-notes"

seed_writer = DocumentWriter(
    document_store=document_store,
    policy=DuplicatePolicy.FAIL,
)
seed_writer.run(
    documents=[Document(id=document_id, content="Version 1 release notes")]
)

writer = DocumentWriter(
    document_store=document_store,
    policy=DuplicatePolicy.SKIP,
)

skip_result = writer.run(
    documents=[Document(id=document_id, content="Skipped replacement")]
)
stored_after_skip = document_store.filter_documents()[0].content

overwrite_result = writer.run(
    documents=[Document(id=document_id, content="Version 2 release notes")],
    policy=DuplicatePolicy.OVERWRITE,
)
stored_after_overwrite = document_store.filter_documents()[0].content

try:
    writer.run(
        documents=[Document(id=document_id, content="Rejected replacement")],
        policy=DuplicatePolicy.FAIL,
    )
except DuplicateDocumentError:
    fail_error = "DuplicateDocumentError"
else:
    fail_error = "No error"

final_documents = document_store.filter_documents()

print(f"SKIP wrote: {skip_result['documents_written']}")
print(f"After SKIP: {stored_after_skip}")
print(f"OVERWRITE wrote: {overwrite_result['documents_written']}")
print(f"After OVERWRITE: {stored_after_overwrite}")
print(f"FAIL raised: {fail_error}")
print(f"Final count: {len(final_documents)}")

assert skip_result["documents_written"] == 0
assert stored_after_skip == "Version 1 release notes"
assert overwrite_result["documents_written"] == 1
assert stored_after_overwrite == "Version 2 release notes"
assert fail_error == "DuplicateDocumentError"
assert len(final_documents) == 1
assert final_documents[0].content == "Version 2 release notes"