#
# 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 (
    WaitPolicyNotSupported,
    find_registered_custom_partition_mapper,
    find_registered_custom_timetable,
    find_registered_custom_window,
    is_core_partition_mapper_import_path,
    is_core_timetable_import_path,
    is_core_wait_policy_import_path,
    is_core_window_import_path,
)

if TYPE_CHECKING:
    from airflow.partition_mappers.base import PartitionMapper
    from airflow.partition_mappers.wait_policy import WaitPolicy
    from airflow.partition_mappers.window import Window
    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

    d = _normalize_stringified_encoding_keys(d)
    if not isinstance(d, dict):
        return d
    if Encoding.TYPE in d and Encoding.VAR in d:
        return BaseSerialization.deserialize(d)
    return {k: smart_decode_trigger_kwargs(v) for k, v in d.items()}


def _normalize_stringified_encoding_keys(value):
    if isinstance(value, list):
        return [_normalize_stringified_encoding_keys(v) for v in value]
    if not isinstance(value, dict):
        return value

    if str(Encoding.TYPE) in value and str(Encoding.VAR) in value:
        return {
            Encoding.TYPE if k == str(Encoding.TYPE) else Encoding.VAR if k == str(Encoding.VAR) else k: (
                _normalize_stringified_encoding_keys(v)
            )
            for k, v in value.items()
        }
    return {k: _normalize_stringified_encoding_keys(v) for k, v in value.items()}


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
        ],
        access_control=var.get("access_control", {}),
    )


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.definitions.deadline import VariableInterval
    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)

    raw_interval = data[DeadlineAlertFields.INTERVAL]

    if raw_interval is None:
        raise ValueError(
            "DeadlineAlert interval is missing. This can happen after downgrading "
            "from a version that supports VariableInterval. Downgrade is not fully reversible."
        )

    interval: datetime.timedelta | VariableInterval

    # Backward compatibility: previously interval was stored as total_seconds() (float/int).
    # Handle numeric values by converting to timedelta.
    if isinstance(raw_interval, (int, float)):
        interval = datetime.timedelta(seconds=raw_interval)
    else:
        deserialized = deserialize(raw_interval)
        if isinstance(deserialized, (datetime.timedelta, VariableInterval)):
            interval = deserialized
        else:
            raise TypeError(f"Invalid interval type: {type(deserialized).__name__}")

    return SerializedDeadlineAlert(
        reference=reference,
        interval=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])


def decode_window(var: dict[str, Any]) -> Window:
    """
    Decode a previously serialized :class:`Window`.

    Custom windows must be registered via the ``windows`` plugin attribute;
    unregistered import paths are rejected up-front instead of being handed to
    ``import_string``. See :func:`encode_window` for the matching encode-side
    restriction.

    :meta private:
    """
    importable_string = var[Encoding.TYPE]
    if is_core_window_import_path(importable_string):
        window_cls: type[Window] = import_string(importable_string)
    else:
        window_cls = find_registered_custom_window(importable_string)
    return window_cls.deserialize(var[Encoding.VAR])


def decode_wait_policy(var: dict[str, Any]) -> WaitPolicy:
    """
    Decode a previously serialized :class:`WaitPolicy`.

    Only built-in trigger policies are accepted — a tampered serialized Dag
    naming a non-core import path is rejected up-front instead of being handed
    to ``import_string``. See :func:`encode_wait_policy` for the matching
    encode-side restriction.

    :meta private:
    """
    importable_string = var[Encoding.TYPE]
    if not is_core_wait_policy_import_path(importable_string):
        raise WaitPolicyNotSupported(importable_string)
    policy_cls: type[WaitPolicy] = import_string(importable_string)
    return policy_cls.deserialize(var[Encoding.VAR])
