from typing import List, NamedTuple, Optional, Sequence, Set, cast

from dagster import (
    AssetKey,
    DagsterEventType,
    DagsterInstance,
    EventRecordsFilter,
    _check as check,
)
from dagster._core.definitions.multi_dimensional_partitions import (
    MultiPartitionKey,
    MultiPartitionsDefinition,
)
from dagster._core.definitions.partition import (
    DynamicPartitionsDefinition,
    PartitionsDefinition,
    PartitionsSubset,
    StaticPartitionsDefinition,
)
from dagster._core.definitions.time_window_partitions import TimeWindowPartitionsDefinition
from dagster._core.instance import DynamicPartitionsStore
from dagster._core.storage.tags import (
    MULTIDIMENSIONAL_PARTITION_PREFIX,
    get_dimension_from_partition_tag,
)
from dagster._serdes import deserialize_json_to_dagster_namedtuple, whitelist_for_serdes

CACHEABLE_PARTITION_TYPES = (
    TimeWindowPartitionsDefinition,
    MultiPartitionsDefinition,
    StaticPartitionsDefinition,
)


@whitelist_for_serdes
class AssetStatusCacheValue(
    NamedTuple(
        "_AssetPartitionsStatusCacheValue",
        [
            ("latest_storage_id", int),
            ("partitions_def_id", Optional[str]),
            ("serialized_materialized_partition_subset", Optional[str]),
            # Future partition subset features will go here, e.g. stale partition subsets
        ],
    )
):
    """
    Set of asset fields that reflect partition materialization status. This is used to display
    global partition status in the asset view.

    Properties:
        latest_storage_id (int): The latest evaluated storage id for the asset.
        partitions_def_id (Optional(str)): The serializable unique identifier for the partitions
            definition. When this value differs from the new partitions definition, this cache
            value needs to be recalculated. None if the asset is unpartitioned.
        serialized_materialized_partition_subset (Optional(str)): The serialized representation of the
            materialized partition subsets, up to the latest storage id. None if the asset is
            unpartitioned.
    """

    def __new__(
        cls,
        latest_storage_id: int,
        partitions_def_id: Optional[str] = None,
        serialized_materialized_partition_subset: Optional[str] = None,
    ):
        check.int_param(latest_storage_id, "latest_storage_id")
        check.opt_inst_param(
            serialized_materialized_partition_subset,
            "serialized_materialized_partition_subset",
            str,
        )
        check.opt_str_param(partitions_def_id, "partitions_def_id")
        return super(AssetStatusCacheValue, cls).__new__(
            cls, latest_storage_id, partitions_def_id, serialized_materialized_partition_subset
        )

    @staticmethod
    def from_db_string(db_string):
        if not db_string:
            return None

        cached_data = deserialize_json_to_dagster_namedtuple(db_string)
        if not isinstance(cached_data, AssetStatusCacheValue):
            return None

        return cached_data

    def deserialize_materialized_partition_subsets(
        self, partitions_def: PartitionsDefinition
    ) -> PartitionsSubset:
        if not self.serialized_materialized_partition_subset:
            return partitions_def.empty_subset()

        return partitions_def.deserialize_subset(self.serialized_materialized_partition_subset)


def get_materialized_multipartitions(
    instance: DagsterInstance, asset_key: AssetKey, partitions_def: MultiPartitionsDefinition
) -> Sequence[str]:
    dimension_names = partitions_def.partition_dimension_names
    materialized_keys: List[MultiPartitionKey] = []
    for event_tags in instance.get_event_tags_for_asset(asset_key):
        event_partition_keys_by_dimension = {
            get_dimension_from_partition_tag(key): value
            for key, value in event_tags.items()
            if key.startswith(MULTIDIMENSIONAL_PARTITION_PREFIX)
        }

        if all(
            dimension_name in event_partition_keys_by_dimension.keys()
            for dimension_name in dimension_names
        ):
            materialized_keys.append(
                MultiPartitionKey(
                    {
                        dimension_names[0]: event_partition_keys_by_dimension[dimension_names[0]],
                        dimension_names[1]: event_partition_keys_by_dimension[dimension_names[1]],
                    }
                )
            )
    return materialized_keys


def get_validated_partition_keys(
    dynamic_partitions_store: DynamicPartitionsStore,
    partitions_def: PartitionsDefinition,
    partition_keys: Set[str],
):
    if isinstance(partitions_def, (DynamicPartitionsDefinition, StaticPartitionsDefinition)):
        validated_partitions = (
            set(
                partitions_def.get_partition_keys(dynamic_partitions_store=dynamic_partitions_store)
            )
            & partition_keys
        )
    elif isinstance(partitions_def, MultiPartitionsDefinition):
        partition_keys_by_dimension = {
            dim.name: dim.partitions_def.get_partition_keys(
                dynamic_partitions_store=dynamic_partitions_store
            )
            for dim in partitions_def.partitions_defs
        }
        validated_partitions = set()
        for partition_key in partition_keys:
            multipartition_key = partitions_def.get_partition_key_from_str(partition_key)
            if all(
                key in partition_keys_by_dimension.get(dim, [])
                for dim, key in cast(
                    MultiPartitionKey, multipartition_key
                ).keys_by_dimension.items()
            ):
                validated_partitions.add(partition_key)
    else:
        if not isinstance(partitions_def, TimeWindowPartitionsDefinition):
            check.failed("Unexpected partitions definition type {partitions_def}")
        validated_partitions = {
            pk for pk in partition_keys if partitions_def.is_valid_partition_key(pk)
        }
    return validated_partitions


def _build_status_cache(
    instance: DagsterInstance,
    asset_key: AssetKey,
    latest_storage_id: int,
    partitions_def: Optional[PartitionsDefinition],
) -> AssetStatusCacheValue:
    """
    This method refreshes the asset status cache for a given asset key. It recalculates
    the materialized partition subset for the asset key and updates the cache value.
    """
    if not partitions_def or not isinstance(partitions_def, CACHEABLE_PARTITION_TYPES):
        return AssetStatusCacheValue(latest_storage_id=latest_storage_id)

    materialized_keys: Sequence[str]
    if isinstance(partitions_def, MultiPartitionsDefinition):
        materialized_keys = get_materialized_multipartitions(instance, asset_key, partitions_def)
    else:
        materialized_keys = [
            key
            for key, count in instance.get_materialization_count_by_partition([asset_key])
            .get(asset_key, {})
            .items()
            if count > 0
        ]

    serialized_materialized_partition_subset = partitions_def.empty_subset()

    serialized_materialized_partition_subset = (
        serialized_materialized_partition_subset.with_partition_keys(
            get_validated_partition_keys(instance, partitions_def, set(materialized_keys))
        )
    )

    return AssetStatusCacheValue(
        latest_storage_id=latest_storage_id,
        partitions_def_id=partitions_def.serializable_unique_identifier,
        serialized_materialized_partition_subset=serialized_materialized_partition_subset.serialize(),
    )


def _get_updated_status_cache(
    instance: DagsterInstance,
    asset_key: AssetKey,
    current_status_cache_value: AssetStatusCacheValue,
    partitions_def: Optional[PartitionsDefinition],
) -> AssetStatusCacheValue:
    """
    This method accepts the current asset status cache value, and fetches unevaluated
    records from the event log. It then updates the cache value with the new materializations.
    """
    unevaluated_event_records = instance.get_event_records(
        event_records_filter=EventRecordsFilter(
            event_type=DagsterEventType.ASSET_MATERIALIZATION,
            asset_key=asset_key,
            after_cursor=current_status_cache_value.latest_storage_id,
        )
    )

    if not unevaluated_event_records:
        return current_status_cache_value

    latest_storage_id = max([record.storage_id for record in unevaluated_event_records])
    if not partitions_def or not isinstance(partitions_def, CACHEABLE_PARTITION_TYPES):
        return AssetStatusCacheValue(latest_storage_id=latest_storage_id)

    check.invariant(
        current_status_cache_value.partitions_def_id
        == partitions_def.serializable_unique_identifier
    )
    materialized_subset: PartitionsSubset = (
        partitions_def.deserialize_subset(
            current_status_cache_value.serialized_materialized_partition_subset
        )
        if current_status_cache_value
        and current_status_cache_value.serialized_materialized_partition_subset
        else partitions_def.empty_subset()
    )
    newly_materialized_partitions = set()

    for unevaluated_materializations in unevaluated_event_records:
        if not (
            unevaluated_materializations.event_log_entry.dagster_event
            and unevaluated_materializations.event_log_entry.dagster_event.is_step_materialization
        ):
            check.failed("Expected materialization event")
        if unevaluated_materializations.event_log_entry.dagster_event.partition:
            newly_materialized_partitions.add(
                unevaluated_materializations.event_log_entry.dagster_event.partition
            )

    materialized_subset = materialized_subset.with_partition_keys(
        get_validated_partition_keys(instance, partitions_def, newly_materialized_partitions)
    )
    return AssetStatusCacheValue(
        latest_storage_id=latest_storage_id,
        partitions_def_id=current_status_cache_value.partitions_def_id,
        serialized_materialized_partition_subset=materialized_subset.serialize(),
    )


def _fetch_stored_asset_status_cache_value(
    instance: DagsterInstance, asset_key: AssetKey
) -> Optional[AssetStatusCacheValue]:
    asset_records = (
        instance.get_asset_records()
        if not asset_key
        else instance.get_asset_records(asset_keys=[asset_key])
    )
    if not asset_records:
        return None
    else:
        return list(asset_records)[0].asset_entry.cached_status


def _get_fresh_asset_status_cache_value(
    instance: DagsterInstance,
    asset_key: AssetKey,
    partitions_def: Optional[PartitionsDefinition] = None,
) -> Optional[AssetStatusCacheValue]:
    cached_status_data = _fetch_stored_asset_status_cache_value(instance, asset_key)

    updated_cache_value = None
    if cached_status_data is None or cached_status_data.partitions_def_id != (
        partitions_def.serializable_unique_identifier if partitions_def else None
    ):
        event_records = instance.get_event_records(
            event_records_filter=EventRecordsFilter(
                event_type=DagsterEventType.ASSET_MATERIALIZATION,
                asset_key=asset_key,
            ),
            limit=1,
        )
        if event_records:
            updated_cache_value = _build_status_cache(
                instance=instance,
                asset_key=asset_key,
                partitions_def=partitions_def,
                latest_storage_id=next(iter(event_records)).storage_id,
            )
    else:
        updated_cache_value = _get_updated_status_cache(
            instance=instance,
            asset_key=asset_key,
            partitions_def=partitions_def,
            current_status_cache_value=cached_status_data,
        )

    return updated_cache_value


def get_and_update_asset_status_cache_value(
    instance: DagsterInstance,
    asset_key: AssetKey,
    partitions_def: Optional[PartitionsDefinition] = None,
) -> Optional[AssetStatusCacheValue]:
    updated_cache_value = _get_fresh_asset_status_cache_value(instance, asset_key, partitions_def)
    if updated_cache_value:
        instance.update_asset_cached_status_data(asset_key, updated_cache_value)

    return updated_cache_value
