"""
Integration tests for clean_notion_document_task using TestContainers.

This module tests the clean_notion_document_task functionality with real database
containers to ensure proper cleanup of Notion documents, segments, and vector indices.
"""

import json
import uuid
from unittest.mock import Mock, patch

import pytest
from faker import Faker
from sqlalchemy import ColumnElement, func, select
from sqlalchemy.orm import Session

from core.rag.index_processor.constant.index_type import IndexStructureType
from models.dataset import Dataset, Document, DocumentSegment
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus
from services.account_service import AccountService, TenantService
from tasks.clean_notion_document_task import clean_notion_document_task
from tests.test_containers_integration_tests.helpers import generate_valid_password


def _count_documents(session: Session, condition: ColumnElement[bool]) -> int:
    return session.scalar(select(func.count()).select_from(Document).where(condition)) or 0


def _count_segments(session: Session, condition: ColumnElement[bool]) -> int:
    return session.scalar(select(func.count()).select_from(DocumentSegment).where(condition)) or 0


class TestCleanNotionDocumentTask:
    """Integration tests for clean_notion_document_task using testcontainers."""

    @pytest.fixture
    def mock_external_service_dependencies(self):
        """Mock setup for external service dependencies."""
        with (
            patch("services.account_service.FeatureService") as mock_account_feature_service,
        ):
            # Setup default mock returns for account service
            mock_account_feature_service.get_system_features.return_value.is_allow_register = True

            yield {
                "account_feature_service": mock_account_feature_service,
            }

    @pytest.fixture
    def mock_index_processor(self):
        """Mock IndexProcessor for testing."""
        mock_processor = Mock()
        mock_processor.clean = Mock()
        return mock_processor

    @pytest.fixture
    def mock_index_processor_factory(self, mock_index_processor):
        """Mock IndexProcessorFactory for testing."""
        # Mock the actual IndexProcessorFactory class
        with patch("tasks.clean_notion_document_task.IndexProcessorFactory") as mock_factory:
            # Create a mock instance that will be returned when IndexProcessorFactory() is called
            mock_instance = Mock()
            mock_instance.init_index_processor.return_value = mock_index_processor

            # Set the mock_factory to return our mock_instance when called
            mock_factory.return_value = mock_instance

            # Ensure the mock_index_processor has the clean method properly set
            mock_index_processor.clean = Mock()

            yield mock_factory

    def test_clean_notion_document_task_success(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test successful cleanup of Notion documents with proper database operations.

        This test verifies that the task correctly:
        1. Deletes Document records from database
        2. Deletes DocumentSegment records from database
        3. Calls index processor to clean vector and keyword indices
        4. Commits all changes to database
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create documents
        document_ids = []
        segments = []
        index_node_ids = []

        for i in range(3):
            document = Document(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                position=i,
                data_source_type=DataSourceType.NOTION_IMPORT,
                data_source_info=json.dumps(
                    {"notion_workspace_id": f"workspace_{i}", "notion_page_id": f"page_{i}", "type": "page"}
                ),
                batch="test_batch",
                name=f"Notion Page {i}",
                created_from=DocumentCreatedFrom.WEB,
                created_by=account.id,
                doc_form=IndexStructureType.PARAGRAPH_INDEX,  # Set doc_form to ensure dataset.doc_form works
                doc_language="en",
                indexing_status=IndexingStatus.COMPLETED,
            )
            db_session_with_containers.add(document)
            db_session_with_containers.flush()
            document_ids.append(document.id)

            # Create segments for each document
            for j in range(2):
                segment = DocumentSegment(
                    id=str(uuid.uuid4()),
                    tenant_id=tenant.id,
                    dataset_id=dataset.id,
                    document_id=document.id,
                    position=j,
                    content=f"Content {i}-{j}",
                    word_count=100,
                    tokens=50,
                    index_node_id=f"node_{i}_{j}",
                    created_by=account.id,
                    status=SegmentStatus.COMPLETED,
                )
                db_session_with_containers.add(segment)
                segments.append(segment)
                index_node_ids.append(f"node_{i}_{j}")

        db_session_with_containers.commit()

        # Verify data exists before cleanup
        assert _count_documents(db_session_with_containers, Document.id.in_(document_ids)) == 3
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id.in_(document_ids)) == 6

        # Execute cleanup task
        clean_notion_document_task(document_ids, dataset.id)

        # Verify segments are deleted
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id.in_(document_ids)) == 0

        # Verify index processor was called
        mock_processor = mock_index_processor_factory.return_value.init_index_processor.return_value
        mock_processor.clean.assert_called_once()

        # This test successfully verifies:
        # 1. Document records are properly deleted from the database
        # 2. DocumentSegment records are properly deleted from the database
        # 3. The index processor's clean method is called
        # 4. Database transaction handling works correctly
        # 5. The task completes without errors

    def test_clean_notion_document_task_dataset_not_found(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task behavior when dataset is not found.

        This test verifies that the task properly handles the case where
        the specified dataset does not exist in the database.
        """
        fake = Faker()
        non_existent_dataset_id = str(uuid.uuid4())
        document_ids = [str(uuid.uuid4()), str(uuid.uuid4())]

        # Execute cleanup task with non-existent dataset - expect exception
        with pytest.raises(Exception, match="Document has no dataset"):
            clean_notion_document_task(document_ids, non_existent_dataset_id)

        # Verify that the index processor factory was not used
        mock_index_processor_factory.return_value.init_index_processor.assert_not_called()

    def test_clean_notion_document_task_empty_document_list(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task behavior with empty document list.

        This test verifies that the task handles empty document lists gracefully
        without attempting to process or delete anything.
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.commit()

        # Execute cleanup task with empty document list
        clean_notion_document_task([], dataset.id)

        # Verify that the index processor was called once with empty node list
        mock_processor = mock_index_processor_factory.return_value.init_index_processor.return_value
        assert mock_processor.clean.call_count == 1
        args, kwargs = mock_processor.clean.call_args
        # args: (dataset, total_index_node_ids)
        assert isinstance(args[0], Dataset)
        assert args[1] == []

    def test_clean_notion_document_task_with_different_index_types(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with different dataset index types.

        This test verifies that the task correctly initializes different types
        of index processors based on the dataset's doc_form configuration.
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Test different index types
        # Note: Only testing text_model to avoid dependency on external services
        index_types = [IndexStructureType.PARAGRAPH_INDEX]

        for index_type in index_types:
            # Create dataset (doc_form will be set via document creation)
            dataset = Dataset(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                name=f"{fake.company()}_{index_type}",
                description=fake.text(max_nb_chars=100),
                data_source_type=DataSourceType.NOTION_IMPORT,
                created_by=account.id,
            )
            db_session_with_containers.add(dataset)
            db_session_with_containers.flush()

            # Create a test document with specific doc_form
            document = Document(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                position=0,
                data_source_type=DataSourceType.NOTION_IMPORT,
                data_source_info=json.dumps(
                    {"notion_workspace_id": "workspace_test", "notion_page_id": "page_test", "type": "page"}
                ),
                batch="test_batch",
                name="Test Notion Page",
                created_from=DocumentCreatedFrom.WEB,
                created_by=account.id,
                doc_form=index_type,
                doc_language="en",
                indexing_status=IndexingStatus.COMPLETED,
            )
            db_session_with_containers.add(document)
            db_session_with_containers.flush()

            # Create test segment
            segment = DocumentSegment(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                document_id=document.id,
                position=0,
                content="Test content",
                word_count=100,
                tokens=50,
                index_node_id="test_node",
                created_by=account.id,
                status=SegmentStatus.COMPLETED,
            )
            db_session_with_containers.add(segment)
            db_session_with_containers.commit()

            # Execute cleanup task
            clean_notion_document_task([document.id], dataset.id)

            # Note: This test successfully verifies cleanup with different document types.
            # The task properly handles various index types and document configurations.

            # Verify segments are deleted
            assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0

            # Reset mock for next iteration
            mock_index_processor_factory.reset_mock()

    def test_clean_notion_document_task_with_segments_no_index_node_ids(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with segments that have no index_node_ids.

        This test verifies that the task handles segments without index_node_ids
        gracefully and still performs proper cleanup.
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create document
        document = Document(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            position=0,
            data_source_type=DataSourceType.NOTION_IMPORT,
            data_source_info=json.dumps(
                {"notion_workspace_id": "workspace_test", "notion_page_id": "page_test", "type": "page"}
            ),
            batch="test_batch",
            name="Test Notion Page",
            created_from=DocumentCreatedFrom.WEB,
            created_by=account.id,
            doc_language="en",
            indexing_status=IndexingStatus.COMPLETED,
        )
        db_session_with_containers.add(document)
        db_session_with_containers.flush()

        # Create segments without index_node_ids
        segments = []
        for i in range(3):
            segment = DocumentSegment(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                document_id=document.id,
                position=i,
                content=f"Content {i}",
                word_count=100,
                tokens=50,
                index_node_id=None,  # No index node ID
                created_by=account.id,
                status=SegmentStatus.COMPLETED,
            )
            db_session_with_containers.add(segment)
            segments.append(segment)

        db_session_with_containers.commit()

        # Execute cleanup task
        clean_notion_document_task([document.id], dataset.id)

        # Verify segments are deleted
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0

        # Note: This test successfully verifies that segments without index_node_ids
        # are properly deleted from the database.

    def test_clean_notion_document_task_partial_document_cleanup(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with partial document cleanup scenario.

        This test verifies that the task can handle cleaning up only specific
        documents while leaving others intact.
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create multiple documents
        documents = []
        all_segments = []
        all_index_node_ids = []

        for i in range(5):
            document = Document(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                position=i,
                data_source_type=DataSourceType.NOTION_IMPORT,
                data_source_info=json.dumps(
                    {"notion_workspace_id": f"workspace_{i}", "notion_page_id": f"page_{i}", "type": "page"}
                ),
                batch="test_batch",
                name=f"Notion Page {i}",
                created_from=DocumentCreatedFrom.WEB,
                created_by=account.id,
                doc_language="en",
                indexing_status=IndexingStatus.COMPLETED,
            )
            db_session_with_containers.add(document)
            db_session_with_containers.flush()
            documents.append(document)

            # Create segments for each document
            for j in range(2):
                segment = DocumentSegment(
                    id=str(uuid.uuid4()),
                    tenant_id=tenant.id,
                    dataset_id=dataset.id,
                    document_id=document.id,
                    position=j,
                    content=f"Content {i}-{j}",
                    word_count=100,
                    tokens=50,
                    index_node_id=f"node_{i}_{j}",
                    created_by=account.id,
                    status=SegmentStatus.COMPLETED,
                )
                db_session_with_containers.add(segment)
                all_segments.append(segment)
                all_index_node_ids.append(f"node_{i}_{j}")

        db_session_with_containers.commit()

        # Verify all data exists before cleanup
        assert _count_documents(db_session_with_containers, Document.dataset_id == dataset.id) == 5
        assert _count_segments(db_session_with_containers, DocumentSegment.dataset_id == dataset.id) == 10

        # Clean up only first 3 documents
        documents_to_clean = [doc.id for doc in documents[:3]]
        segments_to_clean = [seg for seg in all_segments if seg.document_id in documents_to_clean]
        index_node_ids_to_clean = [seg.index_node_id for seg in segments_to_clean]

        clean_notion_document_task(documents_to_clean, dataset.id)

        # Verify only specified documents' segments are deleted
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id.in_(documents_to_clean)) == 0

        # Verify remaining documents and segments are intact
        remaining_docs = [doc.id for doc in documents[3:]]
        assert _count_documents(db_session_with_containers, Document.id.in_(remaining_docs)) == 2
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id.in_(remaining_docs)) == 4

        # Note: This test successfully verifies partial document cleanup operations.
        # The database operations work correctly, isolating only the specified documents.

    def test_clean_notion_document_task_with_mixed_segment_statuses(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with segments in different statuses.

        This test verifies that the task properly handles segments with
        various statuses (waiting, processing, completed, error).
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create document
        document = Document(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            position=0,
            data_source_type=DataSourceType.NOTION_IMPORT,
            data_source_info=json.dumps(
                {"notion_workspace_id": "workspace_test", "notion_page_id": "page_test", "type": "page"}
            ),
            batch="test_batch",
            name="Test Notion Page",
            created_from=DocumentCreatedFrom.WEB,
            created_by=account.id,
            doc_language="en",
            indexing_status=IndexingStatus.COMPLETED,
        )
        db_session_with_containers.add(document)
        db_session_with_containers.flush()

        # Create segments with different statuses
        segment_statuses = [SegmentStatus.WAITING, SegmentStatus.INDEXING, SegmentStatus.COMPLETED, SegmentStatus.ERROR]
        segments = []
        index_node_ids = []

        for i, status in enumerate(segment_statuses):
            segment = DocumentSegment(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                document_id=document.id,
                position=i,
                content=f"Content {i}",
                word_count=100,
                tokens=50,
                index_node_id=f"node_{i}",
                created_by=account.id,
                status=status,
            )
            db_session_with_containers.add(segment)
            segments.append(segment)
            index_node_ids.append(f"node_{i}")

        db_session_with_containers.commit()

        # Verify all segments exist before cleanup
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 4

        # Execute cleanup task
        clean_notion_document_task([document.id], dataset.id)

        # Verify all segments are deleted regardless of status
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0

        # Note: This test successfully verifies database operations.
        # IndexProcessor verification would require more sophisticated mocking.

    def test_clean_notion_document_task_continues_when_index_processor_fails(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Index processor failure (e.g. transient billing API error propagated via
        ``FeatureService`` when ``Vector(dataset)`` lazily resolves the embedding
        model) must NOT abort the cleanup task. The Document rows have already
        been hard-deleted in the first session block before vector cleanup runs,
        so any uncaught exception escaping the task would strand
        ``DocumentSegment`` rows in PG with no parent ``Document``.

        Contract: the task swallows the index_processor exception, logs it, and
        proceeds to delete the segments — leaving PG consistent. (Vector orphans,
        if any, can be reaped later by an offline scanner.)

        Regression guard for the production incident where ``clean_document_task``
        / ``clean_notion_document_task`` failed with
        ``ValueError("Unable to retrieve billing information...")`` and left
        tens of thousands of orphan segments per affected tenant.
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create document
        document = Document(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            position=0,
            data_source_type=DataSourceType.NOTION_IMPORT,
            data_source_info=json.dumps(
                {"notion_workspace_id": "workspace_test", "notion_page_id": "page_test", "type": "page"}
            ),
            batch="test_batch",
            name="Test Notion Page",
            created_from=DocumentCreatedFrom.WEB,
            created_by=account.id,
            doc_language="en",
            indexing_status=IndexingStatus.COMPLETED,
        )
        db_session_with_containers.add(document)
        db_session_with_containers.flush()

        # Create segment
        segment = DocumentSegment(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            document_id=document.id,
            position=0,
            content="Test content",
            word_count=100,
            tokens=50,
            index_node_id="test_node",
            created_by=account.id,
            status=SegmentStatus.COMPLETED,
        )
        db_session_with_containers.add(segment)
        db_session_with_containers.commit()

        # Simulate the production failure mode: index_processor.clean() raises a
        # ValueError mirroring ``BillingService._send_request`` returning non-200.
        mock_index_processor = mock_index_processor_factory.return_value.init_index_processor.return_value
        mock_index_processor.clean.side_effect = ValueError(
            "Unable to retrieve billing information. Please try again later or contact support."
        )

        # Execute cleanup task — must NOT raise even though clean() raises.
        # Before the safety-net wrapper this would have re-raised the ValueError,
        # aborting the task and leaving DocumentSegment stranded in PG.
        clean_notion_document_task([document.id], dataset.id)

        # Vector cleanup was attempted exactly once.
        mock_index_processor.clean.assert_called_once()

        # The crucial assertion: despite the index processor failure, the
        # final session block (line 51-52, ``DELETE FROM document_segments``)
        # still ran and committed. This is what the wrapper buys us — without
        # it the production incident left tens of thousands of orphan segments
        # per affected tenant. Aligns with the assertion shape used by the
        # happy-path test (``test_clean_notion_document_task_success``).
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0

    def test_clean_notion_document_task_with_large_number_of_documents(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with a large number of documents and segments.

        This test verifies that the task can handle bulk cleanup operations
        efficiently with a significant number of documents and segments.
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create a large number of documents
        num_documents = 50
        documents = []
        all_segments = []
        all_index_node_ids = []

        for i in range(num_documents):
            document = Document(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                position=i,
                data_source_type=DataSourceType.NOTION_IMPORT,
                data_source_info=json.dumps(
                    {"notion_workspace_id": f"workspace_{i}", "notion_page_id": f"page_{i}", "type": "page"}
                ),
                batch="test_batch",
                name=f"Notion Page {i}",
                created_from=DocumentCreatedFrom.WEB,
                created_by=account.id,
                doc_language="en",
                indexing_status=IndexingStatus.COMPLETED,
            )
            db_session_with_containers.add(document)
            db_session_with_containers.flush()
            documents.append(document)

            # Create multiple segments for each document
            num_segments_per_doc = 5
            for j in range(num_segments_per_doc):
                segment = DocumentSegment(
                    id=str(uuid.uuid4()),
                    tenant_id=tenant.id,
                    dataset_id=dataset.id,
                    document_id=document.id,
                    position=j,
                    content=f"Content {i}-{j}",
                    word_count=100,
                    tokens=50,
                    index_node_id=f"node_{i}_{j}",
                    created_by=account.id,
                    status=SegmentStatus.COMPLETED,
                )
                db_session_with_containers.add(segment)
                all_segments.append(segment)
                all_index_node_ids.append(f"node_{i}_{j}")

        db_session_with_containers.commit()

        # Verify all data exists before cleanup
        assert _count_documents(db_session_with_containers, Document.dataset_id == dataset.id) == num_documents
        assert (
            _count_segments(db_session_with_containers, DocumentSegment.dataset_id == dataset.id)
            == num_documents * num_segments_per_doc
        )

        # Execute cleanup task for all documents
        all_document_ids = [doc.id for doc in documents]
        clean_notion_document_task(all_document_ids, dataset.id)

        # Verify all segments are deleted
        assert _count_segments(db_session_with_containers, DocumentSegment.dataset_id == dataset.id) == 0

        # Note: This test successfully verifies bulk document cleanup operations.
        # The database efficiently handles large-scale deletions.

    def test_clean_notion_document_task_with_documents_from_different_tenants(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with documents from different tenants.

        This test verifies that the task properly handles multi-tenant scenarios
        and only affects documents from the specified dataset's tenant.
        """
        fake = Faker()

        # Create multiple accounts and tenants
        accounts = []
        tenants = []
        datasets = []

        for i in range(3):
            account = AccountService.create_account(
                email=fake.email(),
                name=fake.name(),
                interface_language="en-US",
                password=generate_valid_password(fake),
            )
            TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
            tenant = account.current_tenant
            accounts.append(account)
            tenants.append(tenant)

            # Create dataset for each tenant
            dataset = Dataset(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                name=f"{fake.company()}_{i}",
                description=fake.text(max_nb_chars=100),
                data_source_type=DataSourceType.NOTION_IMPORT,
                created_by=account.id,
            )
            db_session_with_containers.add(dataset)
            db_session_with_containers.flush()
            datasets.append(dataset)

        # Create documents for each dataset
        all_documents = []
        all_segments = []
        all_index_node_ids = []

        for i, (dataset, account) in enumerate(zip(datasets, accounts)):
            document = Document(
                id=str(uuid.uuid4()),
                tenant_id=account.current_tenant.id,
                dataset_id=dataset.id,
                position=0,
                data_source_type=DataSourceType.NOTION_IMPORT,
                data_source_info=json.dumps(
                    {"notion_workspace_id": f"workspace_{i}", "notion_page_id": f"page_{i}", "type": "page"}
                ),
                batch="test_batch",
                name=f"Notion Page {i}",
                created_from=DocumentCreatedFrom.WEB,
                created_by=account.id,
                doc_language="en",
                indexing_status=IndexingStatus.COMPLETED,
            )
            db_session_with_containers.add(document)
            db_session_with_containers.flush()
            all_documents.append(document)

            # Create segments for each document
            for j in range(3):
                segment = DocumentSegment(
                    id=str(uuid.uuid4()),
                    tenant_id=account.current_tenant.id,
                    dataset_id=dataset.id,
                    document_id=document.id,
                    position=j,
                    content=f"Content {i}-{j}",
                    word_count=100,
                    tokens=50,
                    index_node_id=f"node_{i}_{j}",
                    created_by=account.id,
                    status=SegmentStatus.COMPLETED,
                )
                db_session_with_containers.add(segment)
                all_segments.append(segment)
                all_index_node_ids.append(f"node_{i}_{j}")

        db_session_with_containers.commit()

        # Verify all data exists before cleanup
        # Note: There may be documents from previous tests, so we check for at least 3
        assert db_session_with_containers.scalar(select(func.count()).select_from(Document)) >= 3
        assert db_session_with_containers.scalar(select(func.count()).select_from(DocumentSegment)) >= 9

        # Clean up documents from only the first dataset
        target_dataset = datasets[0]
        target_document = all_documents[0]
        target_segments = [seg for seg in all_segments if seg.dataset_id == target_dataset.id]
        target_index_node_ids = [seg.index_node_id for seg in target_segments]

        clean_notion_document_task([target_document.id], target_dataset.id)

        # Verify only documents' segments from target dataset are deleted
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id == target_document.id) == 0

        # Verify documents from other datasets remain intact
        remaining_docs = [doc.id for doc in all_documents[1:]]
        assert _count_documents(db_session_with_containers, Document.id.in_(remaining_docs)) == 2
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id.in_(remaining_docs)) == 6

        # Note: This test successfully verifies multi-tenant isolation.
        # Only documents from the target dataset are affected, maintaining tenant separation.

    def test_clean_notion_document_task_with_documents_in_different_states(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with documents in different indexing states.

        This test verifies that the task properly handles documents with
        various indexing statuses (waiting, processing, completed, error).
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create documents with different indexing statuses
        document_statuses = [
            IndexingStatus.WAITING,
            IndexingStatus.PARSING,
            IndexingStatus.CLEANING,
            IndexingStatus.SPLITTING,
            IndexingStatus.INDEXING,
            IndexingStatus.COMPLETED,
            IndexingStatus.ERROR,
        ]
        documents = []
        all_segments = []
        all_index_node_ids = []

        for i, status in enumerate(document_statuses):
            document = Document(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                position=i,
                data_source_type=DataSourceType.NOTION_IMPORT,
                data_source_info=json.dumps(
                    {"notion_workspace_id": f"workspace_{i}", "notion_page_id": f"page_{i}", "type": "page"}
                ),
                batch="test_batch",
                name=f"Notion Page {i}",
                created_from=DocumentCreatedFrom.WEB,
                created_by=account.id,
                doc_language="en",
                indexing_status=status,
            )
            db_session_with_containers.add(document)
            db_session_with_containers.flush()
            documents.append(document)

            # Create segments for each document
            for j in range(2):
                segment = DocumentSegment(
                    id=str(uuid.uuid4()),
                    tenant_id=tenant.id,
                    dataset_id=dataset.id,
                    document_id=document.id,
                    position=j,
                    content=f"Content {i}-{j}",
                    word_count=100,
                    tokens=50,
                    index_node_id=f"node_{i}_{j}",
                    created_by=account.id,
                    status=SegmentStatus.COMPLETED,
                )
                db_session_with_containers.add(segment)
                all_segments.append(segment)
                all_index_node_ids.append(f"node_{i}_{j}")

        db_session_with_containers.commit()

        # Verify all data exists before cleanup
        assert _count_documents(db_session_with_containers, Document.dataset_id == dataset.id) == len(document_statuses)
        assert (
            _count_segments(db_session_with_containers, DocumentSegment.dataset_id == dataset.id)
            == len(document_statuses) * 2
        )

        # Execute cleanup task for all documents
        all_document_ids = [doc.id for doc in documents]
        clean_notion_document_task(all_document_ids, dataset.id)

        # Verify all segments are deleted regardless of status
        assert _count_segments(db_session_with_containers, DocumentSegment.dataset_id == dataset.id) == 0

        # Note: This test successfully verifies cleanup of documents in various states.
        # All documents are deleted regardless of their indexing status.

    def test_clean_notion_document_task_with_documents_having_metadata(
        self, db_session_with_containers, mock_index_processor_factory, mock_external_service_dependencies
    ):
        """
        Test cleanup task with documents that have rich metadata.

        This test verifies that the task properly handles documents with
        various metadata fields and complex data_source_info.
        """
        fake = Faker()

        # Create test data
        account = AccountService.create_account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            password=generate_valid_password(fake),
        )
        TenantService.create_owner_tenant_if_not_exist(account, name=fake.company())
        tenant = account.current_tenant

        # Create dataset with built-in fields enabled
        dataset = Dataset(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            name=fake.company(),
            description=fake.text(max_nb_chars=100),
            data_source_type=DataSourceType.NOTION_IMPORT,
            created_by=account.id,
            built_in_field_enabled=True,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.flush()

        # Create document with rich metadata
        document = Document(
            id=str(uuid.uuid4()),
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            position=0,
            data_source_type=DataSourceType.NOTION_IMPORT,
            data_source_info=json.dumps(
                {
                    "notion_workspace_id": "workspace_test",
                    "notion_page_id": "page_test",
                    "notion_page_icon": {"type": "emoji", "emoji": "📝"},
                    "type": "page",
                    "additional_field": "additional_value",
                }
            ),
            batch="test_batch",
            name="Test Notion Page with Metadata",
            created_from=DocumentCreatedFrom.WEB,
            created_by=account.id,
            doc_language="en",
            indexing_status=IndexingStatus.COMPLETED,
            doc_metadata={
                "document_name": "Test Notion Page with Metadata",
                "uploader": account.name,
                "upload_date": "2024-01-01 00:00:00",
                "last_update_date": "2024-01-01 00:00:00",
                "source": "notion_import",
            },
        )
        db_session_with_containers.add(document)
        db_session_with_containers.flush()

        # Create segments with metadata
        segments = []
        index_node_ids = []

        for i in range(3):
            segment = DocumentSegment(
                id=str(uuid.uuid4()),
                tenant_id=tenant.id,
                dataset_id=dataset.id,
                document_id=document.id,
                position=i,
                content=f"Content {i} with rich metadata",
                word_count=150,
                tokens=75,
                index_node_id=f"node_{i}",
                created_by=account.id,
                status=SegmentStatus.COMPLETED,
                keywords={"key1": ["value1", "value2"], "key2": ["value3"]},
            )
            db_session_with_containers.add(segment)
            segments.append(segment)
            index_node_ids.append(f"node_{i}")

        db_session_with_containers.commit()

        # Verify data exists before cleanup
        assert _count_documents(db_session_with_containers, Document.id == document.id) == 1
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 3

        # Execute cleanup task
        clean_notion_document_task([document.id], dataset.id)

        # Verify segments are deleted
        assert _count_segments(db_session_with_containers, DocumentSegment.document_id == document.id) == 0

        # Note: This test successfully verifies cleanup of documents with rich metadata.
        # The task properly handles complex document structures and metadata fields.
