"""
Integration tests for disable_segment_from_index_task using TestContainers.

This module provides comprehensive integration tests for the disable_segment_from_index_task
using real database and Redis containers to ensure the task works correctly with actual
data and external dependencies.
"""

import logging
import time
from datetime import UTC, datetime
from unittest.mock import patch

import pytest
from faker import Faker
from sqlalchemy.orm import Session

from core.rag.index_processor.constant.index_type import IndexStructureType, IndexTechniqueType
from extensions.ext_redis import redis_client
from models import Account, Tenant, TenantAccountJoin, TenantAccountRole
from models.dataset import Dataset, Document, DocumentSegment
from models.enums import DataSourceType, DocumentCreatedFrom, IndexingStatus, SegmentStatus
from tasks.disable_segment_from_index_task import disable_segment_from_index_task

logger = logging.getLogger(__name__)


class TestDisableSegmentFromIndexTask:
    """Integration tests for disable_segment_from_index_task using testcontainers."""

    @pytest.fixture
    def mock_index_processor(self):
        """Mock IndexProcessorFactory and its clean method."""
        with patch("tasks.disable_segment_from_index_task.IndexProcessorFactory") as mock_factory:
            mock_processor = mock_factory.return_value.init_index_processor.return_value
            mock_processor.clean.return_value = None
            yield mock_processor

    def _create_test_account_and_tenant(self, db_session_with_containers: Session) -> tuple[Account, Tenant]:
        """
        Helper method to create a test account and tenant for testing.

        Args:
            db_session_with_containers: Database session from testcontainers infrastructure

        Returns:
            tuple: (account, tenant) - Created account and tenant instances
        """
        fake = Faker()

        # Create account
        account = Account(
            email=fake.email(),
            name=fake.name(),
            interface_language="en-US",
            status="active",
        )
        db_session_with_containers.add(account)
        db_session_with_containers.commit()

        # Create tenant
        tenant = Tenant(
            name=fake.company(),
            status="normal",
            plan="basic",
        )
        db_session_with_containers.add(tenant)
        db_session_with_containers.commit()

        # Create tenant-account join with owner role
        join = TenantAccountJoin(
            tenant_id=tenant.id,
            account_id=account.id,
            role=TenantAccountRole.OWNER,
            current=True,
        )
        db_session_with_containers.add(join)
        db_session_with_containers.commit()

        # Set current tenant for account
        account.current_tenant = tenant

        return account, tenant

    def _create_test_dataset(self, db_session_with_containers: Session, tenant: Tenant, account: Account) -> Dataset:
        """
        Helper method to create a test dataset.

        Args:
            tenant: Tenant instance
            account: Account instance

        Returns:
            Dataset: Created dataset instance
        """
        fake = Faker()

        dataset = Dataset(
            tenant_id=tenant.id,
            name=fake.sentence(nb_words=3),
            description=fake.text(max_nb_chars=200),
            data_source_type=DataSourceType.UPLOAD_FILE,
            indexing_technique=IndexTechniqueType.HIGH_QUALITY,
            created_by=account.id,
        )
        db_session_with_containers.add(dataset)
        db_session_with_containers.commit()

        return dataset

    def _create_test_document(
        self,
        db_session_with_containers: Session,
        dataset: Dataset,
        tenant: Tenant,
        account: Account,
        doc_form: str = IndexStructureType.PARAGRAPH_INDEX,
    ) -> Document:
        """
        Helper method to create a test document.

        Args:
            dataset: Dataset instance
            tenant: Tenant instance
            account: Account instance
            doc_form: Document form type

        Returns:
            Document: Created document instance
        """
        fake = Faker()

        document = Document(
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            position=1,
            data_source_type=DataSourceType.UPLOAD_FILE,
            batch=fake.uuid4(),
            name=fake.file_name(),
            created_from=DocumentCreatedFrom.API,
            created_by=account.id,
            indexing_status=IndexingStatus.COMPLETED,
            enabled=True,
            archived=False,
            doc_form=doc_form,
            word_count=1000,
            tokens=500,
            completed_at=datetime.now(UTC),
        )
        db_session_with_containers.add(document)
        db_session_with_containers.commit()

        return document

    def _create_test_segment(
        self,
        db_session_with_containers: Session,
        document: Document,
        dataset: Dataset,
        tenant: Tenant,
        account: Account,
        status: SegmentStatus = SegmentStatus.COMPLETED,
        enabled: bool = True,
    ) -> DocumentSegment:
        """
        Helper method to create a test document segment.

        Args:
            document: Document instance
            dataset: Dataset instance
            tenant: Tenant instance
            account: Account instance
            status: Segment status
            enabled: Whether segment is enabled

        Returns:
            DocumentSegment: Created segment instance
        """
        fake = Faker()

        segment = DocumentSegment(
            tenant_id=tenant.id,
            dataset_id=dataset.id,
            document_id=document.id,
            position=1,
            content=fake.text(max_nb_chars=500),
            word_count=100,
            tokens=50,
            index_node_id=fake.uuid4(),
            index_node_hash=fake.sha256(),
            status=status,
            enabled=enabled,
            created_by=account.id,
            completed_at=datetime.now(UTC) if status == SegmentStatus.COMPLETED else None,
        )
        db_session_with_containers.add(segment)
        db_session_with_containers.commit()

        return segment

    def test_disable_segment_success(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test successful segment disabling from index.

        This test verifies:
        - Segment is found and validated
        - Index processor clean method is called with correct parameters
        - Redis cache is cleared
        - Task completes successfully
        """
        # Arrange: Create test data
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Set up Redis cache
        indexing_cache_key = f"segment_{segment.id}_indexing"
        redis_client.setex(indexing_cache_key, 600, 1)

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task completed successfully
        assert result is None  # Task returns None on success

        # Verify index processor was called correctly
        mock_index_processor.clean.assert_called_once()
        call_args = mock_index_processor.clean.call_args
        assert call_args[0][0].id == dataset.id  # Check dataset ID
        assert call_args[0][1] == [segment.index_node_id]  # Check index node IDs

        # Verify Redis cache was cleared
        assert redis_client.get(indexing_cache_key) is None

        # Verify segment is still in database
        db_session_with_containers.refresh(segment)
        assert segment.id is not None

    def test_disable_segment_not_found(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test handling when segment is not found.

        This test verifies:
        - Task handles non-existent segment gracefully
        - No index processor operations are performed
        - Task returns early without errors
        """
        # Arrange: Use a non-existent segment ID
        fake = Faker()
        non_existent_segment_id = fake.uuid4()

        # Act: Execute the task with non-existent segment
        result = disable_segment_from_index_task(non_existent_segment_id)

        # Assert: Verify the task handled the error gracefully
        assert result is None

        # Verify index processor was not called
        mock_index_processor.clean.assert_not_called()

    def test_disable_segment_not_completed(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test handling when segment is not in completed status.

        This test verifies:
        - Task rejects segments that are not completed
        - No index processor operations are performed
        - Task returns early without errors
        """
        # Arrange: Create test data with non-completed segment
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(
            db_session_with_containers, document, dataset, tenant, account, status=SegmentStatus.INDEXING, enabled=True
        )

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task handled the invalid status gracefully
        assert result is None

        # Verify index processor was not called
        mock_index_processor.clean.assert_not_called()

    def test_disable_segment_no_dataset(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test handling when segment has no associated dataset.

        This test verifies:
        - Task handles segments without dataset gracefully
        - No index processor operations are performed
        - Task returns early without errors
        """
        # Arrange: Create test data
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Manually remove dataset association
        segment.dataset_id = "00000000-0000-0000-0000-000000000000"
        db_session_with_containers.commit()

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task handled the missing dataset gracefully
        assert result is None

        # Verify index processor was not called
        mock_index_processor.clean.assert_not_called()

    def test_disable_segment_no_document(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test handling when segment has no associated document.

        This test verifies:
        - Task handles segments without document gracefully
        - No index processor operations are performed
        - Task returns early without errors
        """
        # Arrange: Create test data
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Manually remove document association
        segment.document_id = "00000000-0000-0000-0000-000000000000"
        db_session_with_containers.commit()

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task handled the missing document gracefully
        assert result is None

        # Verify index processor was not called
        mock_index_processor.clean.assert_not_called()

    def test_disable_segment_document_disabled(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test handling when document is disabled.

        This test verifies:
        - Task handles disabled documents gracefully
        - No index processor operations are performed
        - Task returns early without errors
        """
        # Arrange: Create test data with disabled document
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        document.enabled = False
        db_session_with_containers.commit()

        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task handled the disabled document gracefully
        assert result is None

        # Verify index processor was not called
        mock_index_processor.clean.assert_not_called()

    def test_disable_segment_document_archived(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test handling when document is archived.

        This test verifies:
        - Task handles archived documents gracefully
        - No index processor operations are performed
        - Task returns early without errors
        """
        # Arrange: Create test data with archived document
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        document.archived = True
        db_session_with_containers.commit()

        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task handled the archived document gracefully
        assert result is None

        # Verify index processor was not called
        mock_index_processor.clean.assert_not_called()

    def test_disable_segment_document_indexing_not_completed(
        self, db_session_with_containers: Session, mock_index_processor
    ):
        """
        Test handling when document indexing is not completed.

        This test verifies:
        - Task handles documents with incomplete indexing gracefully
        - No index processor operations are performed
        - Task returns early without errors
        """
        # Arrange: Create test data with incomplete indexing
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        document.indexing_status = "indexing"
        db_session_with_containers.commit()

        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task handled the incomplete indexing gracefully
        assert result is None

        # Verify index processor was not called
        mock_index_processor.clean.assert_not_called()

    def test_disable_segment_index_processor_exception(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test handling when index processor raises an exception.

        This test verifies:
        - Task handles index processor exceptions gracefully
        - Segment is re-enabled on failure
        - Redis cache is still cleared
        - Database changes are committed
        """
        # Arrange: Create test data
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Set up Redis cache
        indexing_cache_key = f"segment_{segment.id}_indexing"
        redis_client.setex(indexing_cache_key, 600, 1)

        # Configure mock to raise exception
        mock_index_processor.clean.side_effect = Exception("Index processor error")

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify the task handled the exception gracefully
        assert result is None

        # Verify index processor was called
        mock_index_processor.clean.assert_called_once()
        call_args = mock_index_processor.clean.call_args
        # Check that the call was made with the correct parameters
        assert len(call_args[0]) == 2  # Check two arguments were passed
        assert call_args[0][1] == [segment.index_node_id]  # Check index node IDs

        # Verify segment was re-enabled
        db_session_with_containers.refresh(segment)
        assert segment.enabled is True

        # Verify Redis cache was still cleared
        assert redis_client.get(indexing_cache_key) is None

    def test_disable_segment_different_doc_forms(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test disabling segments with different document forms.

        This test verifies:
        - Task works with different document form types
        - Correct index processor is initialized for each form
        - Index processor clean method is called correctly
        """
        # Test different document forms
        doc_forms = [
            IndexStructureType.PARAGRAPH_INDEX,
            IndexStructureType.QA_INDEX,
            IndexStructureType.PARENT_CHILD_INDEX,
        ]

        for doc_form in doc_forms:
            # Arrange: Create test data for each form
            account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
            dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
            document = self._create_test_document(
                db_session_with_containers, dataset, tenant, account, doc_form=doc_form
            )
            segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

            # Reset mock for each iteration
            mock_index_processor.reset_mock()

            # Act: Execute the task
            result = disable_segment_from_index_task(segment.id)

            # Assert: Verify the task completed successfully
            assert result is None

            # Verify correct index processor was initialized
            mock_index_processor.clean.assert_called_once()
            call_args = mock_index_processor.clean.call_args
            assert call_args[0][0].id == dataset.id  # Check dataset ID
            assert call_args[0][1] == [segment.index_node_id]  # Check index node IDs

    def test_disable_segment_redis_cache_handling(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test Redis cache handling during segment disabling.

        This test verifies:
        - Redis cache is properly set before task execution
        - Cache is cleared after task completion
        - Cache handling works with different scenarios
        """
        # Arrange: Create test data
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Test with cache present
        indexing_cache_key = f"segment_{segment.id}_indexing"
        redis_client.setex(indexing_cache_key, 600, 1)
        assert redis_client.get(indexing_cache_key) is not None

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify cache was cleared
        assert result is None
        assert redis_client.get(indexing_cache_key) is None

        # Test with no cache present
        segment2 = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)
        result2 = disable_segment_from_index_task(segment2.id)

        # Assert: Verify task still works without cache
        assert result2 is None

    def test_disable_segment_performance_timing(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test performance timing of segment disabling task.

        This test verifies:
        - Task execution time is reasonable
        - Performance logging works correctly
        - Task completes within expected time bounds
        """
        # Arrange: Create test data
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Act: Execute the task and measure time
        start_time = time.perf_counter()
        result = disable_segment_from_index_task(segment.id)
        end_time = time.perf_counter()

        # Assert: Verify task completed successfully and timing is reasonable
        assert result is None
        execution_time = end_time - start_time
        assert execution_time < 5.0  # Should complete within 5 seconds

    def test_disable_segment_database_session_management(
        self, db_session_with_containers: Session, mock_index_processor
    ):
        """
        Test database session management during task execution.

        This test verifies:
        - Database sessions are properly managed
        - Sessions are closed after task completion
        - No session leaks occur
        """
        # Arrange: Create test data
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)
        segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)

        # Act: Execute the task
        result = disable_segment_from_index_task(segment.id)

        # Assert: Verify task completed and session management worked
        assert result is None

        # Verify segment is still accessible (session was properly managed)
        db_session_with_containers.refresh(segment)
        assert segment.id is not None

    def test_disable_segment_concurrent_execution(self, db_session_with_containers: Session, mock_index_processor):
        """
        Test concurrent execution of segment disabling tasks.

        This test verifies:
        - Multiple tasks can run concurrently
        - Each task processes its own segment correctly
        - No interference between concurrent tasks
        """
        # Arrange: Create multiple test segments
        account, tenant = self._create_test_account_and_tenant(db_session_with_containers)
        dataset = self._create_test_dataset(db_session_with_containers, tenant, account)
        document = self._create_test_document(db_session_with_containers, dataset, tenant, account)

        segments = []
        for i in range(3):
            segment = self._create_test_segment(db_session_with_containers, document, dataset, tenant, account)
            segments.append(segment)

        # Act: Execute tasks concurrently (simulated)
        results = []
        for segment in segments:
            result = disable_segment_from_index_task(segment.id)
            results.append(result)

        # Assert: Verify all tasks completed successfully
        assert all(result is None for result in results)

        # Verify all segments were processed
        assert mock_index_processor.clean.call_count == len(segments)

        # Verify each segment was processed with correct parameters
        for segment in segments:
            # Check that clean was called with this segment's dataset and index_node_id
            found = False
            for call in mock_index_processor.clean.call_args_list:
                if call[0][0].id == dataset.id and call[0][1] == [segment.index_node_id]:
                    found = True
                    break
            assert found, f"Segment {segment.id} was not processed correctly"
