#
# 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.

from __future__ import annotations

import datetime
from typing import TYPE_CHECKING, Any, TypeVar

import dateutil.relativedelta

from airflow._shared.module_loading import import_string
from airflow.serialization.definitions.assets import (
    SerializedAsset,
    SerializedAssetAlias,
    SerializedAssetAll,
    SerializedAssetAny,
    SerializedAssetBase,
    SerializedAssetNameRef,
    SerializedAssetUriRef,
    SerializedAssetWatcher,
)
from airflow.serialization.definitions.deadline import (
    DeadlineAlertFields,
    SerializedDeadlineAlert,
    SerializedReferenceModels,
)
from airflow.serialization.enums import DagAttributeTypes as DAT, Encoding
from airflow.serialization.helpers import (
    find_registered_custom_partition_mapper,
    find_registered_custom_timetable,
    is_core_partition_mapper_import_path,
    is_core_timetable_import_path,
)

if TYPE_CHECKING:
    from airflow.partition_mappers.base import PartitionMapper
    from airflow.timetables.base import Timetable as CoreTimetable

R = TypeVar("R")


def decode_relativedelta(var: dict[str, Any]) -> dateutil.relativedelta.relativedelta:
    """Dencode a relativedelta object."""
    copy_var = var.copy()
    if "weekday" in copy_var:
        copy_var["weekday"] = dateutil.relativedelta.weekday(*copy_var["weekday"])
    return dateutil.relativedelta.relativedelta(**copy_var)


def decode_interval(value: int | dict) -> datetime.timedelta | dateutil.relativedelta.relativedelta:
    if isinstance(value, dict):
        return decode_relativedelta(value)
    return datetime.timedelta(seconds=value)


def decode_run_immediately(value: bool | float) -> bool | datetime.timedelta:
    if isinstance(value, float):
        return datetime.timedelta(seconds=value)
    return value


def smart_decode_trigger_kwargs(d):
    """
    Slightly clean up kwargs for display or execution.

    This detects one level of BaseSerialization and tries to deserialize the
    content, removing some __type __var ugliness when the value is displayed
    in UI to the user and/or while execution.
    """
    from airflow.serialization.serialized_objects import BaseSerialization

    if not isinstance(d, dict) or Encoding.TYPE not in d:
        return d
    return BaseSerialization.deserialize(d)


def _decode_asset(var: dict[str, Any]):
    watchers = var.get("watchers", [])
    return SerializedAsset(
        name=var["name"],
        uri=var["uri"],
        group=var["group"],
        extra=var["extra"],
        watchers=[
            SerializedAssetWatcher(
                name=watcher["name"],
                trigger={
                    "classpath": watcher["trigger"]["classpath"],
                    "kwargs": smart_decode_trigger_kwargs(watcher["trigger"]["kwargs"]),
                },
            )
            for watcher in watchers
        ],
        allow_producer_teams=var.get("allow_producer_teams", []),
    )


def decode_asset_like(var: dict[str, Any]) -> SerializedAssetBase:
    """
    Decode a previously serialized asset-like object.

    :meta private:
    """
    typ = var[Encoding.TYPE]
    if Encoding.VAR in var:
        var = var[Encoding.VAR]
    else:
        var = {k: v for k, v in var.items() if k != Encoding.TYPE}
    match typ:
        case DAT.ASSET:
            return _decode_asset(var)
        case DAT.ASSET_ALL:
            return SerializedAssetAll([decode_asset_like(x) for x in var["objects"]])
        case DAT.ASSET_ANY:
            return SerializedAssetAny([decode_asset_like(x) for x in var["objects"]])
        case DAT.ASSET_ALIAS:
            return SerializedAssetAlias(name=var["name"], group=var["group"])
        case DAT.ASSET_REF:
            if "name" in var:
                return SerializedAssetNameRef(**var)
            return SerializedAssetUriRef(**var)
        case data_type:
            raise ValueError(f"deserialization not implemented for DAT {data_type!r}")


def decode_deadline_reference(reference_data: dict):
    """Decode a previously serialized deadline reference."""
    ref_name = reference_data.get(SerializedReferenceModels.REFERENCE_TYPE_FIELD)

    if ref_name and SerializedReferenceModels.is_builtin_reference(ref_name):
        reference_class = SerializedReferenceModels.get_reference_class(ref_name)
    else:
        reference_class = SerializedReferenceModels.SerializedCustomReference

    return reference_class.deserialize_reference(reference_data)


def decode_deadline_alert(encoded_data: dict):
    """
    Decode a previously serialized deadline alert.

    :meta private:
    """
    from airflow.sdk.serde import deserialize

    data = encoded_data.get(Encoding.VAR, encoded_data)

    reference_data = data[DeadlineAlertFields.REFERENCE]
    reference = decode_deadline_reference(reference_data)

    return SerializedDeadlineAlert(
        reference=reference,
        interval=datetime.timedelta(seconds=data[DeadlineAlertFields.INTERVAL]),
        callback=deserialize(data[DeadlineAlertFields.CALLBACK]),
        name=data.get(DeadlineAlertFields.NAME),
    )


def decode_timetable(var: dict[str, Any]) -> CoreTimetable:
    """
    Decode a previously serialized timetable.

    Most of the deserialization logic is delegated to the actual type, which
    we import from string.

    :meta private:
    """
    if is_core_timetable_import_path(importable_string := var[Encoding.TYPE]):
        timetable_type: type[CoreTimetable] = import_string(importable_string)
    else:
        timetable_type = find_registered_custom_timetable(importable_string)
    return timetable_type.deserialize(var[Encoding.VAR])


def decode_partition_mapper(var: dict[str, Any]) -> PartitionMapper:
    """
    Decode a previously serialized PartitionMapper.

    Most of the deserialization logic is delegated to the actual type, which
    we import from string.

    :meta private:
    """
    importable_string = var[Encoding.TYPE]
    if is_core_partition_mapper_import_path(importable_string):
        partition_mapper_cls = import_string(importable_string)
    else:
        partition_mapper_cls = find_registered_custom_partition_mapper(importable_string)
    return partition_mapper_cls.deserialize(var[Encoding.VAR])
