"""Unit tests for SegmentService behaviors in dataset_service."""

from .dataset_service_test_helpers import (
    Account,
    ChildChunk,
    ChildChunkDeleteIndexError,
    ChildChunkIndexingError,
    ChildChunkUpdateArgs,
    DocumentSegment,
    IndexStructureType,
    MagicMock,
    ModelType,
    SegmentService,
    SegmentUpdateArgs,
    SimpleNamespace,
    _make_child_chunk,
    _make_dataset,
    _make_document,
    _make_lock_context,
    _make_segment,
    create_autospec,
    patch,
    pytest,
)


class TestSegmentServiceChildChunks:
    """Unit tests for child-chunk CRUD helpers."""

    @pytest.fixture
    def account_context(self):
        account = create_autospec(Account, instance=True)
        account.id = "user-1"
        account.current_tenant_id = "tenant-1"

        with patch("services.dataset_service.current_user", account):
            yield account

    def test_create_child_chunk_assigns_next_position_and_commits(self, account_context):
        dataset = SimpleNamespace(id="dataset-1")
        document = _make_document()
        segment = _make_segment()

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.uuid.uuid4", return_value="node-1"),
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-1"),
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            mock_redis.lock.return_value = _make_lock_context()
            mock_db.session.scalar.return_value = 2

            child_chunk = SegmentService.create_child_chunk("child content", segment, document, dataset)

        assert isinstance(child_chunk, ChildChunk)
        assert child_chunk.position == 3
        assert child_chunk.index_node_id == "node-1"
        assert child_chunk.index_node_hash == "hash-1"
        assert child_chunk.word_count == len("child content")
        mock_db.session.add.assert_called_once_with(child_chunk)
        vector_service.create_child_chunk_vector.assert_called_once_with(child_chunk, dataset)
        mock_db.session.commit.assert_called_once()

    def test_create_child_chunk_rolls_back_and_raises_on_vector_failure(self, account_context):
        dataset = SimpleNamespace(id="dataset-1")
        document = _make_document()
        segment = _make_segment()

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.uuid.uuid4", return_value="node-1"),
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-1"),
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            mock_redis.lock.return_value = _make_lock_context()
            mock_db.session.scalar.return_value = None
            vector_service.create_child_chunk_vector.side_effect = RuntimeError("vector failed")

            with pytest.raises(ChildChunkIndexingError, match="vector failed"):
                SegmentService.create_child_chunk("child content", segment, document, dataset)

        mock_db.session.rollback.assert_called_once()
        mock_db.session.commit.assert_not_called()

    def test_update_child_chunks_updates_deletes_and_creates_records(self, account_context):
        dataset = SimpleNamespace(id="dataset-1")
        document = _make_document()
        segment = _make_segment()
        existing_a = ChildChunk(
            tenant_id="tenant-1",
            dataset_id="dataset-1",
            document_id="doc-1",
            segment_id="segment-1",
            position=1,
            content="old content",
            word_count=11,
            created_by="user-1",
        )
        existing_b = ChildChunk(
            tenant_id="tenant-1",
            dataset_id="dataset-1",
            document_id="doc-1",
            segment_id="segment-1",
            position=2,
            content="remove me",
            word_count=9,
            created_by="user-1",
        )
        existing_a.id = "child-a"
        existing_b.id = "child-b"
        with (
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.uuid.uuid4", return_value="node-new"),
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-new"),
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            mock_db.session.scalars.return_value.all.return_value = [existing_a, existing_b]

            result = SegmentService.update_child_chunks(
                [
                    ChildChunkUpdateArgs(id="child-a", content="updated content"),
                    ChildChunkUpdateArgs(content="brand new"),
                ],
                segment,
                document,
                dataset,
            )

        assert [chunk.position for chunk in result] == [1, 3]
        assert existing_a.content == "updated content"
        assert existing_a.updated_by == account_context.id
        assert existing_a.updated_at == "now"
        mock_db.session.bulk_save_objects.assert_called_once_with([existing_a])
        mock_db.session.delete.assert_called_once_with(existing_b)
        new_chunk = result[1]
        assert isinstance(new_chunk, ChildChunk)
        assert new_chunk.position == 3
        assert new_chunk.index_node_id == "node-new"
        vector_service.update_child_chunk_vector.assert_called_once_with(
            [new_chunk], [existing_a], [existing_b], dataset
        )
        mock_db.session.commit.assert_called_once()

    def test_update_child_chunks_rolls_back_on_vector_failure(self, account_context):
        dataset = SimpleNamespace(id="dataset-1")
        document = _make_document()
        segment = _make_segment()
        existing_chunk = _make_child_chunk()

        with (
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            mock_db.session.scalars.return_value.all.return_value = [existing_chunk]
            vector_service.update_child_chunk_vector.side_effect = RuntimeError("vector failed")

            with pytest.raises(ChildChunkIndexingError, match="vector failed"):
                SegmentService.update_child_chunks(
                    [ChildChunkUpdateArgs(id="child-a", content="updated content")],
                    segment,
                    document,
                    dataset,
                )

        mock_db.session.rollback.assert_called_once()

    def test_update_child_chunk_updates_vector_and_commits(self, account_context):
        dataset = SimpleNamespace(id="dataset-1")
        child_chunk = _make_child_chunk()

        with (
            patch("services.dataset_service.current_user", SimpleNamespace(id="user-1")),
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            result = SegmentService.update_child_chunk(
                "new content", child_chunk, _make_segment(), _make_document(), dataset
            )

        assert result is child_chunk
        assert child_chunk.content == "new content"
        assert child_chunk.word_count == len("new content")
        assert child_chunk.updated_by == "user-1"
        assert child_chunk.updated_at == "now"
        mock_db.session.add.assert_called_once_with(child_chunk)
        vector_service.update_child_chunk_vector.assert_called_once_with([], [child_chunk], [], dataset)
        mock_db.session.commit.assert_called_once()

    def test_delete_child_chunk_raises_delete_index_error_on_vector_failure(self):
        dataset = SimpleNamespace(id="dataset-1")
        child_chunk = _make_child_chunk()

        with (
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            vector_service.delete_child_chunk_vector.side_effect = RuntimeError("delete failed")

            with pytest.raises(ChildChunkDeleteIndexError, match="delete failed"):
                SegmentService.delete_child_chunk(child_chunk, dataset)

        mock_db.session.delete.assert_called_once_with(child_chunk)
        mock_db.session.rollback.assert_called_once()


class TestSegmentServiceQueries:
    """Unit tests for child-chunk and segment query helpers."""

    @pytest.fixture
    def account_context(self):
        account = create_autospec(Account, instance=True)
        account.id = "user-1"
        account.current_tenant_id = "tenant-1"

        with patch("services.dataset_service.current_user", account):
            yield account

    def test_get_child_chunks_applies_keyword_filter_and_paginate(self, account_context):
        paginated = SimpleNamespace(items=["chunk"], total=1)

        with (
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.helper.escape_like_pattern", return_value="escaped") as escape_like,
        ):
            mock_db.paginate.return_value = paginated

            result = SegmentService.get_child_chunks(
                segment_id="segment-1",
                document_id="doc-1",
                dataset_id="dataset-1",
                page=2,
                limit=10,
                keyword="needle",
            )

        assert result is paginated
        escape_like.assert_called_once_with("needle")
        mock_db.paginate.assert_called_once()

    def test_get_child_chunk_by_id_returns_only_child_chunk_instances(self):
        child_chunk = _make_child_chunk()

        with patch("services.dataset_service.db") as mock_db:
            mock_db.session.scalar.return_value = child_chunk
            result = SegmentService.get_child_chunk_by_id("child-a", "tenant-1")

        assert result is child_chunk

        with patch("services.dataset_service.db") as mock_db:
            mock_db.session.scalar.return_value = SimpleNamespace()
            result = SegmentService.get_child_chunk_by_id("child-a", "tenant-1")

        assert result is None

    def test_get_segments_uses_status_and_keyword_filters(self):
        paginated = SimpleNamespace(items=["segment"], total=1)

        with (
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.helper.escape_like_pattern", return_value="escaped") as escape_like,
        ):
            mock_db.paginate.return_value = paginated

            items, total = SegmentService.get_segments(
                document_id="doc-1",
                tenant_id="tenant-1",
                status_list=["completed"],
                keyword="needle",
                page=1,
                limit=20,
            )

        assert items == ["segment"]
        assert total == 1
        escape_like.assert_called_once_with("needle")
        mock_db.paginate.assert_called_once()

    def test_get_segment_by_id_returns_only_document_segment_instances(self):
        segment = DocumentSegment(
            id="segment-1",
            tenant_id="tenant-1",
            dataset_id="dataset-1",
            document_id="doc-1",
            position=1,
            content="segment",
            word_count=7,
            tokens=2,
            created_by="user-1",
        )

        with patch("services.dataset_service.db") as mock_db:
            mock_db.session.scalar.return_value = segment
            result = SegmentService.get_segment_by_id("segment-1", "tenant-1")

        assert result is segment

        with patch("services.dataset_service.db") as mock_db:
            mock_db.session.scalar.return_value = SimpleNamespace()
            result = SegmentService.get_segment_by_id("segment-1", "tenant-1")

        assert result is None

    def test_get_segments_by_document_and_dataset_returns_scalars_result(self):
        segment = DocumentSegment(
            id="segment-1",
            tenant_id="tenant-1",
            dataset_id="dataset-1",
            document_id="doc-1",
            position=1,
            content="segment",
            word_count=7,
            tokens=2,
            created_by="user-1",
        )

        with patch("services.dataset_service.db") as mock_db:
            mock_db.session.scalars.return_value.all.return_value = [segment]

            result = SegmentService.get_segments_by_document_and_dataset(
                document_id="doc-1",
                dataset_id="dataset-1",
                status="completed",
                enabled=True,
            )

        assert result == [segment]
        mock_db.session.scalars.assert_called_once()


class TestSegmentServiceValidation:
    """Unit tests for segment-create argument validation."""

    def test_segment_create_args_validate_requires_answer_for_qa_model(self):
        document = _make_document(doc_form=IndexStructureType.QA_INDEX)

        with pytest.raises(ValueError, match="Answer is required"):
            SegmentService.segment_create_args_validate({"content": "question"}, document)

    def test_segment_create_args_validate_requires_non_empty_content(self):
        document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX)

        with pytest.raises(ValueError, match="Content is empty"):
            SegmentService.segment_create_args_validate({"content": "   "}, document)

    def test_segment_create_args_validate_enforces_attachment_limit(self):
        document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX)
        args = {"content": "hello", "attachment_ids": ["a-1", "a-2"]}

        with patch("services.dataset_service.dify_config.SINGLE_CHUNK_ATTACHMENT_LIMIT", 1):
            with pytest.raises(ValueError, match="Exceeded maximum attachment limit of 1"):
                SegmentService.segment_create_args_validate(args, document)

    def test_segment_create_args_validate_requires_attachment_ids_list(self):
        document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX)

        with pytest.raises(ValueError, match="Attachment IDs is invalid"):
            SegmentService.segment_create_args_validate({"content": "hello", "attachment_ids": "bad-type"}, document)


class TestSegmentServiceMutations:
    """Unit tests for segment create, update, delete, and bulk status flows."""

    @pytest.fixture
    def account_context(self):
        account = create_autospec(Account, instance=True)
        account.id = "user-1"
        account.current_tenant_id = "tenant-1"

        with patch("services.dataset_service.current_user", account):
            yield account

    def test_create_segment_creates_bindings_and_marks_segment_error_on_vector_failure(self, account_context):
        dataset = _make_dataset(indexing_technique="economy")
        document = _make_document(
            dataset_id=dataset.id,
            tenant_id=dataset.tenant_id,
            doc_form=IndexStructureType.QA_INDEX,
            word_count=0,
        )
        refreshed_segment = SimpleNamespace(id="segment-1")
        args = {
            "content": "question",
            "answer": "answer",
            "keywords": ["kw-1"],
            "attachment_ids": ["att-1", "att-2"],
        }

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.VectorService") as vector_service,
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-1"),
            patch("services.dataset_service.uuid.uuid4", return_value="node-1"),
            patch("services.dataset_service.naive_utc_now", return_value="now"),
        ):
            mock_redis.lock.return_value = _make_lock_context()

            mock_db.session.scalar.return_value = 2
            mock_db.session.get.return_value = refreshed_segment

            def add_side_effect(obj):
                if obj.__class__.__name__ == "DocumentSegment" and getattr(obj, "id", None) is None:
                    obj.id = "segment-1"

            mock_db.session.add.side_effect = add_side_effect
            vector_service.create_segments_vector.side_effect = RuntimeError("vector failed")

            result = SegmentService.create_segment(args=args, document=document, dataset=dataset)

        created_segment = vector_service.create_segments_vector.call_args.args[1][0]
        attachment_bindings = [
            call.args[0]
            for call in mock_db.session.add.call_args_list
            if call.args and call.args[0].__class__.__name__ == "SegmentAttachmentBinding"
        ]

        assert result is refreshed_segment
        assert created_segment.position == 3
        assert created_segment.answer == "answer"
        assert created_segment.word_count == len("question") + len("answer")
        assert created_segment.status == "error"
        assert created_segment.enabled is False
        assert created_segment.error == "vector failed"
        assert document.word_count == len("question") + len("answer")
        assert len(attachment_bindings) == 2
        assert {binding.attachment_id for binding in attachment_bindings} == {"att-1", "att-2"}
        assert mock_db.session.commit.call_count == 3

    def test_multi_create_segment_high_quality_marks_segments_error_when_vector_creation_fails(self, account_context):
        dataset = _make_dataset(indexing_technique="high_quality")
        document = _make_document(
            dataset_id=dataset.id,
            tenant_id=dataset.tenant_id,
            doc_form=IndexStructureType.QA_INDEX,
            word_count=5,
        )
        segments = [
            {"content": "question-1", "answer": "answer-1", "keywords": ["k1"]},
            {"content": "question-2", "answer": "answer-2"},
        ]
        embedding_model = MagicMock()
        embedding_model.get_text_embedding_num_tokens.side_effect = [[11], [13]]

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.ModelManager") as model_manager_cls,
            patch("services.dataset_service.VectorService") as vector_service,
            patch("services.dataset_service.helper.generate_text_hash", side_effect=["hash-1", "hash-2"]),
            patch("services.dataset_service.uuid.uuid4", side_effect=["node-1", "node-2"]),
            patch("services.dataset_service.naive_utc_now", return_value="now"),
        ):
            mock_redis.lock.return_value = _make_lock_context()
            model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model
            mock_db.session.scalar.return_value = 1
            vector_service.create_segments_vector.side_effect = RuntimeError("vector failed")

            result = SegmentService.multi_create_segment(segments, document, dataset)

        assert len(result) == 2
        assert [segment.position for segment in result] == [2, 3]
        assert [segment.tokens for segment in result] == [11, 13]
        assert all(segment.status == "error" for segment in result)
        assert all(segment.enabled is False for segment in result)
        assert all(segment.error == "vector failed" for segment in result)
        assert document.word_count == 5 + sum(len(item["content"]) + len(item["answer"]) for item in segments)
        vector_service.create_segments_vector.assert_called_once_with(
            [["k1"], None], result, dataset, document.doc_form
        )
        mock_db.session.commit.assert_called_once()

    def test_update_segment_disables_enabled_segment_and_dispatches_index_cleanup(self, account_context):
        segment = _make_segment(enabled=True)
        document = _make_document()
        dataset = _make_dataset()
        args = SegmentUpdateArgs(enabled=False)

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.disable_segment_from_index_task") as disable_task,
        ):
            mock_redis.get.return_value = None

            result = SegmentService.update_segment(args, segment, document, dataset)

        assert result is segment
        assert segment.enabled is False
        assert segment.disabled_at == "now"
        assert segment.disabled_by == account_context.id
        mock_db.session.add.assert_called_once_with(segment)
        mock_db.session.commit.assert_called_once()
        mock_redis.setex.assert_called_once_with(f"segment_{segment.id}_indexing", 600, 1)
        disable_task.delay.assert_called_once_with(segment.id)

    def test_update_segment_rejects_updates_for_disabled_segment(self, account_context):
        segment = _make_segment(enabled=False)
        document = _make_document()
        dataset = _make_dataset()

        with patch("services.dataset_service.redis_client") as mock_redis:
            mock_redis.get.return_value = None

            with pytest.raises(ValueError, match="Can't update disabled segment"):
                SegmentService.update_segment(SegmentUpdateArgs(content="new content"), segment, document, dataset)

    def test_update_segment_rejects_when_indexing_cache_exists(self, account_context):
        segment = _make_segment(enabled=True)
        document = _make_document()
        dataset = _make_dataset()

        with patch("services.dataset_service.redis_client") as mock_redis:
            mock_redis.get.return_value = "1"

            with pytest.raises(ValueError, match="Segment is indexing"):
                SegmentService.update_segment(SegmentUpdateArgs(content="new content"), segment, document, dataset)

    def test_update_segment_updates_keywords_for_same_content_segment(self, account_context):
        segment = _make_segment(content="same content", keywords=["old"])
        document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX, word_count=20)
        dataset = _make_dataset()
        refreshed_segment = SimpleNamespace(id=segment.id)
        args = SegmentUpdateArgs(content="same content", keywords=["new"])

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            mock_redis.get.return_value = None
            mock_db.session.get.return_value = refreshed_segment

            result = SegmentService.update_segment(args, segment, document, dataset)

        assert result is refreshed_segment
        assert segment.keywords == ["new"]
        vector_service.update_segment_vector.assert_called_once_with(["new"], segment, dataset)
        vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset)

    def test_update_segment_regenerates_child_chunks_and_updates_manual_summary(self, account_context):
        segment = _make_segment(content="same content", word_count=len("same content"))
        document = _make_document(
            doc_form=IndexStructureType.PARENT_CHILD_INDEX,
            word_count=20,
        )
        dataset = _make_dataset(indexing_technique="high_quality")
        refreshed_segment = SimpleNamespace(id=segment.id)
        processing_rule = SimpleNamespace(id=document.dataset_process_rule_id)
        existing_summary = SimpleNamespace(summary_content="old summary")
        embedding_model_instance = object()
        args = SegmentUpdateArgs(
            content="same content",
            regenerate_child_chunks=True,
            summary="new summary",
        )

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.ModelManager") as model_manager_cls,
            patch("services.dataset_service.VectorService") as vector_service,
            patch("services.summary_index_service.SummaryIndexService.update_summary_for_segment") as update_summary,
        ):
            mock_redis.get.return_value = None
            model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model_instance

            # get calls: processing_rule, then refreshed_segment
            mock_db.session.get.side_effect = [processing_rule, refreshed_segment]
            # scalar call: existing_summary
            mock_db.session.scalar.return_value = existing_summary

            result = SegmentService.update_segment(args, segment, document, dataset)

        assert result is refreshed_segment
        vector_service.generate_child_chunks.assert_called_once_with(
            segment,
            document,
            dataset,
            embedding_model_instance,
            processing_rule,
            True,
        )
        update_summary.assert_called_once_with(segment, dataset, "new summary")
        vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset)

    def test_update_segment_auto_regenerates_summary_after_content_change(self, account_context):
        segment = _make_segment(content="old", word_count=3)
        document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX, word_count=10)
        dataset = _make_dataset(indexing_technique="high_quality")
        dataset.summary_index_setting = {"enable": True}
        refreshed_segment = SimpleNamespace(id=segment.id)
        existing_summary = SimpleNamespace(summary_content="old summary")
        embedding_model = MagicMock()
        embedding_model.get_text_embedding_num_tokens.return_value = [9]
        args = SegmentUpdateArgs(content="new content", keywords=["kw-1"])

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.ModelManager") as model_manager_cls,
            patch("services.dataset_service.VectorService") as vector_service,
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-1"),
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch(
                "services.summary_index_service.SummaryIndexService.generate_and_vectorize_summary"
            ) as generate_summary,
        ):
            mock_redis.get.return_value = None
            model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model

            mock_db.session.scalar.return_value = existing_summary
            mock_db.session.get.return_value = refreshed_segment

            result = SegmentService.update_segment(args, segment, document, dataset)

        assert result is refreshed_segment
        assert segment.content == "new content"
        assert segment.index_node_hash == "hash-1"
        assert segment.tokens == 9
        assert document.word_count == 18
        vector_service.update_segment_vector.assert_called_once_with(["kw-1"], segment, dataset)
        generate_summary.assert_called_once_with(segment, dataset, {"enable": True})
        vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset)

    def test_update_segment_regenerates_summary_when_manual_summary_is_unchanged(self, account_context):
        segment = _make_segment(content="old", word_count=3)
        document = _make_document(doc_form=IndexStructureType.PARAGRAPH_INDEX, word_count=10)
        dataset = _make_dataset(indexing_technique="high_quality")
        dataset.summary_index_setting = {"enable": True}
        refreshed_segment = SimpleNamespace(id=segment.id)
        existing_summary = SimpleNamespace(summary_content="same summary")
        embedding_model = MagicMock()
        embedding_model.get_text_embedding_num_tokens.return_value = [7]
        args = SegmentUpdateArgs(content="new text", summary="same summary")

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.ModelManager") as model_manager_cls,
            patch("services.dataset_service.VectorService") as vector_service,
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-2"),
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch(
                "services.summary_index_service.SummaryIndexService.generate_and_vectorize_summary"
            ) as generate_summary,
            patch("services.summary_index_service.SummaryIndexService.update_summary_for_segment") as update_summary,
        ):
            mock_redis.get.return_value = None
            model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model

            mock_db.session.scalar.return_value = existing_summary
            mock_db.session.get.return_value = refreshed_segment

            result = SegmentService.update_segment(args, segment, document, dataset)

        assert result is refreshed_segment
        generate_summary.assert_called_once_with(segment, dataset, {"enable": True})
        update_summary.assert_not_called()
        vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset)

    def test_delete_segment_removes_index_and_updates_document_word_count(self):
        segment = _make_segment(word_count=4, index_node_id="parent-node")
        document = _make_document(word_count=10)
        dataset = _make_dataset()

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.delete_segment_from_index_task") as delete_task,
        ):
            mock_redis.get.return_value = None
            mock_db.session.scalars.return_value.all.return_value = ["child-1", "child-2"]

            SegmentService.delete_segment(segment, document, dataset)

        assert document.word_count == 6
        mock_redis.setex.assert_called_once_with(f"segment_{segment.id}_delete_indexing", 600, 1)
        delete_task.delay.assert_called_once_with(
            ["parent-node"],
            dataset.id,
            document.id,
            [segment.id],
            ["child-1", "child-2"],
        )
        mock_db.session.delete.assert_called_once_with(segment)
        mock_db.session.add.assert_called_once_with(document)
        mock_db.session.commit.assert_called_once()

    def test_delete_segment_rejects_when_delete_is_already_in_progress(self):
        segment = _make_segment()
        document = _make_document()
        dataset = _make_dataset()

        with patch("services.dataset_service.redis_client") as mock_redis:
            mock_redis.get.return_value = "1"

            with pytest.raises(ValueError, match="Segment is deleting"):
                SegmentService.delete_segment(segment, document, dataset)

    def test_delete_segments_removes_records_and_clamps_document_word_count(self):
        dataset = _make_dataset()
        document = _make_document(word_count=3)
        current_user = SimpleNamespace(current_tenant_id="tenant-1")

        with (
            patch("services.dataset_service.current_user", current_user),
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.delete_segment_from_index_task") as delete_task,
        ):
            # execute().all() for segments_info (multi-column)
            execute_result = MagicMock()
            execute_result.all.return_value = [
                ("node-1", "segment-1", 2),
                ("node-2", "segment-2", 5),
            ]
            mock_db.session.execute.return_value = execute_result
            # scalars() for child_node_ids
            mock_db.session.scalars.return_value.all.return_value = ["child-1"]

            SegmentService.delete_segments(["segment-1", "segment-2"], document, dataset)

        assert document.word_count == 0
        mock_db.session.add.assert_called_once_with(document)
        delete_task.delay.assert_called_once_with(
            ["node-1", "node-2"],
            dataset.id,
            document.id,
            ["segment-1", "segment-2"],
            ["child-1"],
        )
        mock_db.session.commit.assert_called_once()

    def test_update_segments_status_enables_only_segments_without_indexing_cache(self):
        dataset = _make_dataset()
        document = _make_document()
        segment_a = _make_segment(segment_id="segment-a", enabled=False)
        segment_b = _make_segment(segment_id="segment-b", enabled=False)
        current_user = SimpleNamespace(id="user-1", current_tenant_id="tenant-1")

        with (
            patch("services.dataset_service.current_user", current_user),
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.enable_segments_to_index_task") as enable_task,
        ):
            mock_db.session.scalars.return_value.all.return_value = [segment_a, segment_b]
            mock_redis.get.side_effect = [None, "1"]

            SegmentService.update_segments_status(["segment-a", "segment-b"], "enable", dataset, document)

        assert segment_a.enabled is True
        assert segment_a.disabled_at is None
        assert segment_a.disabled_by is None
        assert segment_b.enabled is False
        mock_db.session.add.assert_called_once_with(segment_a)
        mock_db.session.commit.assert_called_once()
        enable_task.delay.assert_called_once_with(["segment-a"], dataset.id, document.id)

    def test_update_segments_status_disables_only_segments_without_indexing_cache(self):
        dataset = _make_dataset()
        document = _make_document()
        segment_a = _make_segment(segment_id="segment-a", enabled=True)
        segment_b = _make_segment(segment_id="segment-b", enabled=True)
        current_user = SimpleNamespace(id="user-1", current_tenant_id="tenant-1")

        with (
            patch("services.dataset_service.current_user", current_user),
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.disable_segments_from_index_task") as disable_task,
        ):
            mock_db.session.scalars.return_value.all.return_value = [segment_a, segment_b]
            mock_redis.get.side_effect = [None, "1"]

            SegmentService.update_segments_status(["segment-a", "segment-b"], "disable", dataset, document)

        assert segment_a.enabled is False
        assert segment_a.disabled_at == "now"
        assert segment_a.disabled_by == current_user.id
        assert segment_b.enabled is True
        mock_db.session.add.assert_called_once_with(segment_a)
        mock_db.session.commit.assert_called_once()
        disable_task.delay.assert_called_once_with(["segment-a"], dataset.id, document.id)


class TestSegmentServiceChildChunkTailHelpers:
    """Unit tests for the remaining child-chunk helper branches."""

    def test_update_child_chunk_rolls_back_on_vector_failure(self):
        dataset = SimpleNamespace(id="dataset-1")
        child_chunk = _make_child_chunk()

        with (
            patch("services.dataset_service.current_user", SimpleNamespace(id="user-1")),
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            vector_service.update_child_chunk_vector.side_effect = RuntimeError("vector failed")

            with pytest.raises(ChildChunkIndexingError, match="vector failed"):
                SegmentService.update_child_chunk(
                    "new content", child_chunk, SimpleNamespace(), SimpleNamespace(), dataset
                )

        mock_db.session.rollback.assert_called_once()
        mock_db.session.commit.assert_not_called()

    def test_delete_child_chunk_commits_after_successful_vector_delete(self):
        dataset = SimpleNamespace(id="dataset-1")
        child_chunk = _make_child_chunk()

        with (
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            SegmentService.delete_child_chunk(child_chunk, dataset)

        mock_db.session.delete.assert_called_once_with(child_chunk)
        vector_service.delete_child_chunk_vector.assert_called_once_with(child_chunk, dataset)
        mock_db.session.commit.assert_called_once()


class TestSegmentServiceAdditionalRegenerationBranches:
    """Additional unit tests for segment update and regeneration edge cases."""

    @pytest.fixture
    def account_context(self):
        account = create_autospec(Account, instance=True)
        account.id = "user-1"
        account.current_tenant_id = "tenant-1"

        with patch("services.dataset_service.current_user", account):
            yield account

    def test_update_segment_same_content_updates_answer_and_document_word_count_for_qa_segments(self, account_context):
        segment = _make_segment(content="question", word_count=8)
        document = _make_document(doc_form=IndexStructureType.QA_INDEX, word_count=20)
        dataset = _make_dataset()
        refreshed_segment = SimpleNamespace(id=segment.id)

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            mock_redis.get.return_value = None
            mock_db.session.get.return_value = refreshed_segment

            result = SegmentService.update_segment(
                SegmentUpdateArgs(content="question", answer="new answer"),
                segment,
                document,
                dataset,
            )

        assert result is refreshed_segment
        assert segment.answer == "new answer"
        assert segment.word_count == len("question") + len("new answer")
        assert document.word_count == 20 + (len("question") + len("new answer") - 8)
        vector_service.update_segment_vector.assert_not_called()
        vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset)

    def test_update_segment_content_change_uses_answer_when_counting_tokens_for_qa_segments(self, account_context):
        segment = _make_segment(content="old", word_count=3)
        document = _make_document(doc_form=IndexStructureType.QA_INDEX, word_count=10)
        dataset = _make_dataset(indexing_technique="high_quality")
        refreshed_segment = SimpleNamespace(id=segment.id)
        embedding_model = MagicMock()
        embedding_model.get_text_embedding_num_tokens.return_value = [21]

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.ModelManager") as model_manager_cls,
            patch("services.dataset_service.VectorService") as vector_service,
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-qa"),
            patch("services.dataset_service.naive_utc_now", return_value="now"),
        ):
            mock_redis.get.return_value = None
            model_manager_cls.for_tenant.return_value.get_model_instance.return_value = embedding_model
            mock_db.session.scalar.return_value = None
            mock_db.session.get.return_value = refreshed_segment

            result = SegmentService.update_segment(
                SegmentUpdateArgs(content="new question", answer="new answer", keywords=["kw-1"]),
                segment,
                document,
                dataset,
            )

        assert result is refreshed_segment
        embedding_model.get_text_embedding_num_tokens.assert_called_once_with(texts=["new questionnew answer"])
        assert segment.answer == "new answer"
        assert segment.tokens == 21
        assert segment.word_count == len("new question") + len("new answer")
        vector_service.update_segment_vector.assert_called_once_with(["kw-1"], segment, dataset)
        vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset)

    def test_update_segment_content_change_parent_child_uses_default_embedding_and_ignores_summary_failures(
        self, account_context
    ):
        segment = _make_segment(content="old", word_count=3)
        document = _make_document(
            doc_form=IndexStructureType.PARENT_CHILD_INDEX,
            word_count=10,
        )
        dataset = _make_dataset(indexing_technique="high_quality")
        dataset.embedding_model_provider = None
        refreshed_segment = SimpleNamespace(id=segment.id)
        processing_rule = SimpleNamespace(id=document.dataset_process_rule_id)
        existing_summary = SimpleNamespace(summary_content="old summary")
        embedding_model_instance = object()

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.ModelManager") as model_manager_cls,
            patch("services.dataset_service.VectorService") as vector_service,
            patch("services.dataset_service.helper.generate_text_hash", return_value="hash-parent"),
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.summary_index_service.SummaryIndexService.update_summary_for_segment") as update_summary,
        ):
            mock_redis.get.return_value = None
            model_manager_cls.for_tenant.return_value.get_default_model_instance.return_value = embedding_model_instance
            update_summary.side_effect = RuntimeError("summary failed")

            # get calls: processing_rule, then refreshed_segment
            mock_db.session.get.side_effect = [processing_rule, refreshed_segment]
            # scalar call: existing_summary
            mock_db.session.scalar.return_value = existing_summary

            result = SegmentService.update_segment(
                SegmentUpdateArgs(content="new parent content", regenerate_child_chunks=True, summary="new summary"),
                segment,
                document,
                dataset,
            )

        assert result is refreshed_segment
        model_manager_cls.for_tenant.return_value.get_default_model_instance.assert_called_once_with(
            tenant_id="tenant-1",
            model_type=ModelType.TEXT_EMBEDDING,
        )
        vector_service.generate_child_chunks.assert_called_once_with(
            segment,
            document,
            dataset,
            embedding_model_instance,
            processing_rule,
            True,
        )
        update_summary.assert_called_once_with(segment, dataset, "new summary")
        vector_service.update_multimodel_vector.assert_called_once_with(segment, [], dataset)

    def test_update_segment_same_content_parent_child_marks_segment_error_for_non_high_quality_dataset(
        self, account_context
    ):
        segment = _make_segment(content="same content", word_count=len("same content"))
        document = _make_document(
            doc_form=IndexStructureType.PARENT_CHILD_INDEX,
            word_count=20,
        )
        dataset = _make_dataset(indexing_technique="economy")
        refreshed_segment = SimpleNamespace(id=segment.id)

        with (
            patch("services.dataset_service.redis_client") as mock_redis,
            patch("services.dataset_service.db") as mock_db,
            patch("services.dataset_service.naive_utc_now", return_value="now"),
            patch("services.dataset_service.VectorService") as vector_service,
        ):
            mock_redis.get.return_value = None
            mock_db.session.get.return_value = refreshed_segment

            result = SegmentService.update_segment(
                SegmentUpdateArgs(content="same content", regenerate_child_chunks=True),
                segment,
                document,
                dataset,
            )

        assert result is refreshed_segment
        assert segment.enabled is False
        assert segment.disabled_at == "now"
        assert segment.status == "error"
        assert segment.error == "The knowledge base index technique is not high quality!"
        vector_service.update_multimodel_vector.assert_not_called()
