#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Task group utilities for UI API services."""

from __future__ import annotations

from functools import cache
from typing import TYPE_CHECKING

from airflow.configuration import conf
from airflow.serialization.definitions.baseoperator import SerializedBaseOperator
from airflow.serialization.definitions.mappedoperator import SerializedMappedOperator, is_mapped

if TYPE_CHECKING:
    from collections.abc import Callable
    from typing import Any

    from airflow.serialization.definitions.taskgroup import SerializedTaskGroup


@cache
def get_task_group_children_getter() -> Callable:
    """Get the Task Group Children Getter for the Dag."""
    if conf.get("api", "grid_view_sorting_order") == "topological":
        return lambda task_group, group_dict=None: task_group.topological_sort(group_dict=group_dict)
    return lambda task_group, group_dict=None: task_group.hierarchical_alphabetical_sort()


def task_group_to_dict(task_item_or_group, *, group_dict=None, parent_group_is_mapped=False):
    """Create a nested dict representation of this TaskGroup and its children used to construct the Graph."""
    if isinstance(task := task_item_or_group, (SerializedBaseOperator, SerializedMappedOperator)):
        # we explicitly want the short task ID here, not the full doted notation if in a group
        task_display_name = task.task_display_name if task.task_display_name != task.task_id else task.label
        node_operator = {
            "id": task.task_id,
            "label": task_display_name,
            "operator": task.operator_name,
            "type": "task",
        }
        if task.is_setup:
            node_operator["setup_teardown_type"] = "setup"
        elif task.is_teardown:
            node_operator["setup_teardown_type"] = "teardown"
        if is_mapped(task) or parent_group_is_mapped:
            node_operator["is_mapped"] = True
        return node_operator

    task_group = task_item_or_group
    if group_dict is None:
        group_dict = task_group.dag.task_group.get_task_group_dict()
    mapped = is_mapped(task_group)
    children = [
        task_group_to_dict(
            child,
            parent_group_is_mapped=parent_group_is_mapped or mapped,
            group_dict=group_dict,
        )
        for child in get_task_group_children_getter()(task_group, group_dict)
    ]

    if task_group.upstream_group_ids or task_group.upstream_task_ids:
        # This is the join node used to reduce the number of edges between two TaskGroup.
        children.append({"id": task_group.upstream_join_id, "label": "", "type": "join"})

    if task_group.downstream_group_ids or task_group.downstream_task_ids:
        # This is the join node used to reduce the number of edges between two TaskGroup.
        children.append({"id": task_group.downstream_join_id, "label": "", "type": "join"})

    node = {
        "id": task_group.group_id,
        "label": task_group.group_display_name or task_group.label,
        "tooltip": task_group.tooltip,
        "is_mapped": mapped,
        "children": children,
        "type": "task",
    }
    return node


def task_group_to_dict_grid(
    task_item_or_group,
    *,
    group_dict: dict[str | None, SerializedTaskGroup] | None = None,
    parent_group_is_mapped: bool = False,
) -> dict[str, Any]:
    """
    Create a nested dict representation of this TaskGroup and its children used to construct the Grid.

    :param group_dict: A ``{group_id: group}`` map used to resolve cross-group
        dependencies. Built once at the top of a render and threaded through the
        recursion so nested groups reuse it.
    :param parent_group_is_mapped: Whether an ancestor task group is mapped, propagated to children.
    """
    node: dict[str, Any]

    if isinstance(task := task_item_or_group, (SerializedMappedOperator, SerializedBaseOperator)):
        mapped = None
        if parent_group_is_mapped or is_mapped(task):
            mapped = True
        setup_teardown_type = None
        if task.is_setup is True:
            setup_teardown_type = "setup"
        elif task.is_teardown is True:
            setup_teardown_type = "teardown"
        # we explicitly want the short task ID here, not the full doted notation if in a group
        task_display_name = task.task_display_name if task.task_display_name != task.task_id else task.label
        node = {
            "id": task.task_id,
            "label": task_display_name,
            "is_mapped": mapped,
            "children": None,
            "setup_teardown_type": setup_teardown_type,
        }
        return node

    task_group = task_item_or_group
    if group_dict is None:
        group_dict = task_group.dag.task_group.get_task_group_dict()
    task_group_sort = get_task_group_children_getter()
    mapped = is_mapped(task_group)
    children = [
        task_group_to_dict_grid(
            child,
            group_dict=group_dict,
            parent_group_is_mapped=parent_group_is_mapped or mapped,
        )
        for child in task_group_sort(task_group, group_dict)
    ]

    node = {
        "id": task_group.group_id,
        "label": task_group.group_display_name or task_group.label,
        "is_mapped": mapped or None,
        "children": children or None,
    }
    if task_group.doc_md is not None:
        node["doc_md"] = task_group.doc_md
    return node
