import contextlib
import datetime
import itertools
import json
import logging
import os
import random
import sys
from collections.abc import Iterable, Mapping, Sequence
from typing import AbstractSet, NamedTuple, Optional  # noqa: UP035
from unittest import mock

import dagster as dg
import dagster._check as check
import pytest
from dagster import AssetKey, AssetSelection, DagsterInstance, Nothing
from dagster._core.definitions.asset_daemon_cursor import AssetDaemonCursor
from dagster._core.definitions.assets.graph.asset_graph import AssetGraph
from dagster._core.definitions.assets.graph.asset_graph_subset import AssetGraphSubset
from dagster._core.definitions.assets.graph.base_asset_graph import BaseAssetGraph
from dagster._core.definitions.auto_materialize_policy import AutoMaterializePolicy
from dagster._core.definitions.auto_materialize_rule import AutoMaterializeRule
from dagster._core.definitions.auto_materialize_rule_evaluation import (
    AutoMaterializeDecisionType,
    AutoMaterializeRuleEvaluation,
    AutoMaterializeRuleEvaluationData,
)
from dagster._core.definitions.automation_tick_evaluation_context import (
    AutomationTickEvaluationContext,
)
from dagster._core.definitions.events import CoercibleToAssetKey
from dagster._core.definitions.observe import observe
from dagster._core.definitions.partitions.schedule_type import ScheduleType
from dagster._core.definitions.partitions.subset import PartitionsSubset
from dagster._core.definitions.partitions.utils import get_time_partitions_def
from dagster._core.definitions.timestamp import TimestampWithTimezone
from dagster._core.events import AssetMaterializationPlannedData, DagsterEventType
from dagster._core.execution.asset_backfill import AssetBackfillData
from dagster._core.execution.backfill import BulkActionStatus, PartitionBackfill
from dagster._core.remote_origin import InProcessCodeLocationOrigin
from dagster._core.test_utils import (
    InProcessTestWorkspaceLoadTarget,
    create_test_daemon_workspace_context,
    freeze_time,
)
from dagster._core.types.loadable_target_origin import LoadableTargetOrigin
from dagster._daemon.asset_daemon import AssetDaemon
from dagster._time import get_current_datetime
from dagster._utils import SingleInstigatorDebugCrashFlags


class RunSpec(NamedTuple):
    asset_keys: Sequence[dg.AssetKey]
    partition_key: str | None = None
    failed_asset_keys: Sequence[dg.AssetKey] | None = None
    is_observation: bool = False


class AssetEvaluationSpec(NamedTuple):
    asset_key: str
    rule_evaluations: Sequence[tuple[AutoMaterializeRuleEvaluation, Iterable[str] | None]]
    num_requested: int = 0
    num_skipped: int = 0
    num_discarded: int = 0

    @staticmethod
    def empty(asset_key: str) -> "AssetEvaluationSpec":
        return AssetEvaluationSpec(
            asset_key=asset_key,
            rule_evaluations=[],
            num_requested=0,
            num_skipped=0,
            num_discarded=0,
        )

    @staticmethod
    def from_single_rule(
        asset_key: str,
        rule: AutoMaterializeRule,
        evaluation_data: AutoMaterializeRuleEvaluationData | None = None,
    ) -> "AssetEvaluationSpec":
        return AssetEvaluationSpec(
            asset_key=asset_key,
            rule_evaluations=[
                (
                    AutoMaterializeRuleEvaluation(
                        rule_snapshot=rule.to_snapshot(), evaluation_data=evaluation_data
                    ),
                    None,
                )
            ],
            num_requested=1 if rule.decision_type == AutoMaterializeDecisionType.MATERIALIZE else 0,
            num_skipped=1 if rule.decision_type == AutoMaterializeDecisionType.SKIP else 0,
            num_discarded=1 if rule.decision_type == AutoMaterializeDecisionType.DISCARD else 0,
        )


class AssetReconciliationScenario(
    NamedTuple(
        "_AssetReconciliationScenario",
        [
            ("unevaluated_runs", Sequence[RunSpec]),
            ("assets", Sequence[dg.SourceAsset | dg.AssetsDefinition] | None),
            ("asset_checks", Sequence[dg.AssetChecksDefinition] | None),
            ("between_runs_delta", datetime.timedelta | None),
            ("evaluation_delta", datetime.timedelta | None),
            ("cursor_from", Optional["AssetReconciliationScenario"]),
            ("current_time", datetime.datetime | None),
            ("asset_selection", dg.AssetSelection | None),
            (
                "active_backfill_targets",
                Sequence[Mapping[dg.AssetKey, PartitionsSubset] | Sequence[dg.AssetKey]] | None,
            ),
            ("dagster_runs", Sequence[dg.DagsterRun] | None),
            ("event_log_entries", Sequence[dg.EventLogEntry] | None),
            ("expected_run_requests", Sequence[dg.RunRequest] | None),
            (
                "code_locations",
                Mapping[str, Sequence[dg.SourceAsset | dg.AssetsDefinition]] | None,
            ),
            ("expected_evaluations", Sequence[AssetEvaluationSpec] | None),
            ("requires_respect_materialization_data_versions", bool),
            ("supports_with_remote_asset_graph", bool),
            ("expected_error_message", str | None),
        ],
    )
):
    def __new__(
        cls,
        unevaluated_runs: Sequence[RunSpec],
        assets: Sequence[dg.SourceAsset | dg.AssetsDefinition] | None,
        asset_checks: Sequence[dg.AssetChecksDefinition] | None = None,
        between_runs_delta: datetime.timedelta | None = None,
        evaluation_delta: datetime.timedelta | None = None,
        cursor_from: Optional["AssetReconciliationScenario"] = None,
        current_time: datetime.datetime | None = None,
        asset_selection: dg.AssetSelection | None = None,
        active_backfill_targets: Sequence[
            Mapping[dg.AssetKey, PartitionsSubset] | Sequence[dg.AssetKey]
        ]
        | None = None,
        dagster_runs: Sequence[dg.DagsterRun] | None = None,
        event_log_entries: Sequence[dg.EventLogEntry] | None = None,
        expected_run_requests: Sequence[dg.RunRequest] | None = None,
        code_locations: Mapping[str, Sequence[dg.SourceAsset | dg.AssetsDefinition]] | None = None,
        expected_evaluations: Sequence[AssetEvaluationSpec] | None = None,
        requires_respect_materialization_data_versions: bool = False,
        supports_with_remote_asset_graph: bool = True,
        expected_error_message: str | None = None,
    ) -> "AssetReconciliationScenario":
        # For scenarios with no auto-materialize policies, we infer auto-materialize policies
        # and add them to the assets.
        assets_with_implicit_policies = assets
        if assets and all(
            (
                isinstance(a, dg.AssetsDefinition)
                and all(spec.auto_materialize_policy is None for spec in a.specs)
            )
            or isinstance(a, dg.SourceAsset)
            for a in assets
        ):
            asset_graph = AssetGraph.from_assets([*assets, *(asset_checks or [])])
            auto_materialize_asset_keys = (
                asset_selection.resolve(asset_graph)
                if asset_selection is not None
                else asset_graph.materializable_asset_keys
            )
            assets_with_implicit_policies = with_implicit_auto_materialize_policies(
                assets, asset_graph, auto_materialize_asset_keys
            )

        return super().__new__(
            cls,
            unevaluated_runs=unevaluated_runs,
            assets=assets_with_implicit_policies,
            asset_checks=asset_checks,
            between_runs_delta=between_runs_delta,
            evaluation_delta=evaluation_delta,
            cursor_from=cursor_from,
            current_time=current_time,
            asset_selection=asset_selection,
            active_backfill_targets=active_backfill_targets,
            dagster_runs=dagster_runs,
            event_log_entries=event_log_entries,
            expected_run_requests=expected_run_requests,
            code_locations=code_locations,
            expected_evaluations=expected_evaluations,
            requires_respect_materialization_data_versions=requires_respect_materialization_data_versions,
            supports_with_remote_asset_graph=supports_with_remote_asset_graph,
            expected_error_message=expected_error_message,
        )

    def _get_code_location_origin(
        self, scenario_name, location_name=None
    ) -> InProcessCodeLocationOrigin:
        """scenarios.py puts all the scenarios in its namespace under different 'hacky_daemon_repo_...' names."""
        return InProcessCodeLocationOrigin(
            loadable_target_origin=LoadableTargetOrigin(
                executable_path=sys.executable,
                module_name=(
                    "dagster_tests.declarative_automation_tests.legacy_tests.scenarios.scenarios"
                ),
                working_directory=os.getcwd(),
                attribute="hacky_daemon_repo_"
                + scenario_name
                + (f"_{location_name}" if location_name else ""),
            ),
            location_name=location_name or "test_location",
        )

    def do_sensor_scenario(
        self,
        instance,
        scenario_name=None,
        with_remote_asset_graph=False,
        respect_materialization_data_versions=False,
    ):
        if (
            self.requires_respect_materialization_data_versions
            and not respect_materialization_data_versions
        ):
            pytest.skip("requires respect_materialization_data_versions to be True")
        assert not self.code_locations, "setting code_locations not supported for sensor tests"

        test_time = self.current_time or get_current_datetime()

        with freeze_time(test_time):

            @dg.repository  # pyright: ignore[reportArgumentType]
            def repo():
                return self.assets

            # add any runs to the instance
            for dagster_run in self.dagster_runs or []:
                instance.add_run(dagster_run)
                # make sure to log the planned events
                for asset_key in dagster_run.asset_selection:  # pyright: ignore[reportOptionalIterable]
                    event = dg.DagsterEvent(
                        event_type_value=DagsterEventType.ASSET_MATERIALIZATION_PLANNED.value,
                        job_name=dagster_run.job_name,
                        event_specific_data=AssetMaterializationPlannedData(
                            asset_key, partition=(dagster_run.tags or {}).get("dagster/partition")
                        ),
                    )
                    instance.report_dagster_event(event, dagster_run.run_id, logging.DEBUG)

            # add any events to the instance
            for event_log_entry in self.event_log_entries or []:
                instance.store_event(event_log_entry)

            # add any backfills to the instance
            for i, target in enumerate(self.active_backfill_targets or []):
                if isinstance(target, Mapping):
                    target_subset = AssetGraphSubset(
                        partitions_subsets_by_asset_key=target,
                        non_partitioned_asset_keys=set(),
                    )
                else:
                    target_subset = AssetGraphSubset(
                        partitions_subsets_by_asset_key={},
                        non_partitioned_asset_keys=target,  # pyright: ignore[reportArgumentType]
                    )
                empty_subset = AssetGraphSubset(
                    partitions_subsets_by_asset_key={},
                    non_partitioned_asset_keys=set(),
                )
                asset_backfill_data = AssetBackfillData(
                    latest_storage_id=0,
                    target_subset=target_subset,
                    requested_runs_for_target_roots=False,
                    materialized_subset=empty_subset,
                    requested_subset=empty_subset,
                    failed_and_downstream_subset=empty_subset,
                    backfill_start_time=TimestampWithTimezone(test_time.timestamp(), "UTC"),
                )
                backfill = PartitionBackfill(
                    backfill_id=f"backfill{i}",
                    status=BulkActionStatus.REQUESTED,
                    from_failure=False,
                    tags={},
                    backfill_timestamp=test_time.timestamp(),
                    serialized_asset_backfill_data=asset_backfill_data.serialize(
                        dynamic_partitions_store=instance, asset_graph=repo.asset_graph
                    ),
                )
                instance.add_backfill(backfill)

            if self.cursor_from is not None:

                @dg.repository  # pyright: ignore[reportArgumentType]
                def prior_repo():
                    return self.cursor_from.assets  # pyright: ignore[reportOptionalMemberAccess]

                (
                    run_requests,
                    cursor,
                    evaluations,
                ) = self.cursor_from.do_sensor_scenario(
                    instance,
                    scenario_name=scenario_name,
                    with_remote_asset_graph=with_remote_asset_graph,
                )
                for run_request in run_requests:
                    asset_selection = check.not_none(run_request.asset_selection)
                    instance.create_run_for_job(
                        prior_repo.get_implicit_job_def_for_assets(asset_selection),
                        asset_selection=set(asset_selection),
                        tags=run_request.tags,
                    )

            else:
                cursor = AssetDaemonCursor.empty()

        start = datetime.datetime.now()

        def test_time_fn():
            return (test_time + (datetime.datetime.now() - start)).timestamp()

        for run in self.unevaluated_runs:
            if self.between_runs_delta is not None:
                test_time += self.between_runs_delta

            with freeze_time(test_time), mock.patch("time.time", new=test_time_fn):
                if run.is_observation:
                    observe(
                        instance=instance,
                        assets=[
                            a
                            for a in self.assets  # pyright: ignore[reportOptionalIterable]
                            if isinstance(a, dg.SourceAsset) and a.key in run.asset_keys
                        ],
                    )
                else:
                    do_run(
                        asset_keys=run.asset_keys,
                        partition_key=run.partition_key,
                        all_assets=self.assets,  # pyright: ignore[reportArgumentType]
                        instance=instance,
                        failed_asset_keys=run.failed_asset_keys,
                    )

        if self.evaluation_delta is not None:
            test_time += self.evaluation_delta
        with freeze_time(test_time):
            # get asset_graph
            if not with_remote_asset_graph:
                asset_graph = repo.asset_graph
            else:
                assert scenario_name is not None, "scenario_name must be provided for daemon runs"
                with create_test_daemon_workspace_context(
                    workspace_load_target=InProcessTestWorkspaceLoadTarget(
                        self._get_code_location_origin(scenario_name)
                    ),
                    instance=instance,
                ) as workspace_context:
                    workspace = workspace_context.create_request_context()
                    assert workspace.get_code_location_error("test_location") is None, (
                        workspace.get_code_location_error("test_location")
                    )
                    asset_graph = workspace.asset_graph

            with mock.patch.object(
                dg.DagsterInstance,
                "auto_materialize_respect_materialization_data_versions",
                new=lambda: self.respect_materialization_data_versions,  # pyright: ignore[reportAttributeAccessIssue]
            ):
                run_requests, cursor, evaluations = AutomationTickEvaluationContext(
                    evaluation_id=cursor.evaluation_id + 1,
                    asset_graph=asset_graph,
                    asset_selection=self.asset_selection or AssetSelection.all(),
                    instance=instance,
                    materialize_run_tags={},
                    observe_run_tags={},
                    cursor=cursor,
                    emit_backfills=False,
                    auto_observe_asset_keys={
                        key
                        for key in asset_graph.observable_asset_keys
                        if asset_graph.get(key).auto_observe_interval_minutes is not None
                    },
                    logger=logging.getLogger("dagster.amp"),
                ).evaluate()

        for run_request in run_requests:
            base_job = repo.get_implicit_job_def_for_assets(run_request.asset_selection)  # pyright: ignore[reportArgumentType]
            assert base_job is not None

        return run_requests, cursor, evaluations

    def do_daemon_scenario(
        self,
        instance,
        scenario_name,
        debug_crash_flags: SingleInstigatorDebugCrashFlags | None = None,
    ):
        assert bool(self.assets) != bool(self.code_locations), (
            "Must specify either assets or code_locations"
        )

        assert not self.active_backfill_targets, (
            "setting active_backfill_targets not supported for daemon tests"
        )

        test_time = self.current_time or get_current_datetime()

        with freeze_time(test_time) if self.current_time else contextlib.nullcontext():
            if self.cursor_from is not None:
                self.cursor_from.do_daemon_scenario(
                    instance,
                    scenario_name=scenario_name,
                )

        start = datetime.datetime.now()

        def test_time_fn():
            return (test_time + (datetime.datetime.now() - start)).timestamp()

        for run in self.unevaluated_runs:
            if self.between_runs_delta is not None:
                test_time += self.between_runs_delta

            with freeze_time(test_time), mock.patch("time.time", new=test_time_fn):
                assert not run.is_observation, "Observations not supported for daemon tests"
                if self.assets:
                    do_run(
                        asset_keys=run.asset_keys,
                        partition_key=run.partition_key,
                        all_assets=self.assets,
                        instance=instance,
                        failed_asset_keys=run.failed_asset_keys,
                    )
                else:
                    all_assets = [
                        a for assets in check.not_none(self.code_locations).values() for a in assets
                    ]
                    do_run(
                        asset_keys=run.asset_keys,
                        partition_key=run.partition_key,
                        all_assets=all_assets,  # This isn't quite right, it should be filtered to just the assets for the location
                        instance=instance,
                        failed_asset_keys=run.failed_asset_keys,
                    )

        if self.evaluation_delta is not None:
            test_time += self.evaluation_delta
        with freeze_time(test_time):
            assert scenario_name is not None, "scenario_name must be provided for daemon runs"

            if self.code_locations:
                target = InProcessTestWorkspaceLoadTarget(
                    [
                        self._get_code_location_origin(scenario_name, location_name)
                        for location_name in self.code_locations.keys()
                    ]
                )
            else:
                target = InProcessTestWorkspaceLoadTarget(
                    self._get_code_location_origin(scenario_name)
                )

            with create_test_daemon_workspace_context(
                workspace_load_target=target,
                instance=instance,
            ) as workspace_context:
                workspace = workspace_context.create_request_context()
                assert workspace.get_code_location_error("test_location") is None, (
                    workspace.get_code_location_error("test_location")
                )

                try:
                    AssetDaemon(  # noqa: SLF001
                        settings=instance.get_auto_materialize_settings(),
                        pre_sensor_interval_seconds=42,
                    )._run_iteration_impl(
                        workspace_context,
                        threadpool_executor=None,
                        amp_tick_futures={},
                        debug_crash_flags=(debug_crash_flags or {}),
                    )

                    if self.expected_error_message:
                        raise Exception(
                            f"Failed to raise expected error {self.expected_error_message}"
                        )

                except Exception:
                    if not self.expected_error_message:
                        raise

                    assert self.expected_error_message in str(sys.exc_info())


def do_run(
    asset_keys: Sequence[dg.AssetKey],
    partition_key: str | None,
    all_assets: Sequence[dg.SourceAsset | dg.AssetsDefinition],
    instance: DagsterInstance,
    failed_asset_keys: Sequence[dg.AssetKey] | None = None,
    tags: Mapping[str, str] | None = None,
) -> None:
    assets_in_run: list[dg.SourceAsset | dg.AssetsDefinition] = []
    asset_keys_set = set(asset_keys)
    for a in all_assets:
        if isinstance(a, dg.SourceAsset):
            assets_in_run.append(a)
        else:
            selected_keys = asset_keys_set.intersection(a.keys)
            if selected_keys == a.keys:
                assets_in_run.append(a)
            elif not selected_keys:
                assets_in_run.extend(a.to_source_assets())
            else:
                assets_in_run.append(a.subset_for(asset_keys_set, selected_asset_check_keys=None))
                assets_in_run.extend(
                    a.subset_for(
                        a.keys - selected_keys, selected_asset_check_keys=None
                    ).to_source_assets()
                )
    dg.materialize_to_memory(
        instance=instance,
        partition_key=partition_key,
        assets=assets_in_run,
        run_config={
            "ops": {
                failed_asset_key.path[-1]: {"config": {"fail": True}}
                for failed_asset_key in (failed_asset_keys or [])
            }
        },
        raise_on_error=False,
        tags=tags,
    )


def single_asset_run(asset_key: str, partition_key: str | None = None) -> RunSpec:
    return RunSpec(asset_keys=[AssetKey.from_coercible(asset_key)], partition_key=partition_key)


def run(
    asset_keys: Iterable[str],
    partition_key: str | None = None,
    failed_asset_keys: Iterable[str] | None = None,
    is_observation: bool = False,
):
    return RunSpec(
        asset_keys=list(
            map(AssetKey.from_coercible, itertools.chain(asset_keys, failed_asset_keys or []))
        ),
        failed_asset_keys=list(map(AssetKey.from_coercible, failed_asset_keys or [])),
        partition_key=partition_key,
        is_observation=is_observation,
    )


FAIL_TAG = "test/fail"


def run_request(
    asset_keys: dg.AssetKey | Sequence[CoercibleToAssetKey],
    partition_key: str | None = None,
    fail_keys: Sequence[str] | None = None,
    tags: Mapping[str, str] | None = None,
) -> dg.RunRequest:
    if isinstance(asset_keys, dg.AssetKey):
        asset_selection = [asset_keys]
    else:
        asset_selection = [AssetKey.from_coercible(key) for key in asset_keys]

    return dg.RunRequest(
        asset_selection=asset_selection,
        partition_key=partition_key,
        tags={**(tags or {}), **({FAIL_TAG: json.dumps(fail_keys)} if fail_keys else {})},
    )


def asset_def(
    key: str,
    deps: list[str] | Mapping[str, dg.PartitionMapping | None] | None = None,
    partitions_def: dg.PartitionsDefinition | None = None,
    legacy_freshness_policy: dg.LegacyFreshnessPolicy | None = None,
    auto_materialize_policy: dg.AutoMaterializePolicy | None = None,
    code_version: str | None = None,
    config_schema: Mapping[str, dg.Field] | None = None,
    **asset_def_kwargs,
) -> dg.AssetsDefinition:
    if deps is None:
        non_argument_deps = None
        ins = None
    elif isinstance(deps, list):
        non_argument_deps = deps
        ins = None
    else:
        non_argument_deps = None
        ins = {
            dep: dg.AssetIn(partition_mapping=partition_mapping, dagster_type=Nothing)
            for dep, partition_mapping in deps.items()
        }

    @dg.asset(
        name=key,
        partitions_def=partitions_def,
        deps=non_argument_deps,
        ins=ins,
        config_schema=config_schema or {"fail": dg.Field(bool, default_value=False)},
        legacy_freshness_policy=legacy_freshness_policy,
        auto_materialize_policy=auto_materialize_policy,
        code_version=code_version,
        **asset_def_kwargs,
    )
    def _asset(context, **kwargs):
        del kwargs

        if context.op_execution_context.op_config["fail"]:
            raise ValueError("")

    return _asset  # type: ignore


def multi_asset_def(
    keys: list[str],
    deps: list[str] | Mapping[str, set[str]] | None = None,
    can_subset: bool = False,
    legacy_freshness_policies: Mapping[str, dg.LegacyFreshnessPolicy] | None = None,
) -> dg.AssetsDefinition:
    if deps is None:
        non_argument_deps = None
        internal_asset_deps = None
    elif isinstance(deps, list):
        non_argument_deps = deps
        internal_asset_deps = None
    else:
        non_argument_deps = list(set().union(*deps.values()) - set(deps.keys()))
        internal_asset_deps = {k: {dg.AssetKey(vv) for vv in v} for k, v in deps.items()}

    @dg.multi_asset(
        outs={
            key: dg.AssetOut(
                is_required=not can_subset,
                legacy_freshness_policy=legacy_freshness_policies.get(key)
                if legacy_freshness_policies
                else None,
            )
            for key in keys
        },
        name="_".join(keys),
        deps=non_argument_deps,
        internal_asset_deps=internal_asset_deps,
        can_subset=can_subset,
    )
    def _assets(context):
        for output in keys:
            if output in context.op_execution_context.selected_output_names:
                yield dg.Output(output, output)

    return _assets


def observable_source_asset_def(
    key: str, partitions_def: dg.PartitionsDefinition | None = None, minutes_to_change: int = 0
):
    def _data_version() -> dg.DataVersion:
        return (
            dg.DataVersion(str(get_current_datetime().minute // minutes_to_change))
            if minutes_to_change
            else dg.DataVersion(str(random.random()))
        )

    @dg.observable_source_asset(name=key, partitions_def=partitions_def)
    def _observable():
        if partitions_def is None:
            return _data_version()
        else:
            return dg.DataVersionsByPartition(
                {partition: _data_version() for partition in partitions_def.get_partition_keys()}
            )

    return _observable


def with_auto_materialize_policy(
    assets_defs: Sequence[dg.AssetsDefinition], auto_materialize_policy: AutoMaterializePolicy
) -> Sequence[dg.AssetsDefinition]:
    """Note: this should be implemented in core dagster at some point, and this implementation is
    a lazy hack.
    """
    ret = []
    for assets_def in assets_defs:
        ret.append(
            assets_def.with_attributes(
                automation_condition=auto_materialize_policy.to_automation_condition()
            )
        )
    return ret


def get_implicit_auto_materialize_policy(
    asset_key: AssetKey, asset_graph: BaseAssetGraph
) -> dg.AutoMaterializePolicy | None:
    """For backcompat with pre-auto materialize policy graphs, assume a default scope of 1 day."""
    auto_materialize_policy = asset_graph.get(asset_key).auto_materialize_policy
    if auto_materialize_policy is None:
        time_partitions_def = get_time_partitions_def(asset_graph.get(asset_key).partitions_def)
        if time_partitions_def is None:
            max_materializations_per_minute = None
        elif time_partitions_def.schedule_type == ScheduleType.HOURLY:
            max_materializations_per_minute = 24
        else:
            max_materializations_per_minute = 1
        rules = {
            AutoMaterializeRule.materialize_on_missing(),
            AutoMaterializeRule.materialize_on_required_for_freshness(),
            AutoMaterializeRule.skip_on_parent_outdated(),
            AutoMaterializeRule.skip_on_parent_missing(),
            AutoMaterializeRule.skip_on_required_but_nonexistent_parents(),
            AutoMaterializeRule.skip_on_backfill_in_progress(),
        }
        if not bool(asset_graph.get_downstream_legacy_freshness_policies(asset_key=asset_key)):
            rules.add(AutoMaterializeRule.materialize_on_parent_updated())
        return dg.AutoMaterializePolicy(
            rules=rules,
            max_materializations_per_minute=max_materializations_per_minute,
        )
    return auto_materialize_policy


def with_implicit_auto_materialize_policies(
    assets_defs: Sequence[dg.SourceAsset | dg.AssetsDefinition],
    asset_graph: BaseAssetGraph,
    targeted_assets: AbstractSet[dg.AssetKey] | None = None,
) -> Sequence[dg.AssetsDefinition]:
    """Accepts a list of assets, adding implied auto-materialize policies to targeted assets
    if policies do not exist.
    """
    ret = []
    for assets_def in assets_defs:
        if (
            isinstance(assets_def, dg.AssetsDefinition)
            and not assets_def.auto_materialize_policies_by_key
        ):
            targeted_keys = (
                assets_def.keys & targeted_assets if targeted_assets else assets_def.keys
            )
            automation_conditions_by_key = {}
            for key in targeted_keys:
                policy = get_implicit_auto_materialize_policy(key, asset_graph)
                if policy:
                    automation_conditions_by_key[key] = policy.to_automation_condition()

            ret.append(
                assets_def.with_attributes(automation_condition=automation_conditions_by_key)
            )
        else:
            ret.append(assets_def)
    return ret
