# 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 contextlib
import datetime
import functools
from typing import TYPE_CHECKING, Any, TypeVar, overload

import attrs
import pendulum

from airflow._shared.module_loading import qualname
from airflow.partition_mappers.base import PartitionMapper as CorePartitionMapper
from airflow.sdk import (
    AllowedKeyMapper,
    Asset,
    AssetAlias,
    AssetAll,
    AssetAny,
    AssetOrTimeSchedule,
    ChainMapper,
    CronDataIntervalTimetable,
    CronTriggerTimetable,
    DeltaDataIntervalTimetable,
    DeltaTriggerTimetable,
    EventsTimetable,
    IdentityMapper,
    MultipleCronTriggerTimetable,
    PartitionMapper,
    ProductMapper,
    StartOfDayMapper,
    StartOfMonthMapper,
    StartOfQuarterMapper,
    StartOfWeekMapper,
    StartOfYearMapper,
)
from airflow.sdk.bases.timetable import BaseTimetable
from airflow.sdk.definitions.asset import AssetRef
from airflow.sdk.definitions.partition_mappers.temporal import StartOfHourMapper
from airflow.sdk.definitions.timetables.assets import (
    AssetTriggeredTimetable,
    PartitionedAssetTimetable,
)
from airflow.sdk.definitions.timetables.simple import ContinuousTimetable, NullTimetable, OnceTimetable
from airflow.sdk.definitions.timetables.trigger import CronPartitionTimetable
from airflow.serialization.decoders import decode_deadline_alert
from airflow.serialization.definitions.assets import (
    SerializedAsset,
    SerializedAssetAlias,
    SerializedAssetAll,
    SerializedAssetAny,
    SerializedAssetBase,
    SerializedAssetRef,
)
from airflow.serialization.definitions.deadline import SerializedDeadlineAlert
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,
)
from airflow.timetables.base import Timetable as CoreTimetable
from airflow.utils.docs import get_docs_url

if TYPE_CHECKING:
    from dateutil.relativedelta import relativedelta

    from airflow.sdk.definitions._internal.expandinput import ExpandInput
    from airflow.sdk.definitions.asset import BaseAsset
    from airflow.sdk.definitions.deadline import DeadlineAlert
    from airflow.triggers.base import BaseEventTrigger

    T = TypeVar("T")


def encode_expand_input(var: ExpandInput) -> dict[str, Any]:
    from airflow.serialization.serialized_objects import BaseSerialization

    return {"type": var.EXPAND_INPUT_TYPE, "value": BaseSerialization.serialize(var.value)}


def encode_relativedelta(var: relativedelta) -> dict[str, Any]:
    """Encode a relativedelta object."""
    encoded = {k: v for k, v in var.__dict__.items() if not k.startswith("_") and v}
    if var.weekday and var.weekday.n:
        # Every n'th Friday for example
        encoded["weekday"] = [var.weekday.weekday, var.weekday.n]
    elif var.weekday:
        encoded["weekday"] = [var.weekday.weekday]
    return encoded


def encode_timezone(var: str | pendulum.Timezone | pendulum.FixedTimezone) -> str | int:
    """
    Encode a Pendulum Timezone for serialization.

    Airflow only supports timezone objects that implements Pendulum's Timezone
    interface. We try to keep as much information as possible to make conversion
    round-tripping possible (see ``decode_timezone``). We need to special-case
    UTC; Pendulum implements it as a FixedTimezone (i.e. it gets encoded as
    0 without the special case), but passing 0 into ``pendulum.timezone`` does
    not give us UTC (but ``+00:00``).
    """
    if isinstance(var, str):
        return var
    if isinstance(var, pendulum.FixedTimezone):
        if var.offset == 0:
            return "UTC"
        return var.offset
    if isinstance(var, pendulum.Timezone):
        return var.name
    raise ValueError(
        f"DAG timezone should be a pendulum.tz.Timezone, not {var!r}. "
        f"See {get_docs_url('timezone.html#time-zone-aware-dags')}"
    )


def encode_interval(interval: datetime.timedelta | relativedelta) -> float | dict:
    if isinstance(interval, datetime.timedelta):
        return interval.total_seconds()
    return encode_relativedelta(interval)


def encode_run_immediately(value: bool | datetime.timedelta) -> bool | float:
    if isinstance(value, datetime.timedelta):
        return value.total_seconds()
    return value


def encode_trigger(trigger: BaseEventTrigger | dict):
    from airflow.serialization.serialized_objects import BaseSerialization

    def _ensure_serialized(d):
        """
        Make sure the kwargs dict is JSON-serializable.

        This is done with BaseSerialization logic. A simple check is added to
        ensure we don't double-serialize, which is possible when a trigger goes
        through multiple serialization layers.
        """
        if isinstance(d, dict) and Encoding.TYPE in d:
            return d
        return BaseSerialization.serialize(d)

    if isinstance(trigger, dict):
        classpath = trigger["classpath"]
        kwargs = trigger["kwargs"]
        # unwrap any kwargs that are themselves serialized objects, to avoid double-serialization in the trigger's own serialize() method.
        unwrapped = {}
        for k, v in kwargs.items():
            if isinstance(v, dict) and Encoding.TYPE in v:
                unwrapped[k] = BaseSerialization.deserialize(v)
            else:
                unwrapped[k] = v
        kwargs = unwrapped
    else:
        classpath, kwargs = trigger.serialize()
    return {
        "classpath": classpath,
        "kwargs": {k: _ensure_serialized(v) for k, v in kwargs.items()},
    }


def encode_asset_like(a: BaseAsset | SerializedAssetBase) -> dict[str, Any]:
    """
    Encode an asset-like object.

    :meta private:
    """
    d: dict[str, Any]
    match a:
        case Asset() | SerializedAsset():
            d = {"__type": DAT.ASSET, "name": a.name, "uri": a.uri, "group": a.group, "extra": a.extra}
            if a.watchers:
                d["watchers"] = [{"name": w.name, "trigger": encode_trigger(w.trigger)} for w in a.watchers]
            if a.allow_producer_teams:
                d["allow_producer_teams"] = a.allow_producer_teams
            return d
        case AssetAlias() | SerializedAssetAlias():
            return {"__type": DAT.ASSET_ALIAS, "name": a.name, "group": a.group}
        case AssetAll() | SerializedAssetAll():
            return {"__type": DAT.ASSET_ALL, "objects": [encode_asset_like(x) for x in a.objects]}
        case AssetAny() | SerializedAssetAny():
            return {"__type": DAT.ASSET_ANY, "objects": [encode_asset_like(x) for x in a.objects]}
        case AssetRef() | SerializedAssetRef():
            return {"__type": DAT.ASSET_REF, **attrs.asdict(a)}
    raise ValueError(f"serialization not implemented for {type(a).__name__!r}")


def encode_deadline_alert(d: DeadlineAlert | SerializedDeadlineAlert) -> dict[str, Any]:
    """
    Encode a deadline alert.

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

    return {
        "name": d.name,
        "reference": encode_deadline_reference(d.reference),
        "interval": d.interval.total_seconds(),
        "callback": serialize(d.callback),
    }


_BUILTIN_DEADLINE_MODULES = (
    "airflow.sdk.definitions.deadline",
    "airflow.serialization.definitions.deadline",
    # Include airflow.models.deadline to treat core's deadline references as builtins.
    # This is to maintain backcompat with 3.1.x custom refs that inherit from
    # airflow.models.deadline.ReferenceModels.BaseDeadlineReference.
    "airflow.models.deadline",
)


def encode_deadline_reference(ref) -> dict[str, Any]:
    """
    Encode a deadline reference.

    For custom (non-builtin) deadline references, includes the class path
    so the decoder can import the user's class at runtime.

    :meta private:
    """
    from airflow._shared.module_loading import qualname

    serialized = ref.serialize_reference()

    # Custom types (not built-in) need __class_path so the decoder can import them.
    # Unlike built-in types which are looked up in SerializedReferenceModels,
    # custom types are discovered via import_string(__class_path) at deserialization time.
    module = type(ref).__module__
    if module not in _BUILTIN_DEADLINE_MODULES:
        serialized["__class_path"] = qualname(ref)

    return serialized


def _get_serialized_timetable_import_path(var: BaseTimetable | CoreTimetable) -> str:
    # Find SDK classes.
    with contextlib.suppress(KeyError):
        return _serializer.BUILTIN_TIMETABLES[var_type := type(var)]

    # Check Core classes.
    if is_core_timetable_import_path(importable_string := qualname(var_type)):
        return importable_string

    # Find user-registered classes.
    find_registered_custom_timetable(importable_string)  # This raises if not found.
    return importable_string


def encode_timetable(var: BaseTimetable | CoreTimetable) -> dict[str, Any]:
    """
    Encode a timetable instance.

    See ``_Serializer.serialize()`` for more implementation detail.

    :meta private:
    """
    importable_string = _get_serialized_timetable_import_path(var)
    return {Encoding.TYPE: importable_string, Encoding.VAR: _serializer.serialize_timetable(var)}


class _Serializer:
    """Serialization logic."""

    BUILTIN_TIMETABLES: dict[type, str] = {
        AssetOrTimeSchedule: "airflow.timetables.assets.AssetOrTimeSchedule",
        AssetTriggeredTimetable: "airflow.timetables.simple.AssetTriggeredTimetable",
        ContinuousTimetable: "airflow.timetables.simple.ContinuousTimetable",
        CronDataIntervalTimetable: "airflow.timetables.interval.CronDataIntervalTimetable",
        CronTriggerTimetable: "airflow.timetables.trigger.CronTriggerTimetable",
        CronPartitionTimetable: "airflow.timetables.trigger.CronPartitionTimetable",
        DeltaDataIntervalTimetable: "airflow.timetables.interval.DeltaDataIntervalTimetable",
        DeltaTriggerTimetable: "airflow.timetables.trigger.DeltaTriggerTimetable",
        EventsTimetable: "airflow.timetables.events.EventsTimetable",
        MultipleCronTriggerTimetable: "airflow.timetables.trigger.MultipleCronTriggerTimetable",
        NullTimetable: "airflow.timetables.simple.NullTimetable",
        OnceTimetable: "airflow.timetables.simple.OnceTimetable",
        PartitionedAssetTimetable: "airflow.timetables.simple.PartitionedAssetTimetable",
    }

    @functools.singledispatchmethod
    def serialize_timetable(self, timetable: BaseTimetable | CoreTimetable) -> dict[str, Any]:
        """
        Serialize a timetable into a JSON-compatible dict for storage.

        All timetables defined in the SDK should be handled by registered
        single-dispatch variants below.

        This function's body should only be
        called on timetables defined in Core (under ``airflow.timetables``),
        and user-defined custom timetables registered via plugins, which also
        inherit from the Core timetable base class.

        For timetables in Core, serialization work is delegated to the type.
        """
        if not isinstance(timetable, CoreTimetable):
            raise NotImplementedError(f"can not serialize timetable {type(timetable).__name__}")
        return timetable.serialize()

    @serialize_timetable.register(ContinuousTimetable)
    @serialize_timetable.register(NullTimetable)
    @serialize_timetable.register(OnceTimetable)
    def _(self, timetable: ContinuousTimetable | NullTimetable | OnceTimetable) -> dict[str, Any]:
        return {}

    @serialize_timetable.register
    def _(self, timetable: AssetTriggeredTimetable) -> dict[str, Any]:
        return {"asset_condition": encode_asset_like(timetable.asset_condition)}

    @serialize_timetable.register
    def _(self, timetable: EventsTimetable) -> dict[str, Any]:
        return {
            "event_dates": [x.isoformat(sep="T") for x in timetable.event_dates],
            "restrict_to_events": timetable.restrict_to_events,
            "description": timetable.description,
        }

    @serialize_timetable.register
    def _(self, timetable: CronDataIntervalTimetable) -> dict[str, Any]:
        return {"expression": timetable.expression, "timezone": encode_timezone(timetable.timezone)}

    @serialize_timetable.register
    def _(self, timetable: DeltaDataIntervalTimetable) -> dict[str, Any]:
        return {"delta": encode_interval(timetable.delta)}

    @serialize_timetable.register
    def _(self, timetable: CronTriggerTimetable) -> dict[str, Any]:
        return {
            "expression": timetable.expression,
            "timezone": encode_timezone(timetable.timezone),
            "interval": encode_interval(timetable.interval),
            "run_immediately": encode_run_immediately(timetable.run_immediately),
        }

    @serialize_timetable.register
    def _(self, timetable: CronPartitionTimetable) -> dict[str, Any]:
        return {
            "expression": timetable.expression,
            "timezone": encode_timezone(timetable.timezone),
            "run_immediately": encode_run_immediately(timetable.run_immediately),
            "run_offset": timetable.run_offset,
            "key_format": timetable.key_format,
        }

    @serialize_timetable.register
    def _(self, timetable: DeltaTriggerTimetable) -> dict[str, Any]:
        return {
            "delta": encode_interval(timetable.delta),
            "interval": encode_interval(timetable.interval),
        }

    @serialize_timetable.register
    def _(self, timetable: MultipleCronTriggerTimetable) -> dict[str, Any]:
        # All timetables share the same timezone, interval, and run_immediately
        # values, so we can just use the first to represent them.
        representitive = timetable.timetables[0]
        return {
            "expressions": [t.expression for t in timetable.timetables],
            "timezone": encode_timezone(representitive.timezone),
            "interval": encode_interval(representitive.interval),
            "run_immediately": encode_run_immediately(representitive.run_immediately),
        }

    @serialize_timetable.register
    def _(self, timetable: AssetOrTimeSchedule) -> dict[str, Any]:
        return {
            "asset_condition": encode_asset_like(timetable.asset_condition),
            "timetable": encode_timetable(timetable.timetable),
        }

    @serialize_timetable.register
    def _(self, timetable: CoreTimetable) -> dict[str, Any]:
        return timetable.serialize()

    @serialize_timetable.register
    def _(self, timetable: PartitionedAssetTimetable) -> dict[str, Any]:
        return {
            "asset_condition": encode_asset_like(timetable.asset_condition),
            "default_partition_mapper": encode_partition_mapper(timetable.default_partition_mapper),
            "partition_mapper_config": [
                (encode_asset_like(asset), encode_partition_mapper(partition_mapper))
                for asset, partition_mapper in timetable.partition_mapper_config.items()
            ],
        }

    BUILTIN_PARTITION_MAPPERS: dict[type, str] = {
        AllowedKeyMapper: "airflow.partition_mappers.allowed_key.AllowedKeyMapper",
        ChainMapper: "airflow.partition_mappers.chain.ChainMapper",
        IdentityMapper: "airflow.partition_mappers.identity.IdentityMapper",
        ProductMapper: "airflow.partition_mappers.product.ProductMapper",
        StartOfDayMapper: "airflow.partition_mappers.temporal.StartOfDayMapper",
        StartOfHourMapper: "airflow.partition_mappers.temporal.StartOfHourMapper",
        StartOfMonthMapper: "airflow.partition_mappers.temporal.StartOfMonthMapper",
        StartOfQuarterMapper: "airflow.partition_mappers.temporal.StartOfQuarterMapper",
        StartOfWeekMapper: "airflow.partition_mappers.temporal.StartOfWeekMapper",
        StartOfYearMapper: "airflow.partition_mappers.temporal.StartOfYearMapper",
    }

    @functools.singledispatchmethod
    def serialize_partition_mapper(
        self, partition_mapper: PartitionMapper | CorePartitionMapper
    ) -> dict[str, Any]:
        if not isinstance(partition_mapper, CorePartitionMapper):
            raise NotImplementedError(f"can not serialize timetable {type(partition_mapper).__name__}")
        return partition_mapper.serialize()

    @serialize_partition_mapper.register
    def _(self, partition_mapper: ChainMapper) -> dict[str, Any]:
        return {"mappers": [encode_partition_mapper(m) for m in partition_mapper.mappers]}

    @serialize_partition_mapper.register
    def _(self, partition_mapper: IdentityMapper) -> dict[str, Any]:
        return {}

    @serialize_partition_mapper.register(StartOfHourMapper)
    @serialize_partition_mapper.register(StartOfDayMapper)
    @serialize_partition_mapper.register(StartOfWeekMapper)
    @serialize_partition_mapper.register(StartOfMonthMapper)
    @serialize_partition_mapper.register(StartOfQuarterMapper)
    @serialize_partition_mapper.register(StartOfYearMapper)
    def _(
        self,
        partition_mapper: StartOfHourMapper
        | StartOfDayMapper
        | StartOfWeekMapper
        | StartOfMonthMapper
        | StartOfQuarterMapper
        | StartOfYearMapper,
    ) -> dict[str, Any]:
        return {
            "input_format": partition_mapper.input_format,
            "output_format": partition_mapper.output_format,
        }

    @serialize_partition_mapper.register
    def _(self, partition_mapper: ProductMapper) -> dict[str, Any]:
        return {
            "delimiter": partition_mapper.delimiter,
            "mappers": [encode_partition_mapper(m) for m in partition_mapper.mappers],
        }

    @serialize_partition_mapper.register
    def _(self, partition_mapper: AllowedKeyMapper) -> dict[str, Any]:
        return {"allowed_keys": partition_mapper.allowed_keys}


_serializer = _Serializer()


@overload
def coerce_to_core_timetable(obj: BaseTimetable | CoreTimetable) -> CoreTimetable: ...


@overload
def coerce_to_core_timetable(obj: T) -> T: ...


def coerce_to_core_timetable(obj: object) -> object:
    """
    Convert *obj* from an SDK timetable to a Core tiemtable instance if possible.

    :meta private:
    """
    if isinstance(obj, CoreTimetable) or not isinstance(obj, BaseTimetable):
        return obj

    from airflow.serialization.decoders import decode_timetable

    return decode_timetable(encode_timetable(obj))


@overload
def ensure_serialized_asset(obj: Asset | SerializedAsset) -> SerializedAsset: ...


@overload
def ensure_serialized_asset(obj: AssetAlias | SerializedAssetAlias) -> SerializedAssetAlias: ...


@overload
def ensure_serialized_asset(obj: BaseAsset | SerializedAssetBase) -> SerializedAssetBase: ...


def ensure_serialized_asset(obj: BaseAsset | SerializedAssetBase) -> SerializedAssetBase:
    """
    Convert *obj* from an SDK asset to a Core asset instance if needed.

    :meta private:
    """
    if isinstance(obj, SerializedAssetBase):
        return obj

    from airflow.serialization.decoders import decode_asset_like

    return decode_asset_like(encode_asset_like(obj))


def ensure_serialized_deadline_alert(obj: DeadlineAlert | SerializedDeadlineAlert) -> SerializedDeadlineAlert:
    """
    Convert *obj* from an SDK deadline alert to a serialized deadline alert if needed.

    :meta private:
    """
    if isinstance(obj, SerializedDeadlineAlert):
        return obj
    return decode_deadline_alert(encode_deadline_alert(obj))


def encode_partition_mapper(var: PartitionMapper | CorePartitionMapper) -> dict[str, Any]:
    """
    Encode a PartitionMapper instance.

    This delegates most of the serialization work to the type, so the behavior
    can be completely controlled by a custom subclass.

    :meta private:
    """
    var_type = type(var)
    importable_string = _serializer.BUILTIN_PARTITION_MAPPERS.get(var_type)
    if importable_string is not None:
        return {
            Encoding.TYPE: importable_string,
            Encoding.VAR: _serializer.serialize_partition_mapper(var),
        }

    qn = qualname(var)
    if is_core_partition_mapper_import_path(qn) is False:
        # This raises if not found.
        find_registered_custom_partition_mapper(qn)

    return {
        Encoding.TYPE: qn,
        Encoding.VAR: _serializer.serialize_partition_mapper(var),
    }
