from collections import defaultdict
from collections.abc import Mapping, Sequence
from functools import cached_property
from typing import AbstractSet, Union  # noqa: UP035

from dagster import AssetKey, AssetsDefinition, AssetSpec
from dagster._annotations import beta, public
from dagster._core.definitions.job_definition import JobDefinition
from dagster._core.definitions.repository_definition.repository_definition import (
    RepositoryDefinition,
)
from dagster._record import record

from dagster_airlift.core.airflow_instance import AirflowInstance
from dagster_airlift.core.serialization.compute import (
    AirliftMetadataMappingInfo,
    build_airlift_metadata_mapping_info,
)
from dagster_airlift.core.serialization.serialized_data import DagHandle, TaskHandle
from dagster_airlift.core.utils import (
    dag_handle_from_job,
    dag_handles_for_spec,
    get_producing_dag_ids,
    is_airflow_mapped_job,
    is_dag_mapped_asset_spec,
    is_peered_dag_asset_spec,
    is_task_mapped_asset_spec,
    peered_dag_handles_for_spec,
    spec_iterator,
    task_handles_for_spec,
)

MappedAsset = Union[AssetSpec, AssetsDefinition]


def _is_mapped_asset_spec(asset: AssetSpec) -> bool:
    return (
        is_task_mapped_asset_spec(asset)
        or is_dag_mapped_asset_spec(asset)
        or is_peered_dag_asset_spec(asset)
    )


@beta
@record
class AirflowDefinitionsData:
    """A class that holds data about the assets that are mapped to Airflow dags and tasks, and
    provides methods for retrieving information about the mappings.
    The user should not instantiate this class directly. It is provided when customizing the events
    that are generated by the Airflow sensor using the `event_transformer_fn` argument of
    :py:func:`build_defs_from_airflow_instance`.
    """

    airflow_instance: AirflowInstance
    resolved_repository: RepositoryDefinition

    @property
    def airflow_mapped_asset_specs(self) -> Mapping[AssetKey, AssetSpec]:
        """The assets that are mapped to Airflow tasks and dags."""
        return {
            spec.key: spec
            for spec in spec_iterator(self.resolved_repository.assets_defs_by_key.values())
            if _is_mapped_asset_spec(spec)
        }

    @property
    def airflow_mapped_jobs(self) -> Sequence[JobDefinition]:
        """Jobs mapping to Airflow dags."""
        return [
            job for job in self.resolved_repository.get_all_jobs() if is_airflow_mapped_job(job)
        ]

    @property
    def airflow_mapped_jobs_by_dag_handle(
        self,
    ) -> Mapping[DagHandle, JobDefinition]:
        """Jobs mapping to Airflow dags by dag_id."""
        return {dag_handle_from_job(job): job for job in self.airflow_mapped_jobs}

    @property
    def assets_per_job(self) -> Mapping[str, AbstractSet[AssetKey]]:
        """Assets per job mapping to Airflow dags."""
        return {
            job.name: self.assets_produced_by_dags[dag_handle.dag_id]
            for dag_handle, job in self.airflow_mapped_jobs_by_dag_handle.items()
        }

    @property
    def assets_produced_by_dags(self) -> Mapping[str, AbstractSet[AssetKey]]:
        """Assets produced by Airflow dags."""
        result = defaultdict(set)
        for spec in self.airflow_mapped_asset_specs.values():
            for dag_id in get_producing_dag_ids(spec):
                result[dag_id].add(spec.key)
        return result

    @public
    @property
    def instance_name(self) -> str:
        """The name of the Airflow instance."""
        return self.airflow_instance.name

    @cached_property
    def mapping_info(self) -> AirliftMetadataMappingInfo:
        return build_airlift_metadata_mapping_info(self.airflow_mapped_asset_specs.values())

    @public
    def task_ids_in_dag(self, dag_id: str) -> set[str]:
        """Returns the task ids within the given dag_id.

        Args:
            dag_id (str): The dag id.
        """
        return self.mapping_info.task_id_map[dag_id]

    @property
    def dag_ids_with_mapped_asset_keys(self) -> AbstractSet[str]:
        """All dag_ids that have asset keys explicitly mapped to them. This include peered dag assets."""
        # dag ids that have asset keys explicitly mapped to them, or tasks within them
        explicitly_mapped_dag_ids = self.mapping_info.dag_ids
        # dag ids that have a "peered" dag asset
        peered_dag_ids = {
            handle.dag_id for handle in self.peered_dag_asset_keys_by_dag_handle.keys()
        }
        return explicitly_mapped_dag_ids.union(peered_dag_ids)

    @cached_property
    def mapped_asset_keys_by_task_handle(self) -> Mapping[TaskHandle, AbstractSet[AssetKey]]:
        asset_keys_per_handle = defaultdict(set)
        for spec in self.airflow_mapped_asset_specs.values():
            if is_task_mapped_asset_spec(spec):
                task_handles = task_handles_for_spec(spec)
                for task_handle in task_handles:
                    asset_keys_per_handle[task_handle].add(spec.key)
        return asset_keys_per_handle

    # these dag handle properties are ripe for consolidation
    @cached_property
    def mapped_asset_keys_by_dag_handle(self) -> Mapping[DagHandle, AbstractSet[AssetKey]]:
        """Assets specifically mapped to each dag."""
        asset_keys_per_handle = defaultdict(set)
        for spec in self.airflow_mapped_asset_specs.values():
            if is_dag_mapped_asset_spec(spec):
                dag_handles = dag_handles_for_spec(spec)
                for dag_handle in dag_handles:
                    asset_keys_per_handle[dag_handle].add(spec.key)
        return asset_keys_per_handle

    @cached_property
    def peered_dag_asset_keys_by_dag_handle(self) -> Mapping[DagHandle, AbstractSet[AssetKey]]:
        """Autogenerated "peered" dag assets."""
        asset_keys_per_handle = defaultdict(set)
        for spec in self.airflow_mapped_asset_specs.values():
            if is_peered_dag_asset_spec(spec):
                dag_handles = peered_dag_handles_for_spec(spec)
                for dag_handle in dag_handles:
                    asset_keys_per_handle[dag_handle].add(spec.key)
        return asset_keys_per_handle

    @cached_property
    def all_asset_keys_by_dag_handle(self) -> Mapping[DagHandle, AbstractSet[AssetKey]]:
        """All asset keys mapped to each dag."""
        res = defaultdict(set)
        for handle, keys in self.mapped_asset_keys_by_dag_handle.items():
            res[handle].update(keys)
        for handle, keys in self.peered_dag_asset_keys_by_dag_handle.items():
            res[handle].update(keys)
        return res

    @public
    def asset_keys_in_task(self, dag_id: str, task_id: str) -> AbstractSet[AssetKey]:
        """Returns the asset keys that are mapped to the given task.

        Args:
            dag_id (str): The dag id.
            task_id (str): The task id.
        """
        return self.mapped_asset_keys_by_task_handle[TaskHandle(dag_id=dag_id, task_id=task_id)]
