"""DataHub Documents Source.

This source:
1. Fetches Document entities from DataHub (via GraphQL or MCL events)
2. Partitions Document.text as markdown using unstructured.io
3. Chunks text using semantic chunking strategies
4. Generates embeddings via provider SDKs (Bedrock, Cohere, OpenAI, Vertex AI)
5. Emits SemanticContent aspects to DataHub

Supports both batch (GraphQL) and event-driven (Kafka MCL) modes.
"""

import hashlib
import json
import logging
import re
import time
from dataclasses import field
from datetime import datetime
from itertools import islice
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, cast

from datahub.configuration.common import GraphError, OperationalError
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.decorators import (
    SupportStatus,
    config_class,
    platform_name,
    support_status,
)
from datahub.ingestion.api.source import SourceReport
from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.ingestion.graph.client import DataHubGraph
from datahub.ingestion.graph.config import DatahubClientConfig
from datahub.ingestion.source.datahub_documents.datahub_documents_config import (
    DataHubDocumentsSourceConfig,
)
from datahub.ingestion.source.datahub_documents.document_chunking_state_handler import (
    DocumentChunkingStateHandler,
)
from datahub.ingestion.source.datahub_documents.document_indexing_lock import (
    DocumentIndexingLock,
)
from datahub.ingestion.source.datahub_documents.text_partitioner import TextPartitioner
from datahub.ingestion.source.state.stateful_ingestion_base import (
    StatefulIngestionConfig,
    StatefulIngestionConfigBase,
    StatefulIngestionReport,
    StatefulIngestionSourceBase,
)
from datahub.ingestion.source.unstructured.chunking_config import (
    DataHubConnectionConfig,
    DocumentChunkingSourceConfig,
)
from datahub.ingestion.source.unstructured.chunking_source import DocumentChunkingSource
from datahub.ingestion.source.unstructured.event_consumer import DocumentEventConsumer

logger = logging.getLogger(__name__)

# Aspects whose changes can change what a document's embedding should be built
# from: documentInfo carries the body, the standalone semanticText aspect
# carries the curated override. Passed to DocumentEventConsumer AND checked in
# _process_single_event — keep the two filters in lockstep via this constant.
EMBED_SOURCE_ASPECT_NAMES = ("documentInfo", "semanticText")


class DocumentEnumerationError(Exception):
    """Raised when document enumeration or hydration breaks systemically.

    Distinct from ``GraphError``, which hydration treats as a per-document data
    problem worth retrying URN by URN. This one means the document set can no
    longer be trusted, so the run must abort rather than embed a silent subset.
    """


class DataHubDocumentsReport(StatefulIngestionReport):
    """Report for DataHub documents source."""

    num_documents_fetched: int = 0
    num_documents_processed: int = 0
    num_documents_skipped: int = 0
    num_documents_skipped_unchanged: int = 0
    num_documents_skipped_empty: int = 0
    num_documents_skipped_existing_embeddings: int = 0
    num_documents_skipped_orphaned: int = 0
    num_chunks_created: int = 0
    lock_skipped_run: bool = False
    num_embeddings_generated: int = 0
    num_embedding_failures: int = 0
    embedding_failures: list[str] = field(default_factory=list)
    processing_errors: list[str] = []
    num_documents_limit_reached: bool = False

    def report_document_fetched(self) -> None:
        self.num_documents_fetched += 1

    def report_document_processed(self, num_chunks: int) -> None:
        self.num_documents_processed += 1
        self.num_chunks_created += num_chunks

    def report_document_skipped(self) -> None:
        self.num_documents_skipped += 1

    def report_document_skipped_unchanged(self) -> None:
        """Report document skipped due to unchanged content hash."""
        self.num_documents_skipped_unchanged += 1

    def report_document_skipped_empty(self) -> None:
        """Report document skipped due to empty or too-short text."""
        self.num_documents_skipped_empty += 1

    def report_document_skipped_existing_embeddings(self) -> None:
        """Report document skipped because it already has semanticContent for the model."""
        self.num_documents_skipped_existing_embeddings += 1

    def report_document_skipped_orphaned(self) -> None:
        """Report a URN skipped because it has no resolvable backing entity."""
        self.num_documents_skipped_orphaned += 1

    def report_embeddings_generated(self, count: int) -> None:
        self.num_embeddings_generated += count

    def report_error(self, error: str) -> None:
        self.processing_errors.append(error)


@platform_name("DataHubDocuments", id="datahub-documents")
@support_status(SupportStatus.INCUBATING)
@config_class(DataHubDocumentsSourceConfig)
class DataHubDocumentsSource(StatefulIngestionSourceBase):
    """
    Extract Document entities from DataHub and generate semantic embeddings for semantic search.

    Supports batch mode (GraphQL) and event-driven mode (Kafka MCL) with incremental processing.
    Automatically fetches embedding configuration from server to ensure alignment.
    """

    def __init__(self, ctx: PipelineContext, config: DataHubDocumentsSourceConfig):
        """Initialize DataHubDocumentsSource.

        Args:
            ctx: Pipeline context
            config: Source configuration
        """
        # Cast config to base type for super().__init__ call
        # This is safe because DocumentChunkingStatefulIngestionConfig is a subclass of StatefulIngestionConfig,
        # and StatefulIngestionConfigBase uses invariant generics, so mypy requires explicit cast
        base_config = cast(StatefulIngestionConfigBase[StatefulIngestionConfig], config)
        super().__init__(base_config, ctx)
        self.config = config
        # Override report with our custom report type
        # Type annotation ensures mypy knows the correct type
        self.report: DataHubDocumentsReport = DataHubDocumentsReport()

        # Initialize state handler for stateful ingestion
        self.state_handler: Optional[DocumentChunkingStateHandler] = (
            DocumentChunkingStateHandler(
                source=self,
                config=self.config,
                pipeline_name=self.ctx.pipeline_name,
                run_id=self.ctx.run_id,
            )
            if self.state_provider.is_stateful_ingestion_configured()
            else None
        )

        # Initialize DataHub client
        # Use pipeline context graph if available (automatically configured from sink)
        # This ensures we connect to the correct GMS in managed ingestion environments
        if self.ctx.graph:
            self.graph = self.ctx.graph
        else:
            # Fallback: Create graph from config (for standalone usage)
            # Convert TransparentSecretStr to str for DatahubClientConfig
            token_str = (
                self.config.datahub.token.get_secret_value()
                if self.config.datahub.token
                else None
            )
            graph_config = DatahubClientConfig(
                server=self.config.datahub.server,
                token=token_str,
            )
            self.graph = DataHubGraph(config=graph_config)

        # Initialize text partitioner
        self.text_partitioner = TextPartitioner()

        # Initialize chunking/embedding sub-component (same pattern as NotionSource)
        chunking_config = DocumentChunkingSourceConfig(
            datahub=DataHubConnectionConfig(),
            chunking=self.config.chunking,
            embedding=self.config.embedding,
            max_documents=self.config.max_documents,
        )
        self.chunking_source = DocumentChunkingSource(
            ctx=ctx,
            config=chunking_config,
            standalone=False,
            graph=self.graph,
        )

        # Initialize state tracking for incremental mode
        self.document_state: dict[str, dict[str, Any]] = {}
        self.state_file_path: Optional[Path] = None

        if self.config.incremental.enabled:
            self._initialize_state_tracking()

        # Initialize the distributed lock used to prevent overlapping scheduled runs.
        # We use the internal ``dataHubStepState`` entity as a synchronous key/value store:
        # each lock is a dataHubStepState whose properties map holds the lease payload.
        #
        # To manually clear a stuck lock (rarely needed -- a crashed run's lease expires on
        # its own after locking.lock_ttl_seconds), delete the backing dataHubStepState
        # entity. Its URN is urn:li:dataHubStepState:<lock_id>, logged below and on every
        # acquire/release:
        #     datahub delete --urn "urn:li:dataHubStepState:<lock_id>" --hard -f
        self.lock: Optional[DocumentIndexingLock] = None
        if self.config.locking.enabled:
            lock_id = self.config.locking.lock_id or self._default_lock_id(
                self.ctx.pipeline_name
            )
            self.lock = DocumentIndexingLock(
                graph=self.graph,
                lock_id=lock_id,
                run_id=self.ctx.run_id,
                ttl_seconds=self.config.locking.lock_ttl_seconds,
                renewal_interval_seconds=self.config.locking.lock_renewal_interval_seconds,
            )
            logger.info(
                f"Document indexing lock enabled: urn={self.lock.urn}, "
                f"ttl={self.config.locking.lock_ttl_seconds}s, "
                f"renewal_interval={self.config.locking.lock_renewal_interval_seconds}s. "
                f'To clear it manually: datahub delete --urn "{self.lock.urn}" --hard -f'
            )

        logger.info(
            f"Initialized DataHubDocumentsSource with platforms: {self.config.platform_filter}, "
            f"mode: {'event-driven' if self.config.event_mode.enabled else 'batch'}, "
            f"incremental: {self.config.incremental.enabled}"
        )

    @staticmethod
    def _default_lock_id(pipeline_name: Optional[str]) -> str:
        """Derive a clean, per-source lock id from the pipeline name.

        In managed ingestion, ``pipeline_name`` is the ingestion source URN
        (e.g. ``urn:li:dataHubIngestionSource:datahub-documents``). We key the lock on the
        source id so all scheduled runs of the same source contend for one lock, while
        unrelated sources don't. We strip the URN wrapper and sanitize URN-reserved
        characters so the lock's own dataHubStepState URN id stays clean (no nested URNs).
        """
        name = pipeline_name or "default"
        prefix = "urn:li:dataHubIngestionSource:"
        if name.startswith(prefix):
            name = name[len(prefix) :]
        # Replace anything outside the safe id charset (handles colons, parens, commas).
        name = re.sub(r"[^A-Za-z0-9_.-]", "_", name)
        return f"document-indexing-lock-{name}"

    def _initialize_state_tracking(self) -> None:
        """Initialize state tracking for incremental mode."""
        if self.config.incremental.state_file_path:
            self.state_file_path = Path(self.config.incremental.state_file_path)
        else:
            # Default: ~/.datahub/document_chunking_state/{pipeline_name}.json
            state_dir = Path.home() / ".datahub" / "document_chunking_state"
            state_dir.mkdir(parents=True, exist_ok=True)
            pipeline_name = self.ctx.pipeline_name or "default"
            self.state_file_path = state_dir / f"{pipeline_name}.json"

        # Load existing state
        self._load_state()
        logger.info(
            f"Incremental mode enabled, state file: {self.state_file_path}, "
            f"tracking {len(self.document_state)} documents"
        )

    def _load_state(self) -> None:
        """Load state from file."""
        if self.state_file_path and self.state_file_path.exists():
            try:
                with open(self.state_file_path) as f:
                    self.document_state = json.load(f)
                logger.debug(f"Loaded state for {len(self.document_state)} documents")
            except Exception as e:
                logger.warning(f"Failed to load state file: {e}")
                self.document_state = {}

    def _save_state(self) -> None:
        """Save state to file."""
        if self.state_file_path:
            try:
                with open(self.state_file_path, "w") as f:
                    json.dump(self.document_state, f, indent=2)
                logger.debug(f"Saved state for {len(self.document_state)} documents")
            except Exception as e:
                logger.warning(f"Failed to save state file: {e}")

    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
        """Main entry point - route to batch or event mode."""
        # Acquire the distributed lock before doing any work so overlapping scheduled
        # runs do not duplicate scroll + indexing work. If another run holds it, exit
        # cleanly without processing (this is expected, not a failure).
        if self.lock is not None and not self.lock.acquire():
            self.report.lock_skipped_run = True
            self.report.warning(
                title="Run skipped (lock held)",
                message="Another datahub-documents run is in progress; this run exited "
                "without processing to avoid duplicate work. Adjust the schedule or "
                "locking.lock_ttl_seconds if this happens persistently.",
                log=False,
            )
            return

        try:
            if self.config.event_mode.enabled:
                yield from self._process_event_mode()
            else:
                yield from self._process_batch_mode()
        except DocumentEnumerationError as e:
            # Already surfaced by the enumeration/hydration path under a specific
            # title. Adding the generic failure below would give operators two
            # entries for one problem, the vaguer one last.
            logger.error(f"Aborting run: {e}", exc_info=True)
        except Exception as e:
            logger.error(f"Failed to run Unstructured pipeline: {e}", exc_info=True)
            self.report.failure(message="Failed to run Unstructured pipeline", exc=e)
        finally:
            # Save state after processing
            if self.config.incremental.enabled:
                self._save_state()
            if self.lock is not None:
                self.lock.release()

    def _process_batch_mode(self) -> Iterable[MetadataWorkUnit]:
        """Process documents using GraphQL search (batch mode)."""
        logger.info("Running in batch mode")

        # Fetch documents from DataHub
        documents = self._fetch_documents_graphql()

        # Process each document
        for doc in documents:
            # Keep the distributed lock lease alive during long runs.
            if self.lock is not None:
                self.lock.heartbeat()

            # EXTERNAL documents already semantically indexed by their own source (e.g.
            # Notion/Confluence emit semanticContent themselves) are owned by that source:
            # skip them and do NOT record them in our state, so we never claim ownership.
            if self._is_indexed_by_source(
                doc["urn"], doc.get("source_type") == "EXTERNAL"
            ):
                logger.debug(
                    f"Skipping document {doc['urn']} (already indexed by its source)"
                )
                self.report.report_document_skipped_existing_embeddings()
                continue

            # Check if we should process this document (incremental mode)
            if (
                self.config.incremental.enabled
                and not self.config.incremental.force_reprocess
            ):
                if not self._should_process(doc["urn"], doc.get("text", "")):
                    logger.debug(
                        f"Skipping document {doc['urn']} (unchanged content hash)"
                    )
                    self.report.report_document_skipped_unchanged()
                    continue

            # Process document and yield workunits
            yield from self._process_document_with_throttle(doc)

            # Update state after successful processing (we own this document now).
            if self.config.incremental.enabled:
                self._update_document_state(doc["urn"], doc.get("text", ""))

    def _bootstrap_event_mode_offsets(self, consumer_id: str) -> None:
        """Bootstrap event mode by capturing current offsets BEFORE batch mode.

        This prevents race conditions: batch mode processes all documents up to the
        captured offset, and subsequent runs continue from that offset.

        Args:
            consumer_id: Consumer ID for the event consumer
        """
        if not self.state_handler or not self.state_handler.is_checkpointing_enabled():
            return

        logger.info(
            "Bootstrapping event mode: capturing current offsets BEFORE batch processing..."
        )
        # Create temporary consumer to get current offsets
        bootstrap_consumer = DocumentEventConsumer(
            graph=self.graph,
            consumer_id=consumer_id,
            topics=self.config.event_mode.topics,
            lookback_days=None,  # No lookback - just get current position
            reset_offsets=False,
            idle_timeout_seconds=5,
            poll_timeout_seconds=2,
            poll_limit=1,  # Minimal poll - we just want the offset
            state_handler=None,  # Don't need state handler for bootstrap
        )
        try:
            for topic in self.config.event_mode.topics:
                current_offset = bootstrap_consumer.get_current_offset(topic)
                if current_offset:
                    self.state_handler.update_event_offset(topic, current_offset)
                    logger.info(
                        f"Captured current offset for {topic}: {current_offset} "
                        "(before batch processing)"
                    )
                else:
                    logger.warning(
                        f"Failed to capture current offset for {topic} - "
                        "bootstrap incomplete, but batch mode will still process documents"
                    )
        except Exception as e:
            logger.warning(
                f"Error during bootstrap offset capture: {e}. "
                "Batch mode will still process documents.",
                exc_info=True,
            )
        finally:
            bootstrap_consumer.close()

    def _process_single_event(
        self, event: dict[str, Any]
    ) -> Iterable[MetadataWorkUnit]:
        """Process a single MCL event and yield work units if document should be processed.

        Args:
            event: MCL event dictionary with entityUrn, aspectName, aspect, etc.

        Yields:
            MetadataWorkUnit objects for the document if it should be processed, empty otherwise
        """
        entity_urn = event.get("entityUrn")
        aspect_name = event.get("aspectName")
        aspect_data = event.get("aspect")

        if not entity_urn or not aspect_name:
            logger.debug("Event missing entityUrn or aspectName, skipping")
            return

        if aspect_name not in EMBED_SOURCE_ASPECT_NAMES:
            return

        # Parse aspect data (handles Avro union wrapping)
        try:
            aspect_dict = self._parse_mcl_aspect(aspect_data)
        except Exception as e:
            logger.warning(
                f"Failed to parse MCL aspect for {entity_urn}: {e}",
                exc_info=True,
            )
            return

        if aspect_name == "semanticText":
            # The event carries the override; the document body and source type
            # live on documentInfo, so fetch it.
            info_dict = self._fetch_document_info_dict(entity_urn)
            if info_dict is None:
                logger.debug(
                    f"semanticText event for {entity_urn} without readable documentInfo, skipping"
                )
                return
            contents = dict(info_dict.get("contents") or {})
            contents["semanticText"] = aspect_dict.get("text")
            # Downstream source-type filtering reads from the documentInfo shape.
            aspect_dict = info_dict
        else:
            # documentInfo event: the override (if any) lives in the standalone
            # semanticText aspect, so fetch it. Contents may be null for
            # partial entities.
            contents = dict(aspect_dict.get("contents") or {})
            contents["semanticText"] = self._fetch_semantic_text(entity_urn)

        # semanticText overrides text as the embedding source; see _resolve_embed_text.
        text = self._resolve_embed_text(contents)

        if not text:
            logger.debug(f"No text content in document {entity_urn}")
            return

        # Filter by source type (NATIVE vs EXTERNAL)
        # Default to NATIVE if source or sourceType is not set (backward compatibility with old documents)
        source = aspect_dict.get("source") or {}
        source_type = source.get("sourceType") or "NATIVE"

        # Determine if we should process this document based on source type and platform filter
        should_process = self._should_process_by_source_type_event(
            source_type, aspect_dict, entity_urn
        )
        if not should_process:
            return

        # EXTERNAL documents already semantically indexed by their own source (e.g.
        # Notion/Confluence emit semanticContent themselves) are owned by that source:
        # skip them and do NOT record them in our state, so we never claim ownership.
        if self._is_indexed_by_source(entity_urn, source_type == "EXTERNAL"):
            logger.debug(
                f"Skipping document {entity_urn} (already indexed by its source)"
            )
            self.report.report_document_skipped_existing_embeddings()
            return

        # Check if we should process (incremental mode)
        if (
            self.config.incremental.enabled
            and not self.config.incremental.force_reprocess
        ):
            if not self._should_process(entity_urn, text):
                logger.debug(f"Skipping document {entity_urn} (unchanged content hash)")
                self.report.report_document_skipped_unchanged()
                return

        # Build document dict
        doc = {"urn": entity_urn, "text": text, "source_type": source_type}
        self.report.report_document_fetched()

        # Process document and yield work units
        yield from self._process_document_with_throttle(doc)

        # Update state (we own this document now).
        if self.config.incremental.enabled:
            self._update_document_state(entity_urn, text)

    def _process_event_mode(self) -> Iterable[MetadataWorkUnit]:
        """Process documents using Kafka MCL events (event-driven mode).

        Automatically falls back to batch mode if:
        - State not present (first run)
        - Event API errors occur
        - Offset loading fails
        - State handler unavailable

        When falling back, bootstraps by capturing current offset BEFORE batch mode
        to prevent race conditions.
        """
        logger.info("Running in event-driven mode")

        # Set consumer_id if not provided
        consumer_id = self.config.event_mode.consumer_id or (
            f"datahub-documents-{self.ctx.pipeline_name or 'default'}"
        )

        # Check if we should fall back to batch mode before starting
        should_fallback = self._should_fallback_to_batch_mode()
        if should_fallback:
            logger.warning(
                "Event mode conditions not met, falling back to batch mode. "
                "This ensures all documents are processed even if event mode is unavailable."
            )

            # BEFORE batch mode - capture current offset to prevent race conditions
            self._bootstrap_event_mode_offsets(consumer_id)

            # NOW run batch mode - processes all documents up to the captured offset
            yield from self._process_batch_mode()

            if self.state_handler and self.state_handler.is_checkpointing_enabled():
                logger.info(
                    "Event mode bootstrapping complete. Next run will continue from captured offset."
                )
            return

        # Initialize event consumer with state handler for offset persistence
        event_consumer = DocumentEventConsumer(
            graph=self.graph,
            consumer_id=consumer_id,
            topics=self.config.event_mode.topics,
            lookback_days=self.config.event_mode.lookback_days,
            reset_offsets=self.config.event_mode.reset_offsets,
            idle_timeout_seconds=self.config.event_mode.idle_timeout_seconds,
            poll_timeout_seconds=self.config.event_mode.poll_timeout_seconds,
            poll_limit=self.config.event_mode.poll_limit,
            state_handler=self.state_handler,
            aspect_names=EMBED_SOURCE_ASPECT_NAMES,
        )

        try:
            # Track if we've processed any events
            events_processed = False

            # Process events
            # consume_events() already yields parsed MCL dicts with entityUrn, aspectName, etc.
            for event in event_consumer.consume_events():
                # Keep the distributed lock lease alive during long event runs.
                if self.lock is not None:
                    self.lock.heartbeat()
                events_processed = True
                yield from self._process_single_event(event)

            # If no events were processed and no lookback window, fall back to batch
            if not events_processed and not self.config.event_mode.lookback_days:
                logger.warning(
                    "No events processed in event mode and no lookback window set. "
                    "Falling back to batch mode to ensure documents are processed."
                )
                yield from self._process_batch_mode()
                return

        except DocumentEnumerationError:
            # A batch fallback above aborted systemically and already reported a
            # specific failure. Falling back again would replay a full
            # enumeration + hydration storm against GMS for the same outcome.
            raise
        except Exception as e:
            # Catch any errors during event processing and fall back to batch mode
            error_msg = str(e)
            logger.error(
                f"Error in event mode: {error_msg}. "
                "Falling back to batch mode to ensure documents are processed.",
                exc_info=True,
            )
            self.report.report_error(f"Event mode error: {error_msg}")

            # Fall back to batch mode
            yield from self._process_batch_mode()
            return

        finally:
            if event_consumer:
                event_consumer.close()

    def _parse_mcl_aspect(self, aspect_data: Any) -> dict[str, Any]:
        """Parse MCL aspect data (handles Avro union wrapping)."""
        if isinstance(aspect_data, dict):
            # Check for Avro union wrapper
            if "com.linkedin.pegasus2avro.mxe.GenericAspect" in aspect_data:
                generic_aspect = aspect_data[
                    "com.linkedin.pegasus2avro.mxe.GenericAspect"
                ]
                value_str = generic_aspect.get("value", "{}")
                return json.loads(value_str)
            return aspect_data
        elif isinstance(aspect_data, str):
            return json.loads(aspect_data)
        else:
            return {}

    def _fetch_semantic_text(self, document_urn: str) -> Optional[str]:
        """Read the standalone semanticText aspect (curated embedding-source
        override), or None when absent.

        On read failure we fall back to embedding the document body: a transient
        wrong embed self-corrects on a later run because incremental state hashes
        the embedded value.
        """
        from datahub.metadata.schema_classes import SemanticTextClass

        try:
            aspect = self.graph.get_aspect(
                entity_urn=document_urn, aspect_type=SemanticTextClass
            )
        except Exception as e:
            logger.warning(
                f"Failed to read semanticText for {document_urn}: {e}; "
                "embedding document text."
            )
            return None
        return aspect.text if aspect else None

    def _fetch_document_info_dict(self, document_urn: str) -> Optional[dict[str, Any]]:
        """Read the documentInfo aspect as a dict (the MCL payload shape), or
        None when absent or unreadable."""
        from datahub.metadata.schema_classes import DocumentInfoClass

        try:
            aspect = self.graph.get_aspect(
                entity_urn=document_urn, aspect_type=DocumentInfoClass
            )
        except Exception as e:
            logger.warning(
                f"Failed to read documentInfo for {document_urn}: {e}; skipping event."
            )
            return None
        return aspect.to_obj() if aspect else None

    def _extract_platform_from_aspect(
        self, aspect_dict: dict[str, Any]
    ) -> Optional[str]:
        """Extract platform name from documentInfo aspect."""
        # Try to get from dataPlatformInstance (may be null for partial entities)
        platform_instance = aspect_dict.get("dataPlatformInstance") or {}
        platform_urn = platform_instance.get("platform")
        if platform_urn and isinstance(platform_urn, str):
            # Extract platform name from URN (e.g., "urn:li:dataPlatform:notion" → "notion")
            parts = platform_urn.split(":")
            if len(parts) >= 3:
                return parts[-1]
        return None

    def _should_fallback_to_batch_mode(self) -> bool:
        """Check if we should fall back to batch mode instead of event mode.

        Returns True if:
        - State handler not available (stateful ingestion not configured)
        - No offsets found and no lookback window (first run scenario)
        - Stateful ingestion disabled
        """
        # If stateful ingestion is not enabled, we can't track offsets reliably
        if not self.state_handler or not self.state_handler.is_checkpointing_enabled():
            logger.info(
                "Stateful ingestion not enabled - cannot track event offsets reliably. "
                "Falling back to batch mode."
            )
            return True

        # Check if we have any offsets stored
        # If no offsets and no lookback window, it's likely a first run
        has_offsets = False
        for topic in self.config.event_mode.topics:
            offset = self.state_handler.get_event_offset(topic)
            if offset:
                has_offsets = True
                break

        if not has_offsets and not self.config.event_mode.lookback_days:
            logger.info(
                "No event offsets found and no lookback window configured. "
                "This appears to be a first run - falling back to batch mode to process all documents."
            )
            return True

        return False

    def _fetch_platform_from_entity(self, entity_urn: str) -> Optional[str]:
        """Fetch the platform name for a given entity URN via GraphQL.

        In event mode, we only receive documentInfo aspect changes, which doesn't include
        dataPlatformInstance. We need to fetch it separately to determine the platform.

        Args:
            entity_urn: Document URN

        Returns:
            Platform name (e.g., "notion", "confluence") or None if not found
        """
        try:
            query = """
            query getDocumentPlatform($urn: String!) {
                entity(urn: $urn) {
                    ... on Document {
                        dataPlatformInstance {
                            platform {
                                urn
                            }
                        }
                    }
                }
            }
            """
            response = self.graph.execute_graphql(query, {"urn": entity_urn})
            entity = response.get("entity") or {}
            platform_instance = entity.get("dataPlatformInstance") or {}
            platform = platform_instance.get("platform") or {}
            platform_urn = platform.get("urn")

            if platform_urn and isinstance(platform_urn, str):
                # Extract platform name from URN (e.g., "urn:li:dataPlatform:notion" → "notion")
                parts = platform_urn.split(":")
                if len(parts) >= 3:
                    return parts[-1]
            return None
        except Exception as e:
            logger.debug(f"Failed to fetch platform for {entity_urn}: {e}")
            return None

    def _should_process_by_source_type_event(
        self, source_type: Optional[str], aspect_dict: dict[str, Any], entity_urn: str
    ) -> bool:
        """Determine if document should be processed in event mode based on source type and platform filter.

        Logic:
        - NATIVE documents: Always process if platform_filter is empty, or if platform matches
        - EXTERNAL documents: Processed when include_external_documents is True (default).
          If platform_filter is set, EXTERNAL documents are restricted to those platforms.

        Args:
            source_type: Document source type ("NATIVE" or "EXTERNAL")
            aspect_dict: Parsed documentInfo aspect dictionary
            entity_urn: Document URN for logging and fetching platform

        Returns:
            True if document should be processed, False otherwise
        """
        # Wildcard means process everything
        if self.config.platform_filter and (
            "*" in self.config.platform_filter or "ALL" in self.config.platform_filter
        ):
            return True

        # NATIVE documents: Process if platform_filter is empty (all native) or matches platform
        if source_type == "NATIVE":
            # Empty platform_filter means "all native documents"
            if not self.config.platform_filter:
                return True
            # If platform_filter is specified, check if document's platform matches
            # Native documents typically have platform="datahub" or no platform
            # Note: documentInfo doesn't contain platform, but for NATIVE we can default to "datahub"
            # or fetch it if needed. For simplicity, we'll default to "datahub" for NATIVE.
            platform_name = "datahub"  # Default for native documents
            return platform_name in self.config.platform_filter

        # EXTERNAL documents
        if source_type == "EXTERNAL":
            if not self.config.include_external_documents:
                logger.debug(
                    f"Skipping document {entity_urn} (sourceType=EXTERNAL, include_external_documents is False)"
                )
                return False
            # No platform_filter: include all external documents.
            if not self.config.platform_filter:
                return True
            # platform_filter set: restrict external documents to those platforms.
            # documentInfo aspect doesn't contain platform - need to fetch it via GraphQL
            external_platform = self._fetch_platform_from_entity(entity_urn)
            if (
                external_platform is not None
                and external_platform in self.config.platform_filter
            ):
                return True
            logger.debug(
                f"Skipping document {entity_urn} (sourceType=EXTERNAL, platform={external_platform} not in platform_filter: {self.config.platform_filter})"
            )
            return False

        # Unknown source type - default to not processing
        logger.debug(
            f"Skipping document {entity_urn} (unknown sourceType: {source_type})"
        )
        return False

    def _should_process_by_source_type(
        self, entity: dict[str, Any], info: dict[str, Any]
    ) -> bool:
        """Determine if document should be processed in batch mode based on source type and platform filter.

        Logic:
        - NATIVE documents: Always process if platform_filter is empty, or if platform matches
        - EXTERNAL documents: Processed when include_external_documents is True (default).
          If platform_filter is set, EXTERNAL documents are restricted to those platforms.

        Args:
            entity: GraphQL entity response (includes dataPlatformInstance)
            info: Document info from GraphQL (includes source)

        Returns:
            True if document should be processed, False otherwise
        """
        # Wildcard means process everything
        if self.config.platform_filter and (
            "*" in self.config.platform_filter or "ALL" in self.config.platform_filter
        ):
            return True

        # Extract source type from info (source may be null for partial entities)
        source = info.get("source") or {}
        source_type = source.get("sourceType")
        # Default to NATIVE if sourceType is not set (backward compatibility with old documents)
        if source_type is None:
            source_type = "NATIVE"

        # NATIVE documents: Process if platform_filter is empty (all native) or matches platform
        if source_type == "NATIVE":
            # Empty platform_filter means "all native documents"
            if not self.config.platform_filter:
                return True
            # If platform_filter is specified, check if document's platform matches
            # Extract platform from entity (dataPlatformInstance)
            platform_name = self._extract_platform_from_entity(entity)
            if platform_name is None:
                platform_name = "datahub"  # Default for native documents
            return platform_name in self.config.platform_filter

        # EXTERNAL documents
        if source_type == "EXTERNAL":
            if not self.config.include_external_documents:
                logger.debug(
                    f"Skipping document {entity.get('urn', 'unknown')} (sourceType=EXTERNAL, include_external_documents is False)"
                )
                return False
            # No platform_filter: include all external documents.
            if not self.config.platform_filter:
                return True
            # platform_filter set: restrict external documents to those platforms.
            # Extract platform from entity (dataPlatformInstance)
            external_platform = self._extract_platform_from_entity(entity)
            if (
                external_platform is not None
                and external_platform in self.config.platform_filter
            ):
                return True
            logger.debug(
                f"Skipping document {entity.get('urn', 'unknown')} (sourceType=EXTERNAL, platform={external_platform} not in platform_filter: {self.config.platform_filter})"
            )
            return False

        # Unknown source type - default to not processing
        logger.debug(
            f"Skipping document {entity.get('urn', 'unknown')} (unknown sourceType: {source_type})"
        )
        return False

    def _extract_platform_from_entity(self, entity: dict[str, Any]) -> Optional[str]:
        """Extract platform name from GraphQL entity response."""
        # Try to get from dataPlatformInstance (fields may be null for partial entities)
        platform_instance = entity.get("dataPlatformInstance") or {}
        platform = platform_instance.get("platform") or {}
        platform_urn = platform.get("urn")
        if platform_urn and isinstance(platform_urn, str):
            # Extract platform name from URN (e.g., "urn:li:dataPlatform:notion" → "notion")
            parts = platform_urn.split(":")
            if len(parts) >= 3:
                return parts[-1]
        return None

    # URNs hydrated per GraphQL request. Kept modest so a single response stays
    # within GMS's GraphQL limits even for large documents.
    #
    # Deliberately fixed and independent of the configurable `scroll_batch_size`:
    # that one sizes an Elasticsearch page of URNs (cheap, tuned for scroll
    # pressure on ES), this one sizes a payload of full document bodies (tuned
    # for GMS response size). Tying them together would make raising the scroll
    # page size silently inflate hydration responses.
    _HYDRATE_BATCH_SIZE = 100

    def _fetch_documents_graphql(self) -> Iterable[dict[str, Any]]:
        """Stream Document entities from DataHub with their text content.

        URNs are enumerated (``_scroll_document_urns``) and hydrated
        (``_hydrate_documents``) lazily, so a single orphaned index entry can no
        longer abort the whole run and the catalog is never materialized in
        memory: at most one hydration batch is held at a time.

        Enumeration and hydration own their own error reporting, so failures
        propagate from here untouched rather than being re-wrapped around a
        ``yield``.
        """
        num_documents = 0

        for entity in self._hydrate_documents(self._scroll_document_urns()):
            urn = entity.get("urn")
            if not urn:
                continue

            # Filter by specific URNs if provided
            if self.config.document_urns and urn not in self.config.document_urns:
                continue

            # Extract text content (GraphQL returns null for missing aspects).
            # semanticText overrides text as the embedding source; see _resolve_embed_text.
            info = entity.get("info") or {}
            contents = info.get("contents") or {}
            text = self._resolve_embed_text(contents)

            # Default to NATIVE when sourceType is absent (older documents).
            source = info.get("source") or {}
            source_type = source.get("sourceType") or "NATIVE"

            # Filter by source type (NATIVE vs EXTERNAL) for batch mode
            if not self._should_process_by_source_type(entity, info):
                continue

            # Skip if no text or too short
            if not text or (
                self.config.skip_empty_text and len(text) < self.config.min_text_length
            ):
                logger.debug(
                    f"Skipping document {urn} (empty or too short: {len(text)} chars)"
                )
                continue

            num_documents += 1
            self.report.report_document_fetched()
            yield {"urn": urn, "text": text, "source_type": source_type}

        logger.info(
            f"Fetched {num_documents} documents with text content from platforms: {self.config.platform_filter}"
        )

    def _scroll_document_urns(self) -> Iterable[str]:
        """Enumerate Document URNs via scrollAcrossEntities.

        Only the URN is selected here — never an aspect-backed field. An
        orphaned index entry (a URN present in the search index with no backing
        entity) would make GMS null the non-null ``SearchResult.entity`` and
        abort the whole scroll if any aspect field were selected; requesting the
        URN alone reads it straight off the search hit and stays orphan-safe.
        Content is resolved separately in ``_hydrate_documents``, where orphans
        come back null and are skipped.

        scrollAcrossEntities uses cursor-based pagination and is not subject to
        Elasticsearch's max_result_window limit (default 10,000), unlike
        offset-based search.
        """
        query = """
        query scrollDocumentUrns(
            $scrollId: String,
            $batchSize: Int!,
            $orFilters: [AndFilterInput!]
        ) {
          scrollAcrossEntities(input: {
            types: [DOCUMENT],
            query: "*",
            count: $batchSize,
            scrollId: $scrollId,
            orFilters: $orFilters,
            searchFlags: {
              skipHighlighting: true,
              skipAggregates: true
              includeHiddenLifecycleStages: true
            }
          }) {
            nextScrollId
            searchResults {
              entity {
                urn
              }
            }
          }
        }
        """

        or_filters = None
        if (
            self.config.platform_filter
            and "*" not in self.config.platform_filter
            and "ALL" not in self.config.platform_filter
        ):
            or_filters = [
                {
                    "and": [
                        {
                            "field": "platform",
                            "values": [
                                f"urn:li:dataPlatform:{platform}"
                                for platform in self.config.platform_filter
                            ],
                        }
                    ]
                }
            ]

        scroll_id: Optional[str] = None
        first_iter = True
        while first_iter or scroll_id:
            # Keep the distributed lock lease alive while scrolling large sets.
            if self.lock is not None:
                self.lock.heartbeat()
            # Throttle between pages (not before the first) to reduce load on
            # GMS/Elasticsearch when scrolling large document sets.
            if not first_iter and self.config.scroll_delay_seconds > 0:
                time.sleep(self.config.scroll_delay_seconds)
            first_iter = False

            variables: dict[str, Any] = {
                "batchSize": self.config.scroll_batch_size,
                "scrollId": scroll_id,
                "orFilters": or_filters,
            }
            response = self.graph.execute_graphql(query, variables)
            scroll_data = response.get("scrollAcrossEntities")
            if scroll_data is None:
                # scrollAcrossEntities is nullable and execute_graphql only
                # raises when the response carries GraphQL `errors`, so a null
                # payload arrives here as an ordinary response. It means
                # enumeration itself broke (vs. an empty page, which is a valid
                # `searchResults: []`), so abort: a warning would leave the run
                # free to finish SUCCESS having embedded a silent subset.
                self.report.failure(
                    title="Document enumeration returned no scroll payload",
                    message="scrollAcrossEntities was missing from the GraphQL "
                    "response, so the document set is incomplete; aborting the "
                    "run rather than embedding an unknown subset.",
                )
                raise DocumentEnumerationError(
                    "scrollAcrossEntities missing from the GraphQL response"
                )
            scroll_id = scroll_data.get("nextScrollId")
            search_results = scroll_data.get("searchResults") or []
            logger.debug(
                f"Scrolled page of {len(search_results)} document URN(s) (scrollId={scroll_id})"
            )
            # A full page with no continuation cursor is suspicious: the scroll
            # likely truncated, so some documents may never be enumerated. The
            # page itself is usable, so keep embedding what was enumerated, but
            # fail the run — a warning would let a large silent gap pass as a
            # successful run.
            if not scroll_id and len(search_results) >= self.config.scroll_batch_size:
                self.report.failure(
                    title="Document enumeration ended without a scroll cursor",
                    message="A full page returned no nextScrollId; enumeration "
                    "stopped early and some documents are likely missing from "
                    "this run.",
                )
            for result in search_results:
                entity = result.get("entity") or {}
                urn = entity.get("urn")
                if urn:
                    yield urn

    def _hydrate_documents(self, urns: Iterable[str]) -> Iterable[dict[str, Any]]:
        """Resolve document URNs to entities in batches.

        Consumes ``urns`` lazily in ``_HYDRATE_BATCH_SIZE`` windows so the full
        catalog is never materialized: enumeration and hydration interleave, one
        batch at a time.

        Uses the nullable ``entities(urns:)`` query: an orphaned URN comes back
        as null and is counted and skipped, instead of aborting the run as the
        non-null ``SearchResult.entity`` in ``scrollAcrossEntities`` would.
        """
        query = """
        query hydrateDocuments($urns: [String!]!) {
          entities(urns: $urns) {
            urn
            type
            ... on Document {
              info {
                contents {
                  text
                  semanticText
                }
                customProperties {
                  key
                  value
                }
                source {
                  sourceType
                }
              }
              dataPlatformInstance {
                platform {
                  urn
                }
              }
            }
          }
        }
        """
        urn_iter = iter(urns)
        while True:
            # islice pulls one window from the enumerator, so scrolling and
            # hydration interleave instead of enumerating everything up front.
            batch = list(islice(urn_iter, self._HYDRATE_BATCH_SIZE))
            if not batch:
                break
            if self.lock is not None:
                self.lock.heartbeat()

            try:
                response = self.graph.execute_graphql(query, {"urns": batch})
            except OperationalError:
                # Transport/outage error — abort rather than silently skipping
                # every document (which would report success with 0 embeddings).
                raise
            except GraphError as batch_error:
                # One document with an unexpected null in a non-null field (e.g.
                # DocumentContent.text) makes GMS fail the whole batch. Retry the
                # batch one URN at a time so a single bad document can't drop the
                # rest.
                logger.warning(
                    f"Batch hydration failed for {len(batch)} URN(s); retrying individually"
                )
                hydrated_any = False
                for entity in self._hydrate_individually(query, batch):
                    hydrated_any = True
                    yield entity
                if not hydrated_any:
                    # Every URN in the batch failed the same way — a systemic
                    # error (typically a GraphQL schema mismatch, e.g. the GMS is
                    # missing a queried field), not per-document data. Abort:
                    # continuing would repeat batch-fail plus a per-URN retry
                    # storm for every remaining batch, hammering GMS with O(N)
                    # doomed requests before the run finally ends.
                    self.report.failure(
                        title="Document hydration failed for an entire batch",
                        message="Every URN in a hydration batch failed to resolve; "
                        "this usually means a GraphQL schema mismatch (the GMS is "
                        "missing a queried field) rather than per-document issues.",
                        context=batch[0],
                    )
                    raise DocumentEnumerationError(
                        "every URN in a hydration batch failed to resolve"
                    ) from batch_error
                continue

            entities = response.get("entities") or []
            if len(entities) != len(batch):
                # GMS returns exactly one (possibly null) slot per requested URN —
                # BatchGetEntitiesResolver sizes its array from the input list. Any
                # other length means the response cannot be aligned with the
                # request, so nothing in it is trustworthy. That includes an empty
                # list, which would otherwise read as a batch of orphans.
                self.report.failure(
                    title="Document hydration returned an unexpected entity count",
                    message="GMS returned a different number of entities than URNs "
                    "requested, so the hydrated document set is incomplete.",
                    context=f"requested={len(batch)}, returned={len(entities)}",
                )
                raise DocumentEnumerationError(
                    "entities(urns:) length did not match the hydration batch"
                )

            # Match on each entity's own URN rather than zipping positionally: a
            # reordered response would otherwise attribute one document's content
            # to another URN. Requested URNs with no entity in the response are
            # counted as orphans, never silently dropped.
            entities_by_urn = {
                entity["urn"]: entity
                for entity in entities
                if entity and entity.get("urn")
            }
            if not entities_by_urn:
                # Right length, but every slot came back null. Unlike a length
                # mismatch this is interpretable — every requested URN was
                # unresolvable — and the rest of the catalog may be healthy, so
                # fail the run without aborting it. A page of index drift must not
                # take down a run that can still embed everything else.
                self.report.failure(
                    title="Document hydration resolved nothing in a batch",
                    message="Every URN in a hydration batch resolved to null, which "
                    "is usually a serving problem rather than that many orphans; "
                    "the batch was skipped and the run continues.",
                    context=batch[0],
                )
            for urn in batch:
                hydrated = entities_by_urn.get(urn)
                if hydrated is None:
                    self._skip_orphaned(urn)
                    continue
                yield hydrated

    def _hydrate_individually(
        self, query: str, urns: list[str]
    ) -> Iterable[dict[str, Any]]:
        """Hydrate URNs one at a time so a single un-resolvable document is
        skipped instead of failing its whole batch."""
        for urn in urns:
            try:
                response = self.graph.execute_graphql(query, {"urns": [urn]})
            except OperationalError:
                raise
            except GraphError:
                self.report.warning(
                    title="Skipped document that failed to hydrate",
                    message="A document's content could not be resolved from GMS "
                    "(a non-null field returned null); it was skipped.",
                    context=urn,
                )
                continue
            # Match by URN for the same reason the batch path does, rather than
            # trusting the first slot to be the document we asked for.
            entities = response.get("entities") or []
            entity = next(
                (e for e in entities if e and e.get("urn") == urn),
                None,
            )
            if entity is None:
                self._skip_orphaned(urn)
                continue
            yield entity

    def _skip_orphaned(self, urn: str) -> None:
        """Record an orphaned index entry (a URN with no resolvable entity)."""
        self.report.report_document_skipped_orphaned()
        self.report.warning(
            title="Skipped orphaned document",
            message="A document URN is present in the search index but has no "
            "resolvable entity (likely index/entity-store drift); it was skipped.",
            context=urn,
        )

    def _should_process(self, document_urn: str, text: str) -> bool:
        """Check if document should be processed based on content hash."""
        if (
            not self.config.incremental.enabled
            or self.config.incremental.force_reprocess
        ):
            return True

        current_hash = self._calculate_text_hash(text)

        # Use state_handler if available (proper stateful ingestion)
        if self.state_handler and self.state_handler.is_checkpointing_enabled():
            previous_hash = self.state_handler.get_document_hash(document_urn)
            return previous_hash != current_hash

        # Fall back to local document_state (simple incremental mode)
        if document_urn not in self.document_state:
            return True

        previous_hash = self.document_state[document_urn].get("content_hash")
        return previous_hash != current_hash

    def _get_processing_config_fingerprint(self) -> Dict[str, Any]:
        """Extract processing configuration that affects output artifacts.

        This fingerprint is included in the document hash to trigger reprocessing
        when configuration changes (e.g., enabling chunking, changing embedding model).

        Returns:
            Dictionary of config values that affect processing output.
        """
        # Chunking/embedding is enabled when embedding provider is configured
        embedding_enabled = self.config.embedding.provider is not None

        return {
            # Chunking affects chunk boundaries and structure
            "chunking_enabled": embedding_enabled,
            "chunking_strategy": self.config.chunking.strategy
            if embedding_enabled
            else None,
            "chunking_max_characters": self.config.chunking.max_characters
            if embedding_enabled
            else None,
            "chunking_overlap": self.config.chunking.overlap
            if embedding_enabled
            else None,
            "chunking_combine_under": self.config.chunking.combine_text_under_n_chars
            if embedding_enabled
            else None,
            # Embedding affects vector embeddings on chunks
            "embedding_enabled": embedding_enabled,
            "embedding_provider": self.config.embedding.provider
            if embedding_enabled
            else None,
            "embedding_model": self.config.embedding.model
            if embedding_enabled
            else None,
            # Partitioning affects how text is extracted
            "partition_strategy": self.config.partition_strategy,
        }

    @staticmethod
    def _resolve_embed_text(contents: Dict[str, Any]) -> str:
        """Resolve the text to embed from a contents dict.

        `semanticText` overrides `text` as the embedding source when present, so
        a curated representation can drive retrieval while the full body stays in
        `text`. The override lives in the standalone `semanticText` aspect: batch
        mode still receives it under `contents` because the GraphQL surface maps
        the aspect there, and event mode injects it into the dict before calling
        this. Resolving here (rather than in the hash) means the content hash
        tracks the embedded value: editing `semanticText` re-embeds, and a change
        to `text` while `semanticText` is set does not.
        """
        return contents.get("semanticText") or contents.get("text") or ""

    def _calculate_text_hash(self, text: str) -> str:
        """Calculate hash of document content AND processing configuration.

        This ensures documents are reprocessed when:
        1. Content changes (text is different)
        2. Processing config changes (chunking enabled, embedding model changed, etc.)

        Args:
            text: Document text content

        Returns:
            SHA256 hash hex string
        """
        hash_input = {
            "content": text,
            "config": self._get_processing_config_fingerprint(),
        }
        # Deterministic JSON serialization
        hash_str = json.dumps(hash_input, sort_keys=True)
        return hashlib.sha256(hash_str.encode("utf-8")).hexdigest()

    def _update_document_state(self, document_urn: str, text: str) -> None:
        """Update state after processing document."""
        content_hash = self._calculate_text_hash(text)
        last_processed = datetime.utcnow().isoformat()

        # Use state_handler if available (proper stateful ingestion)
        if self.state_handler and self.state_handler.is_checkpointing_enabled():
            self.state_handler.update_document_state(
                document_urn, content_hash, last_processed
            )
        else:
            # Fall back to local document_state (simple incremental mode)
            self.document_state[document_urn] = {
                "content_hash": content_hash,
                "last_processed": last_processed,
            }

    def _state_has_document(self, document_urn: str) -> bool:
        """Whether incremental state already tracks this document.

        When state already has a record, this run reached the document because its content
        changed (the unchanged check in _should_process let it through), so we must
        re-embed and must NOT consult the existing semanticContent aspect.
        """
        if self.state_handler and self.state_handler.is_checkpointing_enabled():
            return self.state_handler.get_document_hash(document_urn) is not None
        return document_urn in self.document_state

    def _is_indexed_by_source(self, document_urn: str, is_external: bool) -> bool:
        """Whether an EXTERNAL document is already semantically indexed by its own source.

        Some ingestion sources (e.g. Notion, Confluence) perform their own semantic indexing
        and emit the semanticContent aspect themselves. For EXTERNAL documents we therefore
        treat an existing semanticContent aspect (for the active model) as a sign that the
        source owns indexing for that document: we leave it untouched -- we neither re-embed it
        nor record it in our incremental state. The purpose of this source for EXTERNAL
        documents is to pick up the ones the source did NOT index and take ownership of them
        (embed + track in state). NATIVE documents are owned by DataHub directly and are never
        skipped this way.

        We only consult the aspect when our incremental state has no record for the document:
        if state already tracks it, this run reached it because its content changed and we must
        re-embed (we already own it).
        """
        if not is_external:
            return False
        if not self.config.skip_external_if_semantic_content_exists:
            return False
        if self.config.incremental.force_reprocess:
            return False
        if self._state_has_document(document_urn):
            return False
        return self._has_existing_semantic_content(document_urn)

    def _has_existing_semantic_content(self, document_urn: str) -> bool:
        """Check whether the document already has a semanticContent aspect for the active model.

        Used to avoid recomputing embeddings when local incremental state is unavailable
        (e.g. after an ingestion of external documents that were indexed in the source). Returns False if no embedding model is configured or
        the aspect is absent for the current model key.
        """
        from datahub.metadata.schema_classes import SemanticContentClass

        model_key = self.chunking_source.get_model_embedding_key()
        if model_key is None:
            return False
        try:
            aspect = self.graph.get_aspect(
                entity_urn=document_urn, aspect_type=SemanticContentClass
            )
        except Exception as e:
            logger.debug(
                f"Failed to read semanticContent for {document_urn}: {e}; will process."
            )
            return False
        return bool(aspect and aspect.embeddings and model_key in aspect.embeddings)

    def _throttle_after_indexing(self) -> None:
        """Pause after indexing a document to smooth write load to GMS."""
        if self.config.index_delay_seconds > 0:
            time.sleep(self.config.index_delay_seconds)

    def _process_document_with_throttle(
        self, doc: dict[str, Any]
    ) -> Iterable[MetadataWorkUnit]:
        """Process one document and throttle only when it produced work units."""
        indexed = False
        for wu in self._process_single_document(doc):
            indexed = True
            yield wu
        if indexed:
            self._throttle_after_indexing()

    def _process_single_document(
        self, doc: dict[str, Any]
    ) -> Iterable[MetadataWorkUnit]:
        """Process a single document: partition text → chunk → embed → emit SemanticContent."""
        try:
            document_urn = doc["urn"]
            text = doc.get("text", "")

            # Skip if empty or too short
            if not text or len(text) < self.config.min_text_length:
                logger.debug(
                    f"Skipping document {document_urn} (text too short: {len(text)} chars)"
                )
                self.report.report_document_skipped_empty()
                return

            # Partition text as markdown into unstructured elements
            logger.debug(f"Partitioning text for document {document_urn}")
            elements = self.text_partitioner.partition_text(text)

            if not elements:
                logger.warning(
                    f"No elements created from text for document {document_urn}"
                )
                self.report.report_document_skipped()
                return

            # Delegate chunking + embedding + SemanticContent emission to chunking_source.
            # DocumentChunkingSource enforces max_documents and raises RuntimeError when exceeded.
            yield from self.chunking_source.process_elements_inline(
                document_urn=document_urn, elements=elements
            )

        except Exception as e:
            if self.chunking_source.report.num_documents_limit_reached:
                self.report.num_documents_limit_reached = True
                raise
            error_msg = f"Failed to process document {doc.get('urn', 'unknown')}: {e}"
            logger.warning(error_msg, exc_info=True)

    def get_report(self) -> SourceReport:
        # Forward stats from the chunking sub-component into our report
        self.report.num_documents_processed = (
            self.chunking_source.report.num_documents_processed
        )
        self.report.num_chunks_created = self.chunking_source.report.num_chunks_created
        self.report.num_embeddings_generated = (
            self.chunking_source.report.num_embeddings_generated
        )
        self.report.num_documents_limit_reached = (
            self.chunking_source.report.num_documents_limit_reached
        )
        self.report.num_embedding_failures = (
            self.chunking_source.report.num_embedding_failures
        )
        self.report.embedding_failures = list(
            self.chunking_source.report.embedding_failures
        )
        self.report.processing_errors = list(
            self.chunking_source.report.processing_errors
        )
        return self.report

    def close(self) -> None:
        """Cleanup resources."""
        if self.config.incremental.enabled:
            self._save_state()
        # Safety net: release the lock if the work generator was never fully consumed.
        if self.lock is not None:
            self.lock.release()
        super().close()
