import itertools
import warnings
from typing import AbstractSet, Any, Dict, Mapping, Optional, Sequence, Tuple, Union, cast

from dagster import check
from dagster.core.definitions.config import ConfigMapping
from dagster.core.definitions.decorators.op import op
from dagster.core.definitions.dependency import (
    DependencyDefinition,
    IDependencyDefinition,
    NodeInvocation,
)
from dagster.core.definitions.events import AssetKey
from dagster.core.definitions.executor_definition import ExecutorDefinition
from dagster.core.definitions.graph_definition import GraphDefinition
from dagster.core.definitions.job_definition import JobDefinition
from dagster.core.definitions.op_definition import OpDefinition
from dagster.core.definitions.output import Out, OutputDefinition
from dagster.core.definitions.partition import PartitionedConfig, PartitionsDefinition
from dagster.core.definitions.partition_key_range import PartitionKeyRange
from dagster.core.definitions.resource_definition import ResourceDefinition
from dagster.core.errors import DagsterInvalidDefinitionError
from dagster.core.execution.context.input import InputContext, build_input_context
from dagster.core.execution.context.output import build_output_context
from dagster.core.storage.fs_asset_io_manager import fs_asset_io_manager
from dagster.core.storage.root_input_manager import RootInputManagerDefinition, root_input_manager
from dagster.utils.backcompat import ExperimentalWarning, experimental
from dagster.utils.merger import merge_dicts

from .asset import AssetsDefinition
from .asset_partitions import get_upstream_partitions_for_partition_range
from .source_asset import SourceAsset


@experimental
def build_assets_job(
    name: str,
    assets: Sequence[AssetsDefinition],
    source_assets: Optional[Sequence[Union[SourceAsset, AssetsDefinition]]] = None,
    resource_defs: Optional[Mapping[str, ResourceDefinition]] = None,
    description: Optional[str] = None,
    config: Optional[Union[ConfigMapping, Dict[str, Any], PartitionedConfig]] = None,
    tags: Optional[Dict[str, Any]] = None,
    executor_def: Optional[ExecutorDefinition] = None,
) -> JobDefinition:
    """Builds a job that materializes the given assets.

    The dependencies between the ops in the job are determined by the asset dependencies defined
    in the metadata on the provided asset nodes.

    Args:
        name (str): The name of the job.
        assets (List[AssetsDefinition]): A list of assets or
            multi-assets - usually constructed using the :py:func:`@asset` or :py:func:`@multi_asset`
            decorator.
        source_assets (Optional[Sequence[Union[SourceAsset, AssetsDefinition]]]): A list of
            assets that are not materialized by this job, but that assets in this job depend on.
        resource_defs (Optional[Dict[str, ResourceDefinition]]): Resource defs to be included in
            this job.
        description (Optional[str]): A description of the job.

    Examples:

        .. code-block:: python

            @asset
            def asset1():
                return 5

            @asset
            def asset2(asset1):
                return my_upstream_asset + 1

            my_assets_job = build_assets_job("my_assets_job", assets=[asset1, asset2])

    Returns:
        JobDefinition: A job that materializes the given assets.
    """
    check.str_param(name, "name")
    check.sequence_param(assets, "assets", of_type=AssetsDefinition)
    check.opt_sequence_param(
        source_assets, "source_assets", of_type=(SourceAsset, AssetsDefinition)
    )
    check.opt_str_param(description, "description")
    source_assets_by_key = build_source_assets_by_key(source_assets)

    op_defs = build_op_deps(assets, source_assets_by_key.keys())
    root_manager = build_root_manager(source_assets_by_key)
    partitioned_config = build_job_partitions_from_assets(assets, source_assets or [])

    return GraphDefinition(
        name=name,
        node_defs=[asset.op for asset in assets],
        dependencies=op_defs,
        description=description,
        input_mappings=None,
        output_mappings=None,
        config=None,
    ).to_job(
        resource_defs=merge_dicts(
            {"io_manager": fs_asset_io_manager}, resource_defs or {}, {"root_manager": root_manager}
        ),
        config=config or partitioned_config,
        tags=tags,
        executor_def=executor_def,
    )


def build_job_partitions_from_assets(
    assets: Sequence[AssetsDefinition],
    source_assets: Sequence[Union[SourceAsset, AssetsDefinition]],
) -> Optional[PartitionedConfig]:
    assets_with_partitions_defs = [assets_def for assets_def in assets if assets_def.partitions_def]

    if len(assets_with_partitions_defs) == 0:
        return None

    first_assets_with_partitions_def: AssetsDefinition = assets_with_partitions_defs[0]
    for assets_def in assets_with_partitions_defs:
        if assets_def.partitions_def != first_assets_with_partitions_def.partitions_def:
            first_asset_key = next(iter(assets_def.asset_keys)).to_string()
            second_asset_key = next(iter(first_assets_with_partitions_def.asset_keys)).to_string()
            raise DagsterInvalidDefinitionError(
                "When an assets job contains multiple partitions assets, they must have the "
                f"same partitions definitions, but asset '{first_asset_key}' and asset "
                f"'{second_asset_key}' have different partitions definitions. "
            )

    partitions_defs_by_asset_key: Dict[AssetKey, PartitionsDefinition] = {}
    asset: Union[AssetsDefinition, SourceAsset]
    for asset in itertools.chain.from_iterable([assets, source_assets]):
        if isinstance(asset, AssetsDefinition) and asset.partitions_def is not None:
            for asset_key in asset.asset_keys:
                partitions_defs_by_asset_key[asset_key] = asset.partitions_def
        elif isinstance(asset, SourceAsset) and asset.partitions_def is not None:
            partitions_defs_by_asset_key[asset.key] = asset.partitions_def

    def asset_partitions_for_job_partition(
        job_partition_key: str,
    ) -> Mapping[AssetKey, PartitionKeyRange]:
        return {
            asset_key: PartitionKeyRange(job_partition_key, job_partition_key)
            for assets_def in assets
            for asset_key in assets_def.asset_keys
            if assets_def.partitions_def
        }

    def run_config_for_partition_fn(partition_key: str) -> Dict[str, Any]:
        ops_config: Dict[str, Any] = {}
        asset_partitions_by_asset_key = asset_partitions_for_job_partition(partition_key)

        for assets_def in assets:
            outputs_dict: Dict[str, Dict[str, Any]] = {}
            if assets_def.partitions_def is not None:
                for asset_key, output_def in assets_def.output_defs_by_asset_key.items():
                    asset_partition_key_range = asset_partitions_by_asset_key[asset_key]
                    outputs_dict[output_def.name] = {
                        "start": asset_partition_key_range.start,
                        "end": asset_partition_key_range.end,
                    }

            inputs_dict: Dict[str, Dict[str, Any]] = {}
            for in_asset_key, input_def in assets_def.input_defs_by_asset_key.items():
                upstream_partitions_def = partitions_defs_by_asset_key.get(in_asset_key)
                if assets_def.partitions_def is not None and upstream_partitions_def is not None:
                    upstream_partition_key_range = get_upstream_partitions_for_partition_range(
                        assets_def, upstream_partitions_def, in_asset_key, asset_partition_key_range
                    )
                    inputs_dict[input_def.name] = {
                        "start": upstream_partition_key_range.start,
                        "end": upstream_partition_key_range.end,
                    }

            ops_config[assets_def.op.name] = {
                "config": {
                    "assets": {
                        "input_partitions": inputs_dict,
                        "output_partitions": outputs_dict,
                    }
                }
            }

        return {"ops": ops_config}

    return PartitionedConfig(
        partitions_def=cast(PartitionsDefinition, first_assets_with_partitions_def.partitions_def),
        run_config_for_partition_fn=lambda p: run_config_for_partition_fn(p.name),
    )


def build_source_assets_by_key(
    source_assets: Optional[Sequence[Union[SourceAsset, AssetsDefinition]]]
) -> Mapping[AssetKey, Union[SourceAsset, OutputDefinition]]:
    source_assets_by_key: Dict[AssetKey, Union[SourceAsset, OutputDefinition]] = {}
    for asset_source in source_assets or []:
        if isinstance(asset_source, SourceAsset):
            source_assets_by_key[asset_source.key] = asset_source
        elif isinstance(asset_source, AssetsDefinition):
            for asset_key, output_def in asset_source.output_defs_by_asset_key.items():
                if asset_key:
                    source_assets_by_key[asset_key] = output_def

    return source_assets_by_key


def build_op_deps(
    multi_asset_defs: Sequence[AssetsDefinition], source_paths: AbstractSet[AssetKey]
) -> Dict[Union[str, NodeInvocation], Dict[str, IDependencyDefinition]]:
    op_outputs_by_asset: Dict[AssetKey, Tuple[OpDefinition, str]] = {}
    for multi_asset_def in multi_asset_defs:
        for asset_key, output_def in multi_asset_def.output_defs_by_asset_key.items():
            if asset_key in op_outputs_by_asset:
                raise DagsterInvalidDefinitionError(
                    f"The same asset key was included for two definitions: '{asset_key.to_string()}'"
                )

            op_outputs_by_asset[asset_key] = (multi_asset_def.op, output_def.name)

    op_deps: Dict[Union[str, NodeInvocation], Dict[str, IDependencyDefinition]] = {}
    for multi_asset_def in multi_asset_defs:
        op_name = multi_asset_def.op.name
        op_deps[op_name] = {}
        for asset_key, input_def in multi_asset_def.input_defs_by_asset_key.items():
            if asset_key in op_outputs_by_asset:
                op_def, output_name = op_outputs_by_asset[asset_key]
                op_deps[op_name][input_def.name] = DependencyDefinition(op_def.name, output_name)
            elif asset_key not in source_paths and not input_def.dagster_type.is_nothing:
                raise DagsterInvalidDefinitionError(
                    f"Input asset '{asset_key.to_string()}' for asset '{op_name}' is not "
                    "produced by any of the provided asset ops and is not one of the provided "
                    "sources"
                )

    return op_deps


def build_root_manager(
    source_assets_by_key: Mapping[AssetKey, Union[SourceAsset, OutputDefinition]]
) -> RootInputManagerDefinition:
    source_asset_io_manager_keys = {
        source_asset.io_manager_key for source_asset in source_assets_by_key.values()
    }
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=ExperimentalWarning)

        @root_input_manager(required_resource_keys=source_asset_io_manager_keys)
        def _root_manager(input_context: InputContext) -> Any:
            source_asset_key = cast(AssetKey, input_context.asset_key)
            source_asset = source_assets_by_key[source_asset_key]
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", category=ExperimentalWarning)

                @op(out={source_asset_key.path[-1]: Out(asset_key=source_asset_key)})
                def _op():
                    pass

            resource_config = input_context.step_context.resolved_run_config.resources[
                source_asset.io_manager_key
            ].config

            output_context = build_output_context(
                name=source_asset_key.path[-1],
                step_key="none",
                solid_def=_op,
                metadata=cast(Dict[str, Any], source_asset.metadata),
                resource_config=resource_config,
            )
            input_context_with_upstream = build_input_context(
                name=input_context.name,
                metadata=input_context.metadata,
                config=input_context.config,
                dagster_type=input_context.dagster_type,
                upstream_output=output_context,
                op_def=input_context.op_def,
                step_context=input_context.step_context,
                resource_config=resource_config,
            )

            io_manager = getattr(cast(Any, input_context.resources), source_asset.io_manager_key)
            return io_manager.load_input(input_context_with_upstream)

    return _root_manager
