import dataclasses
import logging
import os
import re
import threading
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from functools import lru_cache
from json import JSONDecodeError
from typing import (
    Any,
    Dict,
    Iterable,
    Iterator,
    List,
    Optional,
    Set,
    Tuple,
    Union,
)

import dateutil.parser as dp
import psutil
import pydantic
import requests
import tenacity
import yaml
from liquid import Template, Undefined
from pydantic import Field, field_validator
from requests.adapters import HTTPAdapter, Retry
from requests.exceptions import ConnectionError
from requests.models import HTTPBasicAuth, HTTPError
from tenacity import retry_if_exception_type, stop_after_attempt, wait_exponential

import datahub.emitter.mce_builder as builder
from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
from datahub.configuration.source_common import (
    DatasetLineageProviderConfigBase,
)
from datahub.configuration.validate_field_removal import pydantic_removed_field
from datahub.emitter.mcp import MetadataChangeProposalWrapper
from datahub.emitter.mcp_builder import (
    ContainerKey,
    add_dataset_to_container,
    gen_containers,
)
from datahub.emitter.request_helper import make_curl_command
from datahub.ingestion.api.common import PipelineContext
from datahub.ingestion.api.decorators import (
    SourceCapability,
    SupportStatus,
    capability,
    config_class,
    platform_name,
    support_status,
)
from datahub.ingestion.api.source import MetadataWorkUnitProcessor, SourceReport
from datahub.ingestion.api.workunit import MetadataWorkUnit
from datahub.ingestion.source.common.subtypes import (
    BIAssetSubTypes,
    BIContainerSubTypes,
)
from datahub.ingestion.source.state.stale_entity_removal_handler import (
    StaleEntityRemovalHandler,
    StaleEntityRemovalSourceReport,
    StatefulStaleMetadataRemovalConfig,
)
from datahub.ingestion.source.state.stateful_ingestion_base import (
    StatefulIngestionConfigBase,
    StatefulIngestionSourceBase,
)
from datahub.metadata.com.linkedin.pegasus2avro.common import (
    AuditStamp,
    ChangeAuditStamps,
)
from datahub.metadata.com.linkedin.pegasus2avro.metadata.snapshot import (
    ChartSnapshot,
    DashboardSnapshot,
)
from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
from datahub.metadata.schema_classes import (
    BrowsePathEntryClass,
    BrowsePathsClass,
    BrowsePathsV2Class,
    ChartInfoClass,
    ChartQueryClass,
    ChartQueryTypeClass,
    ChartTypeClass,
    DashboardInfoClass,
    DashboardUsageStatisticsClass,
    DatasetLineageTypeClass,
    DatasetPropertiesClass,
    EmbedClass,
    FineGrainedLineageClass,
    FineGrainedLineageDownstreamTypeClass,
    FineGrainedLineageUpstreamTypeClass,
    GlobalTagsClass,
    InputFieldClass,
    InputFieldsClass,
    OperationClass,
    OperationTypeClass,
    OtherSchemaClass,
    OwnerClass,
    OwnershipClass,
    OwnershipTypeClass,
    QueryLanguageClass,
    QueryPropertiesClass,
    QuerySourceClass,
    QueryStatementClass,
    SchemaFieldClass,
    SchemaMetadataClass,
    SubTypesClass,
    TagAssociationClass,
    TagPropertiesClass,
    UpstreamClass,
    UpstreamLineageClass,
    ViewPropertiesClass,
)
from datahub.metadata.urns import QueryUrn
from datahub.sql_parsing.sqlglot_lineage import (
    ColumnLineageInfo,
    SqlParsingResult,
    create_and_cache_schema_resolver,
    infer_output_schema,
    sqlglot_lineage,
)
from datahub.sql_parsing.sqlglot_utils import (
    parse_statements_and_pick,
)
from datahub.utilities import config_clean
from datahub.utilities.lossy_collections import LossyList
from datahub.utilities.perf_timer import PerfTimer
from datahub.utilities.ratelimiter import RateLimiter
from datahub.utilities.serialized_lru_cache import serialized_lru_cache
from datahub.utilities.threaded_iterator_executor import ThreadedIteratorExecutor

logger: logging.Logger = logging.getLogger(__name__)
# Default API limit for items returned per API call
# Used for the default per_page value for paginated API requests
DEFAULT_API_ITEMS_PER_PAGE = 30
MAX_API_ITEMS_PER_PAGE = 1000

# Override Undefined.__str__ so that unresolved Liquid template variables
# render as "NULL" instead of raising. Done at module level (once) rather
# than per-call to avoid redundant global mutation from worker threads.
Undefined.__str__ = lambda self: "NULL"  # type: ignore


class SpaceKey(ContainerKey):
    # Note that Mode has renamed Spaces to Collections.
    space_token: str


class ModeAPIConfig(ConfigModel):
    retry_backoff_multiplier: Union[int, float] = Field(
        default=2,
        ge=0,
        description="Multiplier for exponential backoff when waiting to retry",
    )
    max_retry_interval: Union[int, float] = Field(
        default=60, ge=0, description="Maximum interval to wait when retrying"
    )
    max_attempts: int = Field(
        default=10,
        ge=1,
        description="Maximum number of attempts to retry before failing",
    )
    timeout: int = Field(
        default=40,
        ge=1,
        description="Timeout setting, how long to wait for the Mode rest api to send data before giving up",
    )
    # Mode's API rate limit is ~40 requests per 10 seconds (4 req/s, 240 req/min).
    # See https://mode.com/help/articles/api-reference/#rate-limiting
    requests_per_minute: int = Field(
        default=180,
        ge=1,
        description="Maximum API requests per minute across all threads. "
        "Mode's API limit is ~240 req/min (4 req/s). Default of 180 "
        "leaves headroom to avoid 429 errors.",
    )


class ModeConfig(
    StatefulIngestionConfigBase,
    DatasetLineageProviderConfigBase,
):
    # See https://mode.com/developer/api-reference/authentication/
    # for authentication
    connect_uri: str = Field(
        default="https://app.mode.com", description="Mode host URL."
    )
    token: str = Field(
        description="When creating workspace API key this is the 'Key ID'."
    )
    password: pydantic.SecretStr = Field(
        description="When creating workspace API key this is the 'Secret'."
    )
    exclude_restricted: bool = Field(
        default=False, description="Exclude restricted collections"
    )

    exclude_personal_collections: bool = Field(
        default=True,
        description="Exclude personal collections from ingestion using Mode's "
        "server-side filter (?filter=custom). When True, only shared/custom "
        "collections are fetched from the API. When False, all collections are "
        "fetched (space_pattern still applies for client-side filtering).",
    )

    workspace: str = Field(
        description="The Mode workspace username. If you navigate to Workspace Settings > Details, "
        "the url will be `https://app.mode.com/organizations/<workspace-username>`. "
        # The lowercase comment is derived from a comment in a Mode API example.
        # https://mode.com/developer/api-cookbook/management/get-all-reports/
        # > "Note: workspace_name value should be all lowercase"
        "This is distinct from the workspace's display name, and should be all lowercase."
    )
    _default_schema = pydantic_removed_field("default_schema")

    space_pattern: AllowDenyPattern = Field(
        default=AllowDenyPattern(
            deny=["^Personal$"],
        ),
        description="Regex patterns for mode spaces to filter in ingestion (Spaces named as 'Personal' are filtered by default.) Specify regex to only match the space name. e.g. to only ingest space named analytics, use the regex 'analytics'",
    )

    owner_username_instead_of_email: bool = Field(
        default=True, description="Use username for owner URN instead of Email"
    )
    api_options: ModeAPIConfig = Field(
        default=ModeAPIConfig(),
        description='Retry/Wait settings for Mode API to avoid "Too many Requests" error. See Mode API Options below',
    )

    ingest_embed_url: bool = Field(
        default=True, description="Whether to Ingest embed URL for Reports"
    )

    stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None

    tag_measures_and_dimensions: bool = Field(
        default=True, description="Tag measures and dimensions in the schema"
    )

    items_per_page: HiddenFromDocs[int] = Field(
        DEFAULT_API_ITEMS_PER_PAGE,
        description="Number of items per page for paginated API requests.",
    )

    exclude_archived: bool = Field(
        default=False, description="Exclude archived reports"
    )

    max_threads: int = Field(
        default=1,
        ge=1,
        le=50,
        description="Maximum number of threads to use for parallel API requests. "
        "Increase to speed up ingestion for large workspaces. "
        "Setting too high may trigger Mode API rate limiting (429 errors).",
    )

    @field_validator("connect_uri", mode="after")
    @classmethod
    def remove_trailing_slash(cls, v):
        return config_clean.remove_trailing_slashes(v)

    @field_validator("items_per_page", mode="after")
    @classmethod
    def validate_items_per_page(cls, v):
        if 1 <= v <= MAX_API_ITEMS_PER_PAGE:
            return v
        else:
            raise ValueError(
                f"items_per_page must be between 1 and {MAX_API_ITEMS_PER_PAGE}"
            )


class HTTPError429(HTTPError):
    pass


class HTTPError504(HTTPError):
    pass


ModeRequestError = (HTTPError, JSONDecodeError)


def _is_http_404(error: Exception) -> bool:
    """Check if an exception is an HTTP 404 error with a valid response object."""
    return (
        isinstance(error, HTTPError)
        and getattr(error, "response", None) is not None
        and error.response.status_code == 404
    )


@dataclass
class ModeSourceReport(StaleEntityRemovalSourceReport):
    filtered_spaces: LossyList[str] = dataclasses.field(default_factory=LossyList)
    num_sql_parsed: int = 0
    num_sql_parser_failures: int = 0
    num_sql_parser_success: int = 0
    num_sql_parser_table_error: int = 0
    num_sql_parser_column_error: int = 0
    num_query_template_render: int = 0
    num_query_template_render_failures: int = 0
    num_query_template_render_success: int = 0
    num_requests_exceeding_rate_limit: int = 0
    num_requests_retried_on_timeout: int = 0
    num_spaces_retrieved: int = 0
    space_get_api_called: int = 0
    report_get_api_called: int = 0
    dataset_get_api_called: int = 0
    query_get_api_called: int = 0
    chart_get_api_called: int = 0
    process_memory_used_mb: float = 0
    # Thread-safe cumulative timing accumulators (seconds).
    # Protected by _lock. These track time spent in worker threads.
    query_api_total_sec: float = 0
    chart_api_total_sec: float = 0
    sql_parsing_total_sec: float = 0
    # Entity counts (protected by _lock for thread-safety)
    num_reports_processed: int = 0
    num_queries_processed: int = 0
    num_charts_processed: int = 0
    chart_api_calls_skipped: int = 0
    # PerfTimer is NOT thread-safe. These timers must only be used from the
    # main thread (e.g., in _get_space_name_and_tokens, _get_reports,
    # _get_datasets), never from threaded _process_report workers.
    space_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
    report_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
    dataset_get_timer: PerfTimer = dataclasses.field(default_factory=PerfTimer)
    # Lock protects structured log methods (report_warning, report_failure, etc.)
    # which mutate LossyList/LossyDict with multi-step non-atomic operations,
    # and counter increments from worker threads.
    _lock: threading.Lock = dataclasses.field(default_factory=threading.Lock)

    def report_dropped_space(self, ent_name: str) -> None:
        with self._lock:
            self.filtered_spaces.append(ent_name)

    def report_warning(self, *args: Any, **kwargs: Any) -> None:
        with self._lock:
            super().report_warning(*args, **kwargs)

    def report_failure(self, *args: Any, **kwargs: Any) -> None:
        with self._lock:
            super().report_failure(*args, **kwargs)

    def info(self, *args: Any, **kwargs: Any) -> None:
        with self._lock:
            super().info(*args, **kwargs)


@platform_name("Mode")
@config_class(ModeConfig)
@support_status(SupportStatus.CERTIFIED)
@capability(SourceCapability.CONTAINERS, "Enabled by default")
@capability(SourceCapability.DESCRIPTIONS, "Enabled by default")
@capability(SourceCapability.PLATFORM_INSTANCE, "Enabled by default")
@capability(SourceCapability.LINEAGE_COARSE, "Supported by default")
@capability(SourceCapability.LINEAGE_FINE, "Supported by default")
@capability(SourceCapability.OWNERSHIP, "Enabled by default")
class ModeSource(StatefulIngestionSourceBase):
    """
    Source that extracts reports and charts from Mode Analytics via REST API.

    Implementation notes:
    - Uses Mode API v3 for workspace, reports, queries, and charts
    - Parses SQL from query definitions to extract lineage
    - Only tested with PostgreSQL; other databases may work but are not validated
    - Chart properties extracted vary by visualization type (chart, table, pivot table)
    """

    ctx: PipelineContext
    config: ModeConfig
    report: ModeSourceReport
    platform = "mode"

    DIMENSION_TAG_URN = "urn:li:tag:Dimension"
    MEASURE_TAG_URN = "urn:li:tag:Measure"

    tag_definitions: Dict[str, TagPropertiesClass] = {
        DIMENSION_TAG_URN: TagPropertiesClass(
            name=DIMENSION_TAG_URN.split(":")[-1],
            description="A tag that is applied to all dimension fields.",
        ),
        MEASURE_TAG_URN: TagPropertiesClass(
            name=MEASURE_TAG_URN.split(":")[-1],
            description="A tag that is applied to all measures (metrics). Measures are typically the columns that you aggregate on",
        ),
    }

    def __hash__(self):
        return id(self)

    def __init__(self, ctx: PipelineContext, config: ModeConfig):
        super().__init__(config, ctx)
        self.config = config
        self.report = ModeSourceReport()
        self.ctx = ctx
        self._data_sources_by_id_cache: Optional[Dict[int, dict]] = None
        self._definitions_map_cache: Optional[Dict[str, str]] = None
        self.rate_limiter = RateLimiter(
            max_calls=self.config.api_options.requests_per_minute, period=60
        )

        self.session = requests.Session()
        # Handling retry and backoff
        retries = 3
        backoff_factor = 10
        retry = Retry(total=retries, backoff_factor=backoff_factor)
        pool_size = self.config.max_threads + 10
        adapter = HTTPAdapter(
            max_retries=retry,
            pool_connections=pool_size,
            pool_maxsize=pool_size,
        )
        self.session.mount("http://", adapter)
        self.session.mount("https://", adapter)

        self.session.auth = HTTPBasicAuth(
            self.config.token,
            self.config.password.get_secret_value(),
        )
        self.session.headers.update(
            {
                "Content-Type": "application/json",
                "Accept": "application/hal+json",
            }
        )

        self.workspace_uri = f"{self.config.connect_uri}/api/{self.config.workspace}"

        # Test the connection
        try:
            key_info = self._get_request_json(f"{self.config.connect_uri}/api/verify")
            logger.debug(f"Auth info: {key_info}")
        except ModeRequestError as e:
            self.report.report_failure(
                title="Failed to Connect",
                message="Unable to verify connection to mode.",
                context=f"Error: {str(e)}",
            )
            self.space_tokens: Dict[str, str] = {}
            return

        self.space_tokens = self._get_space_name_and_tokens()

    def _browse_path_space(self) -> List[BrowsePathEntryClass]:
        # TODO: Use containers for the workspace?
        return [
            BrowsePathEntryClass(id=self.config.workspace),
        ]

    def _browse_path_dashboard(self, space_token: str) -> List[BrowsePathEntryClass]:
        space_container_urn = self.gen_space_key(space_token).as_urn()
        return [
            *self._browse_path_space(),
            BrowsePathEntryClass(id=space_container_urn, urn=space_container_urn),
        ]

    def _browse_path_query(
        self, space_token: str, report_info: dict
    ) -> List[BrowsePathEntryClass]:
        dashboard_urn = self._dashboard_urn(report_info)
        return [
            *self._browse_path_dashboard(space_token),
            BrowsePathEntryClass(id=dashboard_urn, urn=dashboard_urn),
        ]

    def _browse_path_chart(
        self, space_token: str, report_info: dict, query_info: dict
    ) -> List[BrowsePathEntryClass]:
        query_urn = self.get_dataset_urn_from_query(query_info)
        return [
            *self._browse_path_query(space_token, report_info),
            BrowsePathEntryClass(id=query_urn, urn=query_urn),
        ]

    def _dashboard_urn(self, report_info: dict) -> str:
        return builder.make_dashboard_urn(self.platform, str(report_info.get("id", "")))

    @staticmethod
    def _parse_timestamp_ms(ts_str: Optional[str]) -> int:
        if ts_str:
            return int(dp.parse(ts_str).timestamp() * 1000)
        return int(datetime.now(tz=timezone.utc).timestamp() * 1000)

    def _parse_last_run_at(self, report_info: dict) -> Optional[int]:
        # Mode queries are refreshed, and that timestamp is reflected correctly here.
        # However, datasets are synced, and that's captured by the sync timestamps.
        # However, this is probably accurate enough for now.
        last_refreshed_ts = None
        last_refreshed_ts_str = report_info.get("last_run_at")
        if last_refreshed_ts_str:
            last_refreshed_ts = int(dp.parse(last_refreshed_ts_str).timestamp() * 1000)

        return last_refreshed_ts

    def construct_dashboard(
        self, space_token: str, report_info: dict, chart_urns: List[str]
    ) -> Optional[Tuple[DashboardSnapshot, MetadataChangeProposalWrapper]]:
        report_token = report_info.get("token", "")
        # logger.debug(f"Processing report {report_info.get('name', '')}: {report_info}")

        if not report_token:
            self.report.report_warning(
                title="Missing Report Token",
                message=f"Report token is missing for {report_info.get('id', '')}",
            )
            return None

        if not report_info.get("id"):
            self.report.report_warning(
                title="Missing Report ID",
                message=f"Report id is missing for {report_info.get('token', '')}",
            )
            return None

        dashboard_urn = self._dashboard_urn(report_info)
        dashboard_snapshot = DashboardSnapshot(
            urn=dashboard_urn,
            aspects=[],
        )

        title = report_info.get("name", "")
        description = report_info.get("description", "")
        last_modified = ChangeAuditStamps()

        # Creator + created ts.
        creator = self._get_creator(
            report_info.get("_links", {}).get("creator", {}).get("href", "")
        )
        if creator:
            creator_actor = builder.make_user_urn(creator)
            created_ts = self._parse_timestamp_ms(report_info.get("created_at"))
            last_modified.created = AuditStamp(time=created_ts, actor=creator_actor)

        # Last modified ts.
        last_modified_ts_str = report_info.get("last_saved_at")
        if not last_modified_ts_str:
            # Sometimes mode returns null for last_saved_at.
            # In that case, we use the edited_at timestamp instead.
            last_modified_ts_str = report_info.get("edited_at")
        if last_modified_ts_str:
            modified_ts = int(dp.parse(last_modified_ts_str).timestamp() * 1000)
            last_modified.lastModified = AuditStamp(
                time=modified_ts, actor="urn:li:corpuser:unknown"
            )

        # Last refreshed ts.
        last_refreshed_ts = self._parse_last_run_at(report_info)

        # Datasets
        datasets = []
        for imported_dataset_name in report_info.get("imported_datasets", []):
            try:
                mode_dataset = self._get_request_json(
                    f"{self.workspace_uri}/reports/{imported_dataset_name.get('token')}"
                )
            except HTTPError as http_error:
                if _is_http_404(http_error):
                    self.report.report_warning(
                        title="Report Not Found",
                        message="Referenced report for reusable dataset was not found.",
                        context=f"Report: {report_info.get('id')}, "
                        f"Imported Dataset Report: {imported_dataset_name.get('token')}",
                    )
                    continue
                else:
                    raise http_error

            dataset_urn = builder.make_dataset_urn_with_platform_instance(
                self.platform,
                str(mode_dataset.get("id")),
                platform_instance=None,
                env=self.config.env,
            )
            datasets.append(dataset_urn)

        dashboard_info_class = DashboardInfoClass(
            description=description if description else "",
            title=title if title else "",
            charts=chart_urns,
            datasets=datasets if datasets else None,
            lastModified=last_modified,
            lastRefreshed=last_refreshed_ts,
            dashboardUrl=f"{self.config.connect_uri}/{self.config.workspace}/reports/{report_token}",
            customProperties={},
        )
        dashboard_snapshot.aspects.append(dashboard_info_class)

        # browse path
        space_name = self.space_tokens[space_token]
        browse_path = BrowsePathsClass(
            paths=[
                f"/mode/{self.config.workspace}/"
                f"{space_name}/"
                f"{title if title else report_info.get('id', '')}"
            ]
        )
        dashboard_snapshot.aspects.append(browse_path)

        browse_path_v2 = BrowsePathsV2Class(
            path=self._browse_path_dashboard(space_token)
        )
        browse_mcp = MetadataChangeProposalWrapper(
            entityUrn=dashboard_urn,
            aspect=browse_path_v2,
        )

        # Ownership
        ownership = self._get_ownership(
            self._get_creator(
                report_info.get("_links", {}).get("creator", {}).get("href", "")
            )
        )
        if ownership is not None:
            dashboard_snapshot.aspects.append(ownership)

        return dashboard_snapshot, browse_mcp

    @lru_cache(maxsize=5000)
    def _get_ownership(self, user: str) -> Optional[OwnershipClass]:
        if user is not None:
            owner_urn = builder.make_user_urn(user)
            ownership: OwnershipClass = OwnershipClass(
                owners=[
                    OwnerClass(
                        owner=owner_urn,
                        type=OwnershipTypeClass.DATAOWNER,
                    )
                ]
            )
            return ownership

        return None

    def _get_creator(self, href: str) -> Optional[str]:
        if not href:
            return None
        try:
            return self._fetch_creator(href)
        except ModeRequestError as e:
            self.report.report_warning(
                title="Failed to retrieve Mode creator",
                message=f"Unable to retrieve user for {href}",
                context=f"Reason: {str(e)}",
            )
            return None

    @serialized_lru_cache(maxsize=5000)
    def _fetch_creator(self, href: str) -> Optional[str]:
        """Fetch creator from Mode API.

        Raises on API failure so the result is not cached and can be retried.
        Returns None (cached) when the user simply has no username/email.
        """
        user_json = self._get_request_json(f"{self.config.connect_uri}{href}")
        return (
            user_json.get("username")
            if self.config.owner_username_instead_of_email
            else user_json.get("email")
        )

    def _get_space_name_and_tokens(self) -> dict:
        space_info = {}
        try:
            logger.debug(f"Retrieving spaces for {self.workspace_uri}")
            with self.report.space_get_timer:
                space_filter = (
                    "custom" if self.config.exclude_personal_collections else "all"
                )
                for spaces_page in self._get_paged_request_json(
                    f"{self.workspace_uri}/spaces?filter={space_filter}",
                    "spaces",
                    self.config.items_per_page,
                ):
                    self.report.space_get_api_called += 1
                    logger.debug(
                        f"Read {len(spaces_page)} spaces records from workspace {self.workspace_uri}"
                    )
                    self.report.num_spaces_retrieved += len(spaces_page)
                    for s in spaces_page:
                        logger.debug(f"Space: {s.get('name')}")
                        space_name = s.get("name", "")
                        # Using both restricted and default_access_level because
                        # there is a current bug with restricted returning False everytime
                        # which has been reported to Mode team
                        if self.config.exclude_restricted and (
                            s.get("restricted")
                            or s.get("default_access_level") == "restricted"
                        ):
                            logger.debug(
                                f"Skipping space {space_name} due to exclude restricted"
                            )
                            continue
                        if not self.config.space_pattern.allowed(space_name):
                            self.report.report_dropped_space(space_name)
                            logger.debug(
                                f"Skipping space {space_name} due to space pattern"
                            )
                            continue
                        space_info[s.get("token", "")] = s.get("name", "")
        except ModeRequestError as e:
            self.report.report_failure(
                title="Failed to Retrieve Spaces",
                message="Unable to retrieve spaces / collections for workspace.",
                context=f"Workspace: {self.workspace_uri}, Error: {str(e)}",
            )

        return space_info

    def _get_chart_type(self, token: str, display_type: str) -> Optional[str]:
        type_mapping = {
            "table": ChartTypeClass.TABLE,
            "bar": ChartTypeClass.BAR,
            "bigNumber": ChartTypeClass.TEXT,
            "line": ChartTypeClass.LINE,
            "stackedBar100": ChartTypeClass.BAR,
            "stackedBar": ChartTypeClass.BAR,
            "hStackedBar": ChartTypeClass.BAR,
            "hStackedBar100": ChartTypeClass.BAR,
            "hBar": ChartTypeClass.BAR,
            "area": ChartTypeClass.AREA,
            "totalArea": ChartTypeClass.AREA,
            "pie": ChartTypeClass.PIE,
            "donut": ChartTypeClass.PIE,
            "scatter": ChartTypeClass.SCATTER,
            "bigValue": ChartTypeClass.TEXT,
            "pivotTable": ChartTypeClass.TABLE,
            "linePlusBar": None,
            "vegas": None,
            "vegasPivotTable": ChartTypeClass.TABLE,
            "histogram": ChartTypeClass.HISTOGRAM,
        }
        if not display_type:
            self.report.info(
                title="Missing chart type found",
                message="Chart type is missing. Setting to None",
                context=f"Token: {token}",
            )
            return None
        try:
            chart_type = type_mapping[display_type]
        except KeyError:
            self.report.info(
                title="Unrecognized chart type found",
                message=f"Chart type {display_type} not supported. Setting to None",
                context=f"Token: {token}",
            )
            chart_type = None

        return chart_type

    def construct_chart_custom_properties(
        self, chart_detail: dict, chart_type: str
    ) -> Dict:
        custom_properties = {
            "ChartType": chart_type,
        }
        metadata = chart_detail.get("encoding", {})
        if chart_type == "table":
            columns = list(chart_detail.get("fieldFormats", {}).keys())
            str_columns = ",".join([c[1:-1] for c in columns])
            filters = metadata.get("filter", [])
            filters = filters[0].get("formula", "") if len(filters) else ""

            custom_properties.update(
                {
                    "Columns": str_columns,
                    "Filters": filters[1:-1] if len(filters) else "",
                }
            )

        elif chart_type == "pivotTable":
            pivot_table = chart_detail.get("pivotTable", {})
            columns = pivot_table.get("columns", [])
            rows = pivot_table.get("rows", [])
            values = pivot_table.get("values", [])
            filters = pivot_table.get("filters", [])

            custom_properties.update(
                {
                    "Columns": ", ".join(columns) if len(columns) else "",
                    "Rows": ", ".join(rows) if len(rows) else "",
                    "Metrics": ", ".join(values) if len(values) else "",
                    "Filters": ", ".join(filters) if len(filters) else "",
                }
            )
            # list filters in their own row
            for filter in filters:
                custom_properties[f"Filter: {filter}"] = ", ".join(
                    pivot_table.get("filterValues", {}).get(filter, "")
                )
        # Chart
        else:
            x = metadata.get("x", [])
            x2 = metadata.get("x2", [])
            y = metadata.get("y", [])
            y2 = metadata.get("y2", [])
            value = metadata.get("value", [])
            filters = metadata.get("filter", [])

            custom_properties.update(
                {
                    "X": x[0].get("formula", "") if len(x) else "",
                    "Y": y[0].get("formula", "") if len(y) else "",
                    "X2": x2[0].get("formula", "") if len(x2) else "",
                    "Y2": y2[0].get("formula", "") if len(y2) else "",
                    "Metrics": value[0].get("formula", "") if len(value) else "",
                    "Filters": filters[0].get("formula", "") if len(filters) else "",
                }
            )

        return custom_properties

    def _get_datahub_friendly_platform(self, adapter, platform):
        # Map adaptor names to what datahub expects in
        # https://github.com/datahub-project/datahub/blob/master/metadata-service/configuration/src/main/resources/bootstrap_mcps/data-platforms.yaml

        platform_mapping = {
            "jdbc:athena": "athena",
            "jdbc:bigquery": "bigquery",
            "jdbc:druid": "druid",
            "jdbc:hive": "hive",
            "jdbc:mysql": "mysql",
            "jdbc:oracle": "oracle",
            "jdbc:postgresql": "postgres",
            "jdbc:presto": "presto",
            "jdbc:redshift": "redshift",
            "jdbc:snowflake": "snowflake",
            "jdbc:spark": "spark",
            "jdbc:trino": "trino",
            "jdbc:sqlserver": "mssql",
            "jdbc:teradata": "teradata",
        }
        if adapter in platform_mapping:
            return platform_mapping[adapter]
        else:
            self.report.report_warning(
                title="Unrecognized Platform Found",
                message=f"Platform was not found in DataHub. "
                f"Using {platform} name as is",
            )

        return platform

    def _get_data_sources_by_id(self) -> Dict[int, dict]:
        """Fetch data sources and index by ID for O(1) lookup.

        Uses a manual cache that only stores successful results so transient
        API failures are retried on the next call instead of being permanently
        cached as empty.
        """
        if self._data_sources_by_id_cache is not None:
            return self._data_sources_by_id_cache
        try:
            ds_json = self._get_request_json(f"{self.workspace_uri}/data_sources")
            data_sources = ds_json.get("_embedded", {}).get("data_sources", [])
            self._data_sources_by_id_cache = {
                int(ds.get("id", -1)): ds for ds in data_sources
            }
            fetched_ids = sorted(self._data_sources_by_id_cache.keys())
            logger.info(
                f"Fetched {len(fetched_ids)} data sources from Mode. IDs: {fetched_ids}"
            )
            return self._data_sources_by_id_cache
        except ModeRequestError as e:
            self.report.report_failure(
                title="Failed to retrieve Data Sources",
                message="Unable to retrieve data sources from Mode.",
                context=f"Error: {str(e)}",
            )
            return {}

    def _get_platform_and_dbname(
        self, data_source_id: int
    ) -> Union[Tuple[str, str], Tuple[None, None]]:
        data_sources_by_id = self._get_data_sources_by_id()

        if not data_sources_by_id:
            self.report.report_failure(
                title="No Data Sources Found",
                message="Could not find data sources matching some ids",
                context=f"Data Source ID: {data_source_id}",
            )
            return None, None

        data_source = data_sources_by_id.get(data_source_id)
        if data_source is None:
            available_ids = sorted(data_sources_by_id.keys())
            self.report.report_warning(
                title="Unable to construct upstream lineage",
                message=f"Data source ID {data_source_id} not found among "
                f"the {len(available_ids)} data sources in the workspace "
                f"(available IDs: {available_ids}). This typically means the "
                f"database connection was deleted from Mode but queries still "
                f"reference it. Lineage will be skipped for affected queries.",
                context=f"Data Source ID: {data_source_id}",
            )
            return None, None

        platform = self._get_datahub_friendly_platform(
            data_source.get("adapter", ""), data_source.get("name", "")
        )
        database = data_source.get("database", "")
        # On bigquery, change the database from "default" to the host (project_id)
        # For lineage we need project_id.db.table
        if platform == "bigquery" and database == "default":
            database = data_source.get("host", "")
        return platform, database

    def _replace_definitions(
        self,
        raw_query: str,
        _depth: int = 0,
        _seen: Optional[Set[str]] = None,
    ) -> str:
        if _depth > 10:
            logger.warning(
                "Definition replacement exceeded max depth of 10, returning query as-is"
            )
            return raw_query

        if _seen is None:
            _seen = set()

        query = raw_query
        definitions = re.findall(r"({{(?:\s+)?@[^}{]+}})", raw_query)
        if not definitions:
            return query

        # Expand each definition's body independently. Only the definition's
        # own name is added to _seen for its recursive expansion, so sibling
        # definitions at the same level don't interfere with each other.
        for definition_variable in definitions:
            definition_name, definition_alias = self._parse_definition_name(
                definition_variable
            )

            if definition_name in _seen:
                logger.warning(
                    f"Circular definition detected for '{definition_name}', skipping"
                )
                fallback = (
                    f"{definition_name} as {definition_alias}"
                    if definition_alias
                    else definition_name
                )
                query = query.replace(definition_variable, fallback)
                continue

            definition_query = self._get_definition(definition_name)
            if definition_query is not None:
                # Recursively expand nested definitions within this body,
                # tracking only this definition's ancestry chain.
                expanded_body = self._replace_definitions(
                    definition_query, _depth + 1, _seen | {definition_name}
                )
                # Newline before the closing paren ensures any trailing
                # SQL line comment (--) in the body doesn't swallow it.
                replacement = (
                    f"({expanded_body}\n) as {definition_alias}"
                    if definition_alias
                    else f"({expanded_body}\n)"
                )
                query = query.replace(definition_variable, replacement)
            else:
                fallback = (
                    f"{definition_name} as {definition_alias}"
                    if definition_alias
                    else definition_name
                )
                query = query.replace(definition_variable, fallback)

        query = query.replace("\\n", "\n")
        query = query.replace("\\t", "\t")
        return query

    def _parse_definition_name(self, definition_variable: str) -> Tuple[str, str]:
        name, alias = "", ""
        # i.e '{{ @join_on_definition as alias}}'
        name_match = re.findall(r"@\w+", definition_variable)
        if len(name_match):
            name = name_match[0][1:]
        alias_match = re.findall(
            r"as\s+\w+", definition_variable
        )  # i.e ['as    alias_name']
        if len(alias_match):
            alias_match = alias_match[0].split(" ")
            alias = alias_match[-1]

        return name, alias

    def _get_definitions_map(self) -> Dict[str, str]:
        """Fetch all definitions and return a {name: source} mapping.

        Uses a manual cache that only stores successful results so transient
        API failures are retried on the next call.
        """
        if self._definitions_map_cache is not None:
            return self._definitions_map_cache
        try:
            definition_json = self._get_request_json(
                f"{self.workspace_uri}/definitions"
            )
            definitions = definition_json.get("_embedded", {}).get("definitions", [])
            self._definitions_map_cache = {
                d.get("name", ""): d.get("source", "") for d in definitions
            }
            return self._definitions_map_cache
        except ModeRequestError as e:
            self.report.report_failure(
                title="Failed to Retrieve Definitions",
                message="Unable to retrieve definitions from Mode.",
                context=f"Error: {str(e)}",
            )
            return {}

    def _get_definition(self, definition_name: str) -> Optional[str]:
        return self._get_definitions_map().get(definition_name)

    def _get_datasource_urn(
        self,
        platform: str,
        platform_instance: Optional[str],
        database: str,
        source_tables: List[str],
    ) -> List[str]:
        if not platform and not database:
            return []
        return [
            builder.make_dataset_urn_with_platform_instance(
                platform,
                f"{database}.{s_table}",
                platform_instance=platform_instance,
                env=self.config.env,
            )
            for s_table in source_tables
        ]

    def get_custom_props_from_dict(self, obj: dict, keys: List[str]) -> Optional[dict]:
        return {key: str(obj[key]) for key in keys if obj.get(key)} or None

    def get_dataset_urn_from_query(self, query_data: dict) -> str:
        return builder.make_dataset_urn_with_platform_instance(
            platform=self.platform,
            name=str(query_data.get("id")),
            platform_instance=None,
            env=self.config.env,
        )

    def get_query_instance_urn_from_query(self, query_data: dict) -> str:
        id = query_data.get("id")
        last_run_id = query_data.get("last_run_id")
        data_source_id = query_data.get("data_source_id")
        return QueryUrn(f"{id}.{data_source_id}.{last_run_id}").urn()

    def set_field_tags(self, fields: List[SchemaFieldClass]) -> None:
        for field in fields:
            # It is not clear how to distinguish between measures and dimensions in Mode.
            # We are using a heuristic to tag fields that are not named `id` or `_id` or ends with '_number' and are not of type `NumberType` as dimensions.
            # This is a heuristic and may not be accurate.
            if field.type.type.RECORD_SCHEMA.name in ["NumberType"] and (
                not field.fieldPath.endswith("_number")
                and not re.findall(r"(^(id)[_\d]?)|([_\d?](id)$)", field.fieldPath)
            ):
                tag = TagAssociationClass(tag=self.MEASURE_TAG_URN)
            else:
                tag = TagAssociationClass(tag=self.DIMENSION_TAG_URN)
            field.globalTags = GlobalTagsClass(tags=[tag])

    def normalize_mode_query(self, query: str) -> str:
        normalized_query: str = query

        with self.report._lock:
            self.report.num_query_template_render += 1
        matches = re.findall(
            r"{% form %}(.*?){% endform %}",
            query,
            re.MULTILINE | re.DOTALL | re.IGNORECASE,
        )
        jinja_params: Dict = {}
        if matches:
            try:
                for match in matches:
                    definition = Template(source=match).render()
                    parameters = yaml.safe_load(definition)
                    for key in parameters:
                        jinja_params[key] = parameters[key].get("default", "")
            except Exception as e:
                logger.warning(
                    f"Parsing form parameters failed, rendering without params: {e}"
                )

            normalized_query = re.sub(
                pattern=r"{% form %}(.*){% endform %}",
                repl="",
                string=query,
                count=0,
                flags=re.MULTILINE | re.DOTALL,
            )

        try:
            rendered_query = Template(normalized_query).render(jinja_params)
            with self.report._lock:
                self.report.num_query_template_render_success += 1
        except Exception as e:
            logger.warning(
                f"Rendering Liquid template failed, lineage may be incomplete: {e}"
            )
            with self.report._lock:
                self.report.num_query_template_render_failures += 1
            # Return form-stripped query so SQL parsers don't choke on template syntax
            return normalized_query

        return rendered_query

    def construct_query_or_dataset(
        self,
        report_token: str,
        query_data: dict,
        space_token: str,
        report_info: dict,
        is_mode_dataset: bool,
    ) -> Iterable[MetadataWorkUnit]:
        query_urn = (
            self.get_dataset_urn_from_query(query_data)
            if not is_mode_dataset
            else self.get_dataset_urn_from_query(report_info)
        )

        query_token = query_data.get("token")

        externalUrl = (
            f"{self.config.connect_uri}/{self.config.workspace}/datasets/{report_token}"
            if is_mode_dataset
            else f"{self.config.connect_uri}/{self.config.workspace}/reports/{report_token}/details/queries/{query_token}"
        )

        dataset_props = DatasetPropertiesClass(
            name=report_info.get("name") if is_mode_dataset else query_data.get("name"),
            description=None,
            externalUrl=externalUrl,
            customProperties=self.get_custom_props_from_dict(
                query_data,
                [
                    "id",
                    "created_at",
                    "updated_at",
                    "last_run_id",
                    "data_source_id",
                    "explorations_count",
                    "report_imports_count",
                    "dbt_metric_id",
                ],
            ),
        )
        yield (
            MetadataChangeProposalWrapper(
                entityUrn=query_urn,
                aspect=dataset_props,
            ).as_workunit()
        )

        if raw_query := query_data.get("raw_query"):
            yield MetadataChangeProposalWrapper(
                entityUrn=query_urn,
                aspect=ViewPropertiesClass(
                    viewLogic=raw_query,
                    viewLanguage=QueryLanguageClass.SQL,
                    materialized=False,
                ),
            ).as_workunit()

        if is_mode_dataset:
            space_container_key = self.gen_space_key(space_token)
            yield from add_dataset_to_container(
                container_key=space_container_key,
                dataset_urn=query_urn,
            )

        subtypes = SubTypesClass(
            typeNames=(
                [
                    BIAssetSubTypes.MODE_DATASET
                    if is_mode_dataset
                    else BIAssetSubTypes.MODE_QUERY
                ]
            )
        )
        yield (
            MetadataChangeProposalWrapper(
                entityUrn=query_urn,
                aspect=subtypes,
            ).as_workunit()
        )

        yield MetadataChangeProposalWrapper(
            entityUrn=query_urn,
            aspect=BrowsePathsV2Class(
                path=self._browse_path_dashboard(space_token)
                if is_mode_dataset
                else self._browse_path_query(space_token, report_info)
            ),
        ).as_workunit()

        data_source_id = query_data.get("data_source_id")
        if data_source_id is None:
            logger.debug(
                f"No data_source_id for report {report_token} query {query_token}, skipping lineage"
            )
            return
        (
            upstream_warehouse_platform,
            upstream_warehouse_db_name,
        ) = self._get_platform_and_dbname(int(data_source_id))
        if upstream_warehouse_platform is None:
            return

        query = query_data.get("raw_query")
        if not query:
            logger.warning(
                f"No raw_query found for report {report_token} query {query_token}, skipping lineage"
            )
            return
        query = self._replace_definitions(query)
        normalized_query = self.normalize_mode_query(query)

        # Detect Snowflake USE WAREHOUSE via regex on the raw string (no parse needed).
        if upstream_warehouse_platform == "snowflake":
            regexp = r"use\s+warehouse\s+(.*)(\s+)?;"
            matches = re.search(
                regexp,
                normalized_query,
                re.MULTILINE | re.DOTALL | re.IGNORECASE,
            )
            if matches and matches.group(1):
                upstream_warehouse_db_name = matches.group(1)

        # Parse multi-statement SQL and pick the last meaningful statement.
        platform_instance = (
            self.config.platform_instance_map.get(upstream_warehouse_platform)
            if upstream_warehouse_platform and self.config.platform_instance_map
            else None
        )
        sql_parse_start = time.perf_counter()
        try:
            parsed_expression = parse_statements_and_pick(
                normalized_query, upstream_warehouse_platform
            )
            query_to_parse = parsed_expression.sql(dialect=upstream_warehouse_platform)
        except Exception as e:
            logger.debug(
                f"parse_statements_and_pick failed on: {normalized_query}, error: {e}"
            )
            query_to_parse = normalized_query

        try:
            schema_resolver = create_and_cache_schema_resolver(
                platform=upstream_warehouse_platform,
                env=self.config.env,
                graph=self.ctx.graph,
                platform_instance=platform_instance,
            )
            parsed_query_object = sqlglot_lineage(
                sql=query_to_parse,
                schema_resolver=schema_resolver,
                default_db=upstream_warehouse_db_name,
            )
        except Exception as e:
            parsed_query_object = SqlParsingResult.make_from_error(e)
        sql_parse_elapsed = time.perf_counter() - sql_parse_start
        with self.report._lock:
            self.report.sql_parsing_total_sec += sql_parse_elapsed

        with self.report._lock:
            self.report.num_sql_parsed += 1
            if parsed_query_object.debug_info.table_error:
                self.report.num_sql_parser_table_error += 1
                self.report.num_sql_parser_failures += 1
            elif parsed_query_object.debug_info.column_error:
                self.report.num_sql_parser_column_error += 1
                self.report.num_sql_parser_failures += 1
            else:
                self.report.num_sql_parser_success += 1
        if parsed_query_object.debug_info.table_error:
            logger.info(
                f"Failed to parse compiled code for report: {report_token} query: {query_token} {parsed_query_object.debug_info.error} the query was [{query_to_parse}]"
            )
        elif parsed_query_object.debug_info.column_error:
            logger.info(
                f"Failed to generate CLL for report: {report_token} query: {query_token}: {parsed_query_object.debug_info.column_error} the query was [{query_to_parse}]"
            )

        schema_fields = infer_output_schema(parsed_query_object)
        if schema_fields:
            schema_metadata = SchemaMetadataClass(
                schemaName="mode_dataset" if is_mode_dataset else "mode_query",
                platform=f"urn:li:dataPlatform:{self.platform}",
                version=0,
                fields=schema_fields,
                hash="",
                platformSchema=OtherSchemaClass(rawSchema=""),
            )
            if self.config.tag_measures_and_dimensions:
                self.set_field_tags(schema_fields)

            yield (
                MetadataChangeProposalWrapper(
                    entityUrn=query_urn,
                    aspect=schema_metadata,
                ).as_workunit()
            )

        yield from self.get_upstream_lineage_for_parsed_sql(
            query_urn, query_data, parsed_query_object
        )

        operation = OperationClass(
            operationType=OperationTypeClass.UPDATE,
            lastUpdatedTimestamp=self._parse_timestamp_ms(query_data.get("updated_at")),
            timestampMillis=int(datetime.now(tz=timezone.utc).timestamp() * 1000),
        )

        yield MetadataChangeProposalWrapper(
            entityUrn=query_urn,
            aspect=operation,
        ).as_workunit()

        creator = self._get_creator(
            query_data.get("_links", {}).get("creator", {}).get("href", "")
        )
        modified_actor = builder.make_user_urn(
            creator if creator is not None else "unknown"
        )

        created_ts = self._parse_timestamp_ms(query_data.get("created_at"))
        modified_ts = self._parse_timestamp_ms(query_data.get("updated_at"))

        query_instance_urn = self.get_query_instance_urn_from_query(query_data)
        value = query_data.get("raw_query")
        if value:
            query_properties = QueryPropertiesClass(
                statement=QueryStatementClass(
                    value=value,
                    language=QueryLanguageClass.SQL,
                ),
                source=QuerySourceClass.SYSTEM,
                created=AuditStamp(time=created_ts, actor=modified_actor),
                lastModified=AuditStamp(time=modified_ts, actor=modified_actor),
            )

            yield MetadataChangeProposalWrapper(
                entityUrn=query_instance_urn,
                aspect=query_properties,
            ).as_workunit()

    def get_upstream_lineage_for_parsed_sql(
        self, query_urn: str, query_data: dict, parsed_query_object: SqlParsingResult
    ) -> Iterable[MetadataWorkUnit]:
        if parsed_query_object is None:
            logger.info(f"Failed to extract lineage from datasource {query_urn}")
            return
        if parsed_query_object.debug_info.table_error:
            logger.info(
                f"Failed to extract table lineage from datasource {query_urn}: {parsed_query_object.debug_info.table_error}"
            )
            return

        fine_grained_lineages: List[FineGrainedLineageClass] = []

        if parsed_query_object.debug_info.column_error:
            logger.info(
                f"Failed to extract column level lineage from datasource {query_urn}: {parsed_query_object.debug_info.column_error}"
            )
        else:
            cll: List[ColumnLineageInfo] = (
                parsed_query_object.column_lineage
                if parsed_query_object.column_lineage is not None
                else []
            )

            for cll_info in cll:
                downstream = (
                    [
                        builder.make_schema_field_urn(
                            query_urn, cll_info.downstream.column
                        )
                    ]
                    if cll_info.downstream is not None
                    and cll_info.downstream.column is not None
                    else []
                )
                upstreams = [
                    builder.make_schema_field_urn(column_ref.table, column_ref.column)
                    for column_ref in cll_info.upstreams
                ]
                fine_grained_lineages.append(
                    FineGrainedLineageClass(
                        downstreamType=FineGrainedLineageDownstreamTypeClass.FIELD,
                        downstreams=downstream,
                        upstreamType=FineGrainedLineageUpstreamTypeClass.FIELD_SET,
                        upstreams=upstreams,
                    )
                )

        upstream_lineage = UpstreamLineageClass(
            upstreams=[
                UpstreamClass(
                    type=DatasetLineageTypeClass.TRANSFORMED,
                    dataset=input_table_urn,
                    query=self.get_query_instance_urn_from_query(query_data),
                )
                for input_table_urn in parsed_query_object.in_tables
            ],
            fineGrainedLineages=fine_grained_lineages or None,
        )

        yield MetadataChangeProposalWrapper(
            entityUrn=query_urn,
            aspect=upstream_lineage,
        ).as_workunit()

    def get_formula_columns(
        self, node: Dict, columns: Optional[Set[str]] = None
    ) -> Set[str]:
        columns = columns if columns is not None else set()
        if isinstance(node, dict):
            for key, item in node.items():
                if isinstance(item, dict):
                    self.get_formula_columns(item, columns)
                elif isinstance(item, list):
                    for i in item:
                        if isinstance(i, dict):
                            self.get_formula_columns(i, columns)
                elif isinstance(item, str):
                    if key == "formula":
                        column_names = re.findall(r"\[(.+?)\]", item)
                        columns.update(column_names)
        return columns

    def get_input_fields(
        self,
        chart_urn: str,
        chart_data: Dict,
        chart_fields: Dict[str, SchemaFieldClass],
        query_urn: str,
    ) -> Iterable[MetadataWorkUnit]:
        # TODO: Identify which fields are used as X, Y, filters, etc and tag them accordingly.
        fields = self.get_formula_columns(chart_data)

        input_fields = []

        for field in fields:
            if field.lower() not in chart_fields:
                continue
            input_field = InputFieldClass(
                schemaFieldUrn=builder.make_schema_field_urn(query_urn, field.lower()),
                schemaField=chart_fields[field.lower()],
            )
            input_fields.append(input_field)

        if not input_fields:
            return

        inputFields = InputFieldsClass(fields=input_fields)

        yield MetadataChangeProposalWrapper(
            entityUrn=chart_urn,
            aspect=inputFields,
        ).as_workunit()

    def construct_chart_from_api_data(
        self,
        index: int,
        chart_data: dict,
        chart_fields: Dict[str, SchemaFieldClass],
        query: dict,
        space_token: str,
        report_info: dict,
        query_name: str,
    ) -> Iterable[MetadataWorkUnit]:
        # logger.debug(f"Processing chart {chart_data.get('token', '')}: {chart_data}")
        chart_urn = builder.make_chart_urn(self.platform, chart_data.get("token", ""))
        chart_snapshot = ChartSnapshot(
            urn=chart_urn,
            aspects=[],
        )

        last_modified = ChangeAuditStamps()
        creator = self._get_creator(
            chart_data.get("_links", {}).get("creator", {}).get("href", "")
        )
        if creator is not None:
            modified_actor = builder.make_user_urn(creator)
            created_ts = self._parse_timestamp_ms(chart_data.get("created_at"))
            modified_ts = self._parse_timestamp_ms(chart_data.get("updated_at"))
            last_modified = ChangeAuditStamps(
                created=AuditStamp(time=created_ts, actor=modified_actor),
                lastModified=AuditStamp(time=modified_ts, actor=modified_actor),
            )

        # Last refreshed ts.
        last_refreshed_ts = self._parse_last_run_at(report_info)

        chart_detail = (
            chart_data.get("view", {})
            if len(chart_data.get("view", {})) != 0
            else chart_data.get("view_vegas", {})
        )

        mode_chart_type = chart_detail.get("chartType", "") or chart_detail.get(
            "selectedChart", ""
        )
        chart_type = self._get_chart_type(chart_data.get("token", ""), mode_chart_type)
        description = (
            chart_detail.get("description")
            or chart_detail.get("chartDescription")
            or ""
        )

        title = (
            chart_detail.get("title")
            or chart_detail.get("chartTitle")
            or f"Chart {index}"
        )

        # create datasource urn
        custom_properties = self.construct_chart_custom_properties(
            chart_detail, mode_chart_type
        )

        query_urn = self.get_dataset_urn_from_query(query)

        # Chart Info
        chart_info = ChartInfoClass(
            type=chart_type,
            description=description,
            title=title,
            lastModified=last_modified,
            lastRefreshed=last_refreshed_ts,
            # The links href starts with a slash already.
            chartUrl=f"{self.config.connect_uri}{chart_data.get('_links', {}).get('report_viz_web', {}).get('href', '')}",
            inputs=[query_urn],
            customProperties=custom_properties,
            inputEdges=[],
        )
        chart_snapshot.aspects.append(chart_info)

        query_urn = self.get_dataset_urn_from_query(query)
        yield from self.get_input_fields(chart_urn, chart_data, chart_fields, query_urn)

        yield MetadataChangeProposalWrapper(
            entityUrn=chart_urn,
            aspect=SubTypesClass(typeNames=[BIAssetSubTypes.MODE_CHART]),
        ).as_workunit()

        # Browse Path
        space_name = self.space_tokens[space_token]
        report_name = report_info.get("name", report_info.get("token", "unknown"))
        path = f"/mode/{self.config.workspace}/{space_name}/{report_name}/{query_name}/{title}"
        browse_path = BrowsePathsClass(paths=[path])
        chart_snapshot.aspects.append(browse_path)

        # Browse path v2
        browse_path_v2 = BrowsePathsV2Class(
            path=self._browse_path_chart(space_token, report_info, query),
        )
        yield MetadataChangeProposalWrapper(
            entityUrn=chart_urn,
            aspect=browse_path_v2,
        ).as_workunit()

        # Query
        chart_query = ChartQueryClass(
            rawQuery=query.get("raw_query", ""),
            type=ChartQueryTypeClass.SQL,
        )
        chart_snapshot.aspects.append(chart_query)

        # Ownership
        ownership = self._get_ownership(
            self._get_creator(
                chart_data.get("_links", {}).get("creator", {}).get("href", "")
            )
        )
        if ownership is not None:
            chart_snapshot.aspects.append(ownership)

        mce = MetadataChangeEvent(proposedSnapshot=chart_snapshot)
        yield MetadataWorkUnit(id=chart_snapshot.urn, mce=mce)

    def _get_reports(self, space_token: str) -> Iterator[List[dict]]:
        try:
            with self.report.report_get_timer:
                for reports_page in self._get_paged_request_json(
                    f"{self.workspace_uri}/spaces/{space_token}/reports?filter=all",
                    "reports",
                    self.config.items_per_page,
                ):
                    self.report.report_get_api_called += 1
                    logger.debug(
                        f"Read {len(reports_page)} reports records from workspace {self.workspace_uri} space {space_token}"
                    )
                    if self.config.exclude_archived:
                        logger.debug(
                            f"Excluding archived reports since exclude_archived: {self.config.exclude_archived}"
                        )
                        reports_page = [
                            report
                            for report in reports_page
                            if not report.get("archived", False)
                        ]
                    yield reports_page
        except ModeRequestError as e:
            if _is_http_404(e):
                self.report.report_warning(
                    title="No Reports Found in Space",
                    message="No reports were found in the space. It may have been recently deleted.",
                    context=f"Space Token: {space_token}, Error: {str(e)}",
                )
            else:
                self.report.report_failure(
                    title="Failed to Retrieve Reports for Space",
                    message="Unable to retrieve reports for space token.",
                    context=f"Space Token: {space_token}, Error: {str(e)}",
                )

    def _get_datasets(self, space_token: str) -> Iterator[List[dict]]:
        try:
            with self.report.dataset_get_timer:
                for dataset_page in self._get_paged_request_json(
                    f"{self.workspace_uri}/spaces/{space_token}/datasets?filter=all",
                    "reports",
                    self.config.items_per_page,
                ):
                    self.report.dataset_get_api_called += 1
                    logger.debug(
                        f"Read {len(dataset_page)} datasets records from workspace {self.workspace_uri} space {space_token}"
                    )
                    yield dataset_page
        except ModeRequestError as e:
            if _is_http_404(e):
                self.report.report_warning(
                    title="No Datasets Found in Space",
                    message="No datasets were found in the space. It may have been recently deleted.",
                    context=f"Space Token: {space_token}, Error: {str(e)}",
                )
            else:
                self.report.report_failure(
                    title="Failed to Retrieve Datasets for Space",
                    message=f"Unable to retrieve datasets for space token {space_token}.",
                    context=f"Space Token: {space_token}, Error: {str(e)}",
                )

    def _get_queries(self, report_token: str) -> List[dict]:
        start = time.perf_counter()
        try:
            # This endpoint does not handle pagination properly
            queries = self._get_request_json(
                f"{self.workspace_uri}/reports/{report_token}/queries"
            )
            with self.report._lock:
                self.report.query_get_api_called += 1
            logger.debug(
                f"Read {len(queries)} queries records from workspace {self.workspace_uri} report {report_token}"
            )
            return queries.get("_embedded", {}).get("queries", [])
        except ModeRequestError as e:
            if _is_http_404(e):
                self.report.report_warning(
                    title="No Queries Found",
                    message="No queries found for the report token. Maybe the report is deleted...",
                    context=f"Report Token: {report_token}, Error: {str(e)}",
                )
            else:
                self.report.report_failure(
                    title="Failed to Retrieve Queries",
                    message="Unable to retrieve queries for report token.",
                    context=f"Report Token: {report_token}, Error: {str(e)}",
                )
            return []
        finally:
            elapsed = time.perf_counter() - start
            with self.report._lock:
                self.report.query_api_total_sec += elapsed

    def _get_charts(self, report_token: str, query_token: str) -> List[dict]:
        start = time.perf_counter()
        try:
            # This endpoint does not handle pagination properly
            charts = self._get_request_json(
                f"{self.workspace_uri}/reports/{report_token}/queries/{query_token}/charts"
            )
            with self.report._lock:
                self.report.chart_get_api_called += 1
            logger.debug(
                f"Read {len(charts)} charts records from workspace {self.workspace_uri} report {report_token} query {query_token}"
            )
            return charts.get("_embedded", {}).get("charts", [])
        except ModeRequestError as e:
            if _is_http_404(e):
                self.report.report_warning(
                    title="No Charts Found for Query",
                    message="No charts were found for the query. The query may have been recently deleted.",
                    context=f"Report Token: {report_token}, Query Token: {query_token}, Error: {str(e)}",
                )
            else:
                self.report.report_failure(
                    title="Failed to Retrieve Charts",
                    message="Unable to retrieve charts from Mode.",
                    context=f"Report Token: {report_token}, Query Token: {query_token}, Error: {str(e)}",
                )
            return []
        finally:
            elapsed = time.perf_counter() - start
            with self.report._lock:
                self.report.chart_api_total_sec += elapsed

    def _get_paged_request_json(
        self, url: str, key: str, per_page: int
    ) -> Iterator[List[Dict]]:
        page: int = 1
        while True:
            page_url = f"{url}&per_page={per_page}&page={page}"
            response = self._get_request_json(page_url)
            data: List[Dict] = response.get("_embedded", {}).get(key, [])
            if not data:
                break
            yield data
            page += 1

    def _get_request_json(self, url: str) -> Dict:
        r = tenacity.Retrying(
            wait=wait_exponential(
                multiplier=self.config.api_options.retry_backoff_multiplier,
                max=self.config.api_options.max_retry_interval,
            ),
            retry=retry_if_exception_type(
                (HTTPError429, HTTPError504, ConnectionError)
            ),
            stop=stop_after_attempt(self.config.api_options.max_attempts),
        )

        @r.wraps
        def get_request():
            curl_command = make_curl_command(self.session, "GET", url, "")
            logger.debug(f"Issuing request; curl equivalent: {curl_command}")

            try:
                with self.rate_limiter:
                    response = self.session.get(
                        url, timeout=self.config.api_options.timeout
                    )
                if response.status_code == 204:  # No content, don't parse json
                    return {}

                response.raise_for_status()
                return response.json()
            except HTTPError as http_error:
                error_response = http_error.response
                if error_response is None:
                    raise http_error
                if error_response.status_code == 429:
                    with self.report._lock:
                        self.report.num_requests_exceeding_rate_limit += 1
                    sleep_time = error_response.headers.get("retry-after")
                    if sleep_time is not None:
                        time.sleep(float(sleep_time))
                    raise HTTPError429(
                        str(http_error), response=error_response
                    ) from http_error
                elif error_response.status_code == 504:
                    with self.report._lock:
                        self.report.num_requests_retried_on_timeout += 1
                    time.sleep(0.1)
                    raise HTTPError504(
                        str(http_error), response=error_response
                    ) from http_error

                logger.debug(
                    f"Error response ({error_response.status_code}): {error_response.text}"
                )
                raise http_error

        return get_request()

    @staticmethod
    def _get_process_memory():
        process = psutil.Process(os.getpid())
        mem_info = process.memory_info()
        return {
            "rss": mem_info.rss / (1024 * 1024),
            "vms": mem_info.vms / (1024 * 1024),
            "shared": getattr(mem_info, "shared", 0) / (1024 * 1024),
            "data": getattr(mem_info, "data", 0) / (1024 * 1024),
        }

    @staticmethod
    def create_embed_aspect_mcp(
        entity_urn: str, embed_url: str
    ) -> MetadataChangeProposalWrapper:
        return MetadataChangeProposalWrapper(
            entityUrn=entity_urn,
            aspect=EmbedClass(renderUrl=embed_url),
        )

    def gen_space_key(self, space_token: str) -> SpaceKey:
        return SpaceKey(platform=self.platform, space_token=space_token)

    def construct_space_container(
        self, space_token: str, space_name: str
    ) -> Iterable[MetadataWorkUnit]:
        key = self.gen_space_key(space_token)
        yield from gen_containers(
            container_key=key,
            name=space_name,
            sub_types=[BIContainerSubTypes.MODE_COLLECTION],
            # TODO: Support extracting the documentation for a space.
        )

        # We have a somewhat atypical browse path here, since we include the workspace name
        # as what's effectively but not officially a platform instance.
        yield MetadataChangeProposalWrapper(
            entityUrn=key.as_urn(),
            aspect=BrowsePathsV2Class(path=self._browse_path_space()),
        ).as_workunit()

    def _process_report(
        self,
        space_token: str,
        report: dict,
    ) -> Iterable[MetadataWorkUnit]:
        """Process a single report: queries, dashboard, and charts."""
        report_token = report.get("token", "")
        report_name = report.get("name", "")
        logger.debug(f"Report: name: {report_name} token: {report_token}")

        report_start = time.perf_counter()
        try:
            yield from self._process_report_inner(space_token, report)
            with self.report._lock:
                self.report.num_reports_processed += 1
            logger.debug(
                f"Processed report {report_name} ({report_token}) "
                f"in {time.perf_counter() - report_start:.2f}s"
            )
        except Exception as e:
            try:
                logger.warning(
                    f"Failed to process report {report_name} ({report_token}) "
                    f"in space {space_token}",
                    exc_info=True,
                )
                self.report.report_failure(
                    title="Failed to Process Report",
                    message=f"Unexpected error processing report {report_name} ({report_token}).",
                    context=f"Space Token: {space_token}, Error: {str(e)}",
                    exc=e,
                )
            except Exception:
                # Guard against the error handler itself raising (e.g., a
                # serialization issue in report_failure). If this propagated,
                # ThreadedIteratorExecutor would abort all remaining workers.
                logger.error(
                    f"Failed to record error for report {report_token}",
                    exc_info=True,
                )

    def _process_dataset(
        self, space_token: str, dataset: dict
    ) -> Iterable[MetadataWorkUnit]:
        """Process a single dataset: queries and lineage."""
        dataset_token = dataset.get("token", "")
        try:
            queries = self._get_queries(dataset_token)
            for query in queries:
                yield from self.construct_query_or_dataset(
                    dataset_token,
                    query,
                    space_token=space_token,
                    report_info=dataset,
                    is_mode_dataset=True,
                )
        except Exception as e:
            logger.warning(
                f"Failed to process dataset {dataset_token} in space {space_token}",
                exc_info=True,
            )
            self.report.report_failure(
                title="Failed to Process Dataset",
                message=f"Unexpected error processing dataset {dataset_token}.",
                context=f"Space Token: {space_token}, Error: {str(e)}",
                exc=e,
            )

    def _process_report_inner(
        self,
        space_token: str,
        report: dict,
    ) -> Iterable[MetadataWorkUnit]:
        report_token = report.get("token", "")
        space_container_key = self.gen_space_key(space_token)

        queries = self._get_queries(report_token)
        with self.report._lock:
            self.report.num_queries_processed += len(queries)
        chart_urns: List[str] = []
        query_chart_data: List[
            Tuple[dict, Dict[str, SchemaFieldClass], List[dict]]
        ] = []

        for query in queries:
            query_mcps = self.construct_query_or_dataset(
                report_token,
                query,
                space_token=space_token,
                report_info=report,
                is_mode_dataset=False,
            )
            chart_fields: Dict[str, SchemaFieldClass] = {}
            for wu in query_mcps:
                if isinstance(
                    wu.metadata, MetadataChangeProposalWrapper
                ) and isinstance(wu.metadata.aspect, SchemaMetadataClass):
                    schema_metadata = wu.metadata.aspect
                    for field in schema_metadata.fields:
                        chart_fields.setdefault(field.fieldPath, field)
                yield wu

            # Skip chart API call when the query has 0 explorations.
            if query.get("explorations_count", None) == 0:
                charts: List[dict] = []
                with self.report._lock:
                    self.report.chart_api_calls_skipped += 1
            else:
                charts = self._get_charts(report_token, query.get("token", ""))
            with self.report._lock:
                self.report.num_charts_processed += len(charts)
            for chart in charts:
                chart_urn = builder.make_chart_urn(
                    self.platform, chart.get("token", "")
                )
                chart_urns.append(chart_urn)
            query_chart_data.append((query, chart_fields, charts))

        dashboard_tuple_from_report = self.construct_dashboard(
            space_token=space_token,
            report_info=report,
            chart_urns=chart_urns,
        )

        if dashboard_tuple_from_report is not None:
            (
                dashboard_snapshot_from_report,
                browse_mcpw,
            ) = dashboard_tuple_from_report

            mce = MetadataChangeEvent(proposedSnapshot=dashboard_snapshot_from_report)

            mcpw = MetadataChangeProposalWrapper(
                entityUrn=dashboard_snapshot_from_report.urn,
                aspect=SubTypesClass(typeNames=[BIAssetSubTypes.MODE_REPORT]),
            )
            yield mcpw.as_workunit()
            yield from add_dataset_to_container(
                container_key=space_container_key,
                dataset_urn=dashboard_snapshot_from_report.urn,
            )
            yield browse_mcpw.as_workunit()

            usage_statistics = DashboardUsageStatisticsClass(
                timestampMillis=round(datetime.now().timestamp() * 1000),
                viewsCount=report.get("view_count", 0),
            )

            yield MetadataChangeProposalWrapper(
                entityUrn=dashboard_snapshot_from_report.urn,
                aspect=usage_statistics,
            ).as_workunit()

            if self.config.ingest_embed_url is True:
                yield self.create_embed_aspect_mcp(
                    entity_urn=dashboard_snapshot_from_report.urn,
                    embed_url=f"{self.config.connect_uri}/{self.config.workspace}/reports/{report.get('token')}/embed",
                ).as_workunit()

            yield MetadataWorkUnit(id=dashboard_snapshot_from_report.urn, mce=mce)

        for query, chart_fields, charts in query_chart_data:
            for i, chart in enumerate(charts):
                yield from self.construct_chart_from_api_data(
                    i,
                    chart,
                    chart_fields,
                    query,
                    space_token=space_token,
                    report_info=report,
                    query_name=query.get("name", query.get("token", "unknown")),
                )

    def _collect_space_work_items(
        self, space_token: str, space_name: str
    ) -> Tuple[
        List[Tuple[str, dict]],
        List[Tuple[str, dict]],
        List[MetadataWorkUnit],
    ]:
        """Collect report and dataset work items for a space.

        Returns (report_args, dataset_args, container_workunits).
        """
        container_wus = list(self.construct_space_container(space_token, space_name))

        report_args: List[Tuple[str, dict]] = []
        for report_page in self._get_reports(space_token):
            for report in report_page:
                report_args.append((space_token, report))

        dataset_args: List[Tuple[str, dict]] = []
        for dataset_page in self._get_datasets(space_token):
            for dataset in dataset_page:
                dataset_args.append((space_token, dataset))

        return report_args, dataset_args, container_wus

    @classmethod
    def create(cls, config_dict: dict, ctx: PipelineContext) -> "ModeSource":
        config: ModeConfig = ModeConfig.model_validate(config_dict)
        return cls(ctx, config)

    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
        return [
            *super().get_workunit_processors(),
            StaleEntityRemovalHandler.create(
                self, self.config, self.ctx
            ).workunit_processor,
        ]

    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
        # Phase 1: Emit space containers and collect all work items
        all_report_args: List[Tuple[str, dict]] = []
        all_dataset_args: List[Tuple[str, dict]] = []
        for space_token, space_name in self.space_tokens.items():
            report_args, dataset_args, container_wus = self._collect_space_work_items(
                space_token, space_name
            )
            yield from container_wus
            all_report_args.extend(report_args)
            all_dataset_args.extend(dataset_args)

        # Phase 2: Process ALL reports across ALL spaces in one global pool
        if self.config.max_threads > 1:
            yield from ThreadedIteratorExecutor.process(
                worker_func=self._process_report,
                args_list=all_report_args,
                max_workers=self.config.max_threads,
            )
        else:
            for args in all_report_args:
                yield from self._process_report(*args)

        # Phase 3: Process ALL datasets across ALL spaces
        if self.config.max_threads > 1 and all_dataset_args:
            yield from ThreadedIteratorExecutor.process(
                worker_func=self._process_dataset,
                args_list=all_dataset_args,
                max_workers=self.config.max_threads,
            )
        else:
            for args in all_dataset_args:
                yield from self._process_dataset(*args)

        memory_used = self._get_process_memory()
        self.report.process_memory_used_mb = round(memory_used["rss"], 2)

    def get_report(self) -> SourceReport:
        return self.report
