import dataclasses
from collections import defaultdict
from collections.abc import Mapping, Sequence
from dataclasses import dataclass

import dagster as dg
import dagster._check as check
from dagster import AssetKey, AutomationContext
from dagster._core.asset_graph_view.asset_graph_view import AssetGraphView
from dagster._core.asset_graph_view.entity_subset import EntitySubset
from dagster._core.definitions.asset_daemon_cursor import AssetDaemonCursor
from dagster._core.definitions.auto_materialize_policy import AutoMaterializePolicy
from dagster._core.definitions.declarative_automation.automation_condition_evaluator import (
    AutomationConditionEvaluator,
)
from dagster._core.definitions.declarative_automation.operators.boolean_operators import (
    AndAutomationCondition,
)
from dagster._core.definitions.declarative_automation.serialized_objects import (
    AutomationConditionCursor,
)
from dagster._core.definitions.events import AssetKeyPartitionKey, CoercibleToAssetKey
from dagster._core.test_utils import freeze_time

from dagster_tests.declarative_automation_tests.scenario_utils.scenario_state import ScenarioState


class FalseAutomationCondition(dg.AutomationCondition):
    """Always returns the empty subset."""

    label: str | None = None

    @property
    def description(self) -> str:
        return ""

    def evaluate(self, context: AutomationContext) -> dg.AutomationResult:
        return dg.AutomationResult(context, true_subset=context.get_empty_subset())


@dataclass(frozen=True)
class AutomationConditionScenarioState(ScenarioState):
    automation_condition: dg.AutomationCondition | None = None
    condition_cursor: AutomationConditionCursor | None = None
    requested_asset_partitions: Sequence[AssetKeyPartitionKey] | None = None
    ensure_empty_result: bool = True
    request_backfills: bool = False

    def _get_request_subsets_by_key(
        self, asset_graph_view: AssetGraphView
    ) -> Mapping[dg.AssetKey, EntitySubset]:
        if self.requested_asset_partitions is None:
            return {}
        ap_by_key = defaultdict(set)
        for ap in self.requested_asset_partitions:
            ap_by_key[ap.asset_key].add(ap)
        return {
            asset_key: asset_graph_view.get_asset_subset_from_asset_partitions(asset_key, aps)
            for asset_key, aps in ap_by_key.items()
        }

    async def evaluate(
        self, asset: CoercibleToAssetKey
    ) -> tuple["AutomationConditionScenarioState", dg.AutomationResult]:
        asset_key = AssetKey.from_coercible(asset)
        # ensure that the top level condition never returns any asset partitions, as otherwise the
        # next evaluation will assume that those asset partitions were requested by the machinery
        asset_condition = (
            AndAutomationCondition(
                operands=[check.not_none(self.automation_condition), FalseAutomationCondition()]
            )
            if self.ensure_empty_result
            else check.not_none(self.automation_condition)
        )
        asset_graph = self.scenario_spec.with_asset_properties(
            keys=[asset],
            auto_materialize_policy=AutoMaterializePolicy.from_automation_condition(
                asset_condition
            ),
        ).asset_graph

        with freeze_time(self.current_time):
            cursor = AssetDaemonCursor.empty().with_updates(
                0, 0, [], [self.condition_cursor] if self.condition_cursor else [], asset_graph
            )
            evaluator = AutomationConditionEvaluator(
                asset_graph=asset_graph,
                instance=self.instance,
                entity_keys=asset_graph.get_all_asset_keys(),
                cursor=cursor,
                logger=self.logger,
                emit_backfills=False,
                evaluation_id=cursor.evaluation_id,
            )
            evaluator.request_subsets_by_key = self._get_request_subsets_by_key(
                evaluator.asset_graph_view
            )  # type: ignore
            context = AutomationContext.create(key=asset_key, evaluator=evaluator)

            full_result = await asset_condition.evaluate(context)  # type: ignore
            new_state = dataclasses.replace(self, condition_cursor=full_result.get_new_cursor())
            result = full_result.child_results[0] if self.ensure_empty_result else full_result

        return new_state, result

    def without_cursor(self) -> "AutomationConditionScenarioState":
        """Removes the previous evaluation state from the state. This is useful for testing
        re-evaluating this data "from scratch" after much computation has occurred.
        """
        return dataclasses.replace(self, condition_cursor=None)

    def with_requested_asset_partitions(
        self, requested_asset_partitions: Sequence[AssetKeyPartitionKey]
    ) -> "AutomationConditionScenarioState":
        return dataclasses.replace(self, requested_asset_partitions=requested_asset_partitions)
