from collections import defaultdict
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    AbstractSet,
    Dict,
    Iterable,
    Mapping,
    Optional,
    Sequence,
    Union,
    cast,
)

import pendulum

import dagster._check as check
from dagster._core.definitions.asset_graph import AssetGraph
from dagster._core.definitions.data_version import (
    DATA_VERSION_TAG,
    DataVersion,
    extract_data_version_from_entry,
)
from dagster._core.definitions.events import AssetKey, AssetKeyPartitionKey
from dagster._core.definitions.partition import DynamicPartitionsDefinition
from dagster._core.definitions.time_window_partitions import TimeWindowPartitionsDefinition
from dagster._core.events import DagsterEventType
from dagster._core.instance import DagsterInstance, DynamicPartitionsStore
from dagster._core.storage.dagster_run import (
    DagsterRun,
    RunRecord,
)
from dagster._core.storage.tags import PARTITION_NAME_TAG
from dagster._utils.cached_method import cached_method

if TYPE_CHECKING:
    from dagster._core.storage.event_log import EventLogRecord
    from dagster._core.storage.event_log.base import AssetRecord


class CachingInstanceQueryer(DynamicPartitionsStore):
    """Provides utility functions for querying for asset-materialization related data from the
    instance which will attempt to limit redundant expensive calls. Intended for use within the
    scope of a single "request" (e.g. GQL request, sensor tick).

    Args:
        instance (DagsterInstance): The instance to query.
    """

    def __init__(
        self,
        instance: DagsterInstance,
        asset_graph: AssetGraph,
        evaluation_time: Optional[datetime] = None,
    ):
        self._instance = instance
        self._asset_graph = asset_graph

        self._asset_record_cache: Dict[AssetKey, Optional[AssetRecord]] = {}
        self._asset_partition_count_cache: Dict[
            Optional[int], Dict[AssetKey, Mapping[str, int]]
        ] = defaultdict(dict)

        self._dynamic_partitions_cache: Dict[str, Sequence[str]] = {}

        self._evaluation_time = evaluation_time if evaluation_time else pendulum.now("UTC")

    @property
    def instance(self) -> DagsterInstance:
        return self._instance

    @property
    def asset_graph(self) -> AssetGraph:
        return self._asset_graph

    @property
    def evaluation_time(self) -> datetime:
        return self._evaluation_time

    ####################
    # QUERY BATCHING
    ####################

    def prefetch_asset_partition_counts(
        self, asset_keys: Sequence[AssetKey], after_cursor: Optional[int]
    ):
        """For performance, batches together queries for selected assets."""
        if after_cursor is not None:
            self._asset_partition_count_cache[after_cursor] = dict(
                self.instance.get_materialization_count_by_partition(
                    asset_keys=asset_keys,
                    after_cursor=after_cursor,
                )
            )

    def prefetch_asset_records(self, asset_keys: Sequence[AssetKey]):
        """For performance, batches together queries for selected assets."""
        # get all asset records for the selected assets
        asset_records = self.instance.get_asset_records(asset_keys)
        for asset_record in asset_records:
            self._asset_record_cache[asset_record.asset_entry.asset_key] = asset_record
        for key in asset_keys:
            if key not in self._asset_record_cache:
                self._asset_record_cache[key] = None

    ####################
    # ASSET RECORDS / STORAGE IDS
    ####################

    def get_asset_record(self, asset_key: AssetKey) -> Optional["AssetRecord"]:
        if asset_key not in self._asset_record_cache:
            self._asset_record_cache[asset_key] = next(
                iter(self.instance.get_asset_records([asset_key])), None
            )
        return self._asset_record_cache[asset_key]

    def _event_type_for_key(self, asset_key: AssetKey) -> DagsterEventType:
        if self.asset_graph.is_source(asset_key):
            return DagsterEventType.ASSET_OBSERVATION
        else:
            return DagsterEventType.ASSET_MATERIALIZATION

    @cached_method
    def _get_latest_materialization_or_observation_record(
        self, *, asset_partition: AssetKeyPartitionKey, before_cursor: Optional[int] = None
    ) -> Optional["EventLogRecord"]:
        """Returns the latest event log record for the given asset partition of an asset. For
        observable source assets, this will be an AssetObservation, otherwise it will be an
        AssetMaterialization.
        """
        from dagster._core.event_api import EventRecordsFilter

        # in the simple case, just use the asset record
        if (
            before_cursor is None
            and asset_partition.partition_key is None
            and not self.asset_graph.is_observable(asset_partition.asset_key)
        ):
            asset_record = self.get_asset_record(asset_partition.asset_key)
            if asset_record is None:
                return None
            return asset_record.asset_entry.last_materialization_record

        records = self.instance.get_event_records(
            EventRecordsFilter(
                event_type=self._event_type_for_key(asset_partition.asset_key),
                asset_key=asset_partition.asset_key,
                asset_partitions=[asset_partition.partition_key]
                if asset_partition.partition_key
                else None,
                before_cursor=before_cursor,
            ),
            ascending=False,
            limit=1,
        )
        return next(iter(records), None)

    @cached_method
    def get_latest_storage_id_for_event_type(
        self, *, event_type: DagsterEventType
    ) -> Optional[int]:
        """Returns the latest storage id across all events of the given event_type.

        Args:
            event_type (DagsterEventType): The event type to query for.
        """
        from dagster._core.event_api import EventRecordsFilter

        latest_record = next(
            iter(
                self.instance.get_event_records(
                    event_records_filter=EventRecordsFilter(event_type=event_type),
                    limit=1,
                )
            ),
            None,
        )
        if latest_record is not None:
            return latest_record.storage_id
        return None

    @cached_method
    def _get_latest_materialization_or_observation_storage_ids_by_asset_partition(
        self, *, asset_key: AssetKey
    ) -> Mapping[AssetKeyPartitionKey, Optional[int]]:
        """Returns a mapping from asset partition to the latest storage id for that asset partition
        for all asset partitions associated with the given asset key.

        Note that for partitioned assets, an asset partition with a None partition key will be
        present in the mapping, representing the latest storage id for the asset as a whole.
        """
        asset_partition = AssetKeyPartitionKey(asset_key)
        latest_record = self._get_latest_materialization_or_observation_record(
            asset_partition=asset_partition
        )
        latest_storage_ids = {
            asset_partition: latest_record.storage_id if latest_record is not None else None
        }
        if self.asset_graph.is_partitioned(asset_key):
            latest_storage_ids.update(
                {
                    AssetKeyPartitionKey(asset_key, partition_key): storage_id
                    for partition_key, storage_id in self.instance.get_latest_storage_id_by_partition(
                        asset_key, event_type=self._event_type_for_key(asset_key)
                    ).items()
                }
            )
        return latest_storage_ids

    def get_latest_materialization_or_observation_storage_id(
        self, asset_partition: AssetKeyPartitionKey
    ) -> Optional[int]:
        """Returns the latest storage id for the given asset partition. If the asset has never been
        materialized, returns None.

        Args:
            asset_partition (AssetKeyPartitionKey): The asset partition to query.
        """
        return self._get_latest_materialization_or_observation_storage_ids_by_asset_partition(
            asset_key=asset_partition.asset_key
        ).get(asset_partition)

    def asset_partition_has_materialization_or_observation(
        self,
        asset_partition: AssetKeyPartitionKey,
        after_cursor: Optional[int] = None,
    ) -> bool:
        """Returns True if there is a materialization record for the given asset partition after
        the specified cursor.

        Args:
            asset_partition (AssetKeyPartitionKey): The asset partition to query.
            after_cursor (Optional[int]): Filter parameter such that only records with a storage_id
                greater than this value will be considered.
        """
        return (self.get_latest_materialization_or_observation_storage_id(asset_partition) or 0) > (
            after_cursor or 0
        )

    def get_latest_materialization_or_observation_record(
        self,
        asset_partition: AssetKeyPartitionKey,
        after_cursor: Optional[int] = None,
        before_cursor: Optional[int] = None,
    ) -> Optional["EventLogRecord"]:
        """Returns the latest record for the given asset partition given the specified cursors.

        Args:
            asset_partition (AssetKeyPartitionKey): The asset partition to query.
            after_cursor (Optional[int]): Filter parameter such that only records with a storage_id
                greater than this value will be considered.
            before_cursor (Optional[int]): Filter parameter such that only records with a storage_id
                less than this value will be considered.
        """
        check.param_invariant(
            not (after_cursor and before_cursor),
            "before_cursor",
            "Cannot set both before_cursor and after_cursor",
        )

        # first, do a quick check to eliminate the case where we know there is no record
        if not self.asset_partition_has_materialization_or_observation(
            asset_partition, after_cursor
        ):
            return None
        # then, if the before_cursor is after our latest record's storage id, we can just return
        # the latest record
        elif (before_cursor or 0) > (
            self.get_latest_materialization_or_observation_storage_id(asset_partition) or 0
        ):
            return self._get_latest_materialization_or_observation_record(
                asset_partition=asset_partition
            )
        # otherwise, do the explicit query
        return self._get_latest_materialization_or_observation_record(
            asset_partition=asset_partition, before_cursor=before_cursor
        )

    ####################
    # OBSERVATIONS
    ####################

    @cached_method
    def next_version_record(
        self,
        *,
        asset_key: AssetKey,
        after_cursor: Optional[int],
        data_version: Optional[DataVersion],
    ) -> Optional["EventLogRecord"]:
        from dagster._core.event_api import EventRecordsFilter

        for record in self.instance.get_event_records(
            EventRecordsFilter(
                event_type=DagsterEventType.ASSET_OBSERVATION,
                asset_key=asset_key,
                after_cursor=after_cursor,
            ),
            ascending=True,
        ):
            record_version = extract_data_version_from_entry(record.event_log_entry)
            if record_version is not None and record_version != data_version:
                return record

        # no records found with a new data version
        return None

    ####################
    # RUNS
    ####################

    @cached_method
    def _get_run_record_by_id(self, *, run_id: str) -> Optional[RunRecord]:
        return self.instance.get_run_record_by_id(run_id)

    def _get_run_by_id(self, run_id: str) -> Optional[DagsterRun]:
        run_record = self._get_run_record_by_id(run_id=run_id)
        if run_record is not None:
            return run_record.dagster_run
        return None

    def run_has_tag(self, run_id: str, tag_key: str, tag_value: str) -> bool:
        return cast(DagsterRun, self._get_run_by_id(run_id)).tags.get(tag_key) == tag_value

    @cached_method
    def _get_planned_materializations_for_run_from_events(
        self, *, run_id: str
    ) -> AbstractSet[AssetKey]:
        """Provides a fallback for fetching the planned materializations for a run from
        the ASSET_MATERIALIZATION_PLANNED events in the event log, in cases where this information
        is not available on the DagsterRun object.

        Args:
            run_id (str): The run id
        """
        materializations_planned = self.instance.get_records_for_run(
            run_id=run_id, of_type=DagsterEventType.ASSET_MATERIALIZATION_PLANNED
        ).records
        return set(cast(AssetKey, record.asset_key) for record in materializations_planned)

    def get_planned_materializations_for_run(self, run_id: str) -> AbstractSet[AssetKey]:
        """Returns the set of asset keys that are planned to be materialized by the run.

        Args:
            run_id (str): The run id
        """
        run = self._get_run_by_id(run_id=run_id)
        if run is None:
            return set()
        elif run.asset_selection:
            return run.asset_selection
        else:
            # must resort to querying the event log
            return self._get_planned_materializations_for_run_from_events(run_id=run_id)

    def is_asset_planned_for_run(
        self, run_id: str, asset: Union[AssetKey, AssetKeyPartitionKey]
    ) -> bool:
        """Returns True if the asset is planned to be materialized by the run."""
        run = self._get_run_by_id(run_id=run_id)
        if not run:
            return False

        if isinstance(asset, AssetKeyPartitionKey):
            asset_key = asset.asset_key
            if run.tags.get(PARTITION_NAME_TAG) != asset.partition_key:
                return False
        else:
            asset_key = asset

        return asset_key in self.get_planned_materializations_for_run(run_id=run_id)

    @cached_method
    def get_current_materializations_for_run(self, *, run_id: str) -> AbstractSet[AssetKey]:
        """Returns the set of asset keys that have been materialized by a given run.

        Args:
            run_id (str): The run id
        """
        materializations = self.instance.get_records_for_run(
            run_id=run_id,
            of_type=DagsterEventType.ASSET_MATERIALIZATION,
        ).records
        return set(cast(AssetKey, record.asset_key) for record in materializations)

    ####################
    # PARTITION COUNTS
    ####################

    def get_materialized_partition_counts(
        self, asset_key: AssetKey, after_cursor: Optional[int] = None
    ) -> Mapping[str, int]:
        """Returns a mapping from partition key to the number of times that partition has been
        materialized for a given asset.

        Args:
            asset_key (AssetKey): The asset key.
            after_cursor (Optional[int]): The cursor after which to look for materializations. If
                not provided, will look at all materializations.
        """
        if (
            after_cursor not in self._asset_partition_count_cache
            or asset_key not in self._asset_partition_count_cache[after_cursor]
        ):
            self._asset_partition_count_cache[after_cursor][
                asset_key
            ] = self.instance.get_materialization_count_by_partition(
                asset_keys=[asset_key], after_cursor=after_cursor
            )[
                asset_key
            ]
        return self._asset_partition_count_cache[after_cursor][asset_key]

    def get_materialized_partitions(
        self, asset_key: AssetKey, after_cursor: Optional[int] = None
    ) -> Iterable[str]:
        """Returns a list of the partitions that have been materialized for the given asset key.

        Args:
            asset_key (AssetKey): The asset key.
            after_cursor (Optional[int]): The cursor after which to look for materializations. If
                not provided, will look at all materializations.
        """
        return [
            partition_key
            for partition_key, count in self.get_materialized_partition_counts(
                asset_key, after_cursor
            ).items()
            if count > 0
        ]

    ####################
    # DYNAMIC PARTITIONS
    ####################

    def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
        """Returns a list of partitions for a partitions definition."""
        if partitions_def_name not in self._dynamic_partitions_cache:
            self._dynamic_partitions_cache[
                partitions_def_name
            ] = self.instance.get_dynamic_partitions(partitions_def_name)
        return self._dynamic_partitions_cache[partitions_def_name]

    def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool:
        return partition_key in self.get_dynamic_partitions(partitions_def_name)

    ####################
    # RECONCILIATION
    ####################

    def _asset_partitions_data_versions(
        self,
        asset_key: AssetKey,
        asset_partitions: Optional[Sequence[AssetKeyPartitionKey]],
        after_cursor: Optional[int] = None,
        before_cursor: Optional[int] = None,
    ) -> Mapping[AssetKeyPartitionKey, Optional[DataVersion]]:
        if not self.asset_graph.is_partitioned(asset_key):
            asset_partition = AssetKeyPartitionKey(asset_key)
            latest_record = self.get_latest_materialization_or_observation_record(
                asset_partition, after_cursor=after_cursor, before_cursor=before_cursor
            )
            return (
                {asset_partition: extract_data_version_from_entry(latest_record.event_log_entry)}
                if latest_record is not None
                else {}
            )
        else:
            query_result = self.instance._event_storage.get_latest_tags_by_partition(  # noqa
                asset_key,
                event_type=self._event_type_for_key(asset_key),
                tag_keys=[DATA_VERSION_TAG],
                after_cursor=after_cursor,
                before_cursor=before_cursor,
                asset_partitions=[
                    asset_partition.partition_key
                    for asset_partition in asset_partitions
                    if asset_partition.partition_key is not None
                ]
                if asset_partitions is not None
                else None,
            )
            return {
                AssetKeyPartitionKey(asset_key, partition_key): DataVersion(tags[DATA_VERSION_TAG])
                if tags.get(DATA_VERSION_TAG)
                else None
                for partition_key, tags in query_result.items()
            }

    def get_asset_partitions_updated_after_cursor(
        self,
        asset_key: AssetKey,
        asset_partitions: Optional[Sequence[AssetKeyPartitionKey]],
        after_cursor: Optional[int],
    ) -> Sequence[AssetKeyPartitionKey]:
        """Returns the set of asset partitions that have been updated after the given cursor.

        Args:
            asset_key (AssetKey): The asset key to check.
            asset_partitions (Optional[Sequence[AssetKeyPartitionKey]]): If supplied, will filter
                the set of checked partitions to the given partitions.
            after_cursor (Optional[int]): The cursor after which to look for updates.
        """
        if not self.asset_partition_has_materialization_or_observation(
            AssetKeyPartitionKey(asset_key), after_cursor=after_cursor
        ):
            return []
        # quick check that just compares latest storage ids
        updated_after_cursor = [
            asset_partition
            for asset_partition, latest_storage_id in self._get_latest_materialization_or_observation_storage_ids_by_asset_partition(
                asset_key=asset_key
            ).items()
            if (asset_partitions is None or asset_partition in asset_partitions)
            and (latest_storage_id or 0) > (after_cursor or 0)
        ]
        if not updated_after_cursor:
            return []
        # for regular assets, don't bother checking versions
        # in the future, we may remove this guard
        if after_cursor is None or not self.asset_graph.is_observable(asset_key):
            return updated_after_cursor

        # more expensive check to explicitly handle data versions
        latest_versions = self._asset_partitions_data_versions(
            asset_key, updated_after_cursor, after_cursor=after_cursor
        )
        previous_versions = self._asset_partitions_data_versions(
            asset_key, updated_after_cursor, before_cursor=after_cursor + 1
        )
        return [
            asset_partition
            for asset_partition, version in latest_versions.items()
            if previous_versions.get(asset_partition) != version
        ]

    @cached_method
    def is_reconciled(self, *, asset_partition: AssetKeyPartitionKey) -> bool:
        """Returns a boolean representing if the given `asset_partition` is currently reconciled.
        An asset (partition) is considered unreconciled if any of:
        - It has never been materialized
        - One of its parents has been updated more recently than it has
        - One of its parents is unreconciled.
        """
        # always treat source assets as reconciled
        if self.asset_graph.is_source(asset_partition.asset_key):
            return True
        elif not self.asset_partition_has_materialization_or_observation(asset_partition):
            return False

        time_or_dynamic_partitioned = isinstance(
            self.asset_graph.get_partitions_def(asset_partition.asset_key),
            (TimeWindowPartitionsDefinition, DynamicPartitionsDefinition),
        )

        parent_asset_partitions_by_key = defaultdict(list)

        for parent in self.asset_graph.get_parents_partitions(
            dynamic_partitions_store=self,
            current_time=self._evaluation_time,
            asset_key=asset_partition.asset_key,
            partition_key=asset_partition.partition_key,
        ).parent_partitions:
            parent_asset_partitions_by_key[parent.asset_key].append(parent)

        for parent_key, parent_asset_partitions in parent_asset_partitions_by_key.items():
            # ignore non-observable source parents
            if self.asset_graph.is_source(parent_key) and not self.asset_graph.is_observable(
                parent_key
            ):
                continue

            # when mapping from time or dynamic downstream to unpartitioned upstream, only check
            # for existence of upstream materialization, do not worry about timestamps
            if time_or_dynamic_partitioned and not self.asset_graph.is_partitioned(parent_key):
                return all(
                    self.asset_partition_has_materialization_or_observation(parent)
                    for parent in parent_asset_partitions
                )
            elif self.get_asset_partitions_updated_after_cursor(
                asset_key=parent_key,
                asset_partitions=parent_asset_partitions,
                after_cursor=self.get_latest_materialization_or_observation_storage_id(
                    asset_partition
                ),
            ):
                return False
            elif any(
                not self.is_reconciled(asset_partition=parent) for parent in parent_asset_partitions
            ):
                return False

        return True
