use super::*;

#[cfg_attr(feature = "ir_serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Clone, PartialEq, Debug, Eq, Hash)]
pub enum IRTemporalFunction {
    Millennium,
    Century,
    Year,
    IsLeapYear,
    IsoYear,
    Quarter,
    Month,
    DaysInMonth,
    Week,
    WeekDay,
    Day,
    OrdinalDay,
    Time,
    Date,
    Datetime,
    #[cfg(feature = "dtype-duration")]
    Duration(TimeUnit),
    Hour,
    Minute,
    Second,
    Millisecond,
    Microsecond,
    Nanosecond,
    #[cfg(feature = "dtype-duration")]
    TotalDays {
        fractional: bool,
    },
    #[cfg(feature = "dtype-duration")]
    TotalHours {
        fractional: bool,
    },
    #[cfg(feature = "dtype-duration")]
    TotalMinutes {
        fractional: bool,
    },
    #[cfg(feature = "dtype-duration")]
    TotalSeconds {
        fractional: bool,
    },
    #[cfg(feature = "dtype-duration")]
    TotalMilliseconds {
        fractional: bool,
    },
    #[cfg(feature = "dtype-duration")]
    TotalMicroseconds {
        fractional: bool,
    },
    #[cfg(feature = "dtype-duration")]
    TotalNanoseconds {
        fractional: bool,
    },
    ToString(String),
    CastTimeUnit(TimeUnit),
    WithTimeUnit(TimeUnit),
    #[cfg(feature = "timezones")]
    ConvertTimeZone(TimeZone),
    TimeStamp(TimeUnit),
    Truncate,
    #[cfg(feature = "offset_by")]
    OffsetBy,
    #[cfg(feature = "month_start")]
    MonthStart,
    #[cfg(feature = "month_end")]
    MonthEnd,
    #[cfg(feature = "timezones")]
    BaseUtcOffset,
    #[cfg(feature = "timezones")]
    DSTOffset,
    Round,
    Replace,
    #[cfg(feature = "timezones")]
    ReplaceTimeZone(Option<TimeZone>, NonExistent),
    Combine(TimeUnit),
    DatetimeFunction {
        time_unit: TimeUnit,
        time_zone: Option<TimeZone>,
    },
}

impl IRTemporalFunction {
    pub(super) fn get_field(&self, mapper: FieldsMapper) -> PolarsResult<Field> {
        use IRTemporalFunction::*;
        match self {
            Millennium | Century | Year | IsoYear => mapper.with_dtype(DataType::Int32),
            OrdinalDay => mapper.with_dtype(DataType::Int16),
            Month | DaysInMonth | Quarter | Week | WeekDay | Day | Hour | Minute | Second => {
                mapper.with_dtype(DataType::Int8)
            },
            Millisecond | Microsecond | Nanosecond => mapper.with_dtype(DataType::Int32),
            #[cfg(feature = "dtype-duration")]
            TotalDays { fractional }
            | TotalHours { fractional }
            | TotalMinutes { fractional }
            | TotalSeconds { fractional }
            | TotalMilliseconds { fractional }
            | TotalMicroseconds { fractional }
            | TotalNanoseconds { fractional } => {
                if *fractional {
                    mapper.with_dtype(DataType::Float64)
                } else {
                    mapper.with_dtype(DataType::Int64)
                }
            },
            ToString(_) => mapper.with_dtype(DataType::String),
            WithTimeUnit(tu) | CastTimeUnit(tu) => mapper.try_map_dtype(|dt| match dt {
                DataType::Duration(_) => Ok(DataType::Duration(*tu)),
                DataType::Datetime(_, tz) => Ok(DataType::Datetime(*tu, tz.clone())),
                dtype => polars_bail!(ComputeError: "expected duration or datetime, got {}", dtype),
            }),
            #[cfg(feature = "timezones")]
            ConvertTimeZone(tz) => mapper.try_map_dtype(|dt| match dt {
                DataType::Datetime(tu, _) => Ok(DataType::Datetime(*tu, Some(tz.clone()))),
                dtype => polars_bail!(ComputeError: "expected Datetime, got {}", dtype),
            }),
            TimeStamp(_) => mapper.with_dtype(DataType::Int64),
            IsLeapYear => mapper.with_dtype(DataType::Boolean),
            Time => mapper.with_dtype(DataType::Time),
            #[cfg(feature = "dtype-duration")]
            Duration(tu) => mapper.with_dtype(DataType::Duration(*tu)),
            Date => mapper.with_dtype(DataType::Date),
            Datetime => mapper.try_map_dtype(|dt| match dt {
                DataType::Datetime(tu, _) => Ok(DataType::Datetime(*tu, None)),
                dtype => polars_bail!(ComputeError: "expected Datetime, got {}", dtype),
            }),
            Truncate => mapper.with_same_dtype(),
            #[cfg(feature = "offset_by")]
            OffsetBy => mapper.with_same_dtype(),
            #[cfg(feature = "month_start")]
            MonthStart => mapper.with_same_dtype(),
            #[cfg(feature = "month_end")]
            MonthEnd => mapper.with_same_dtype(),
            #[cfg(feature = "timezones")]
            BaseUtcOffset => mapper.with_dtype(DataType::Duration(TimeUnit::Milliseconds)),
            #[cfg(feature = "timezones")]
            DSTOffset => mapper.with_dtype(DataType::Duration(TimeUnit::Milliseconds)),
            Round => mapper.with_same_dtype(),
            Replace => mapper.with_same_dtype(),
            #[cfg(feature = "timezones")]
            ReplaceTimeZone(tz, _non_existent) => mapper.map_datetime_dtype_timezone(tz.as_ref()),
            DatetimeFunction {
                time_unit,
                time_zone,
            } => Ok(Field::new(
                PlSmallStr::from_static("datetime"),
                DataType::Datetime(*time_unit, time_zone.clone()),
            )),
            Combine(tu) => mapper.try_map_dtype(|dt| match dt {
                DataType::Datetime(_, tz) => Ok(DataType::Datetime(*tu, tz.clone())),
                DataType::Date => Ok(DataType::Datetime(*tu, None)),
                dtype => {
                    polars_bail!(ComputeError: "expected Date or Datetime, got {}", dtype)
                },
            }),
        }
    }

    pub fn function_options(&self) -> FunctionOptions {
        use IRTemporalFunction as T;
        match self {
            T::Millennium
            | T::Century
            | T::Year
            | T::IsLeapYear
            | T::IsoYear
            | T::Quarter
            | T::Month
            | T::DaysInMonth
            | T::Week
            | T::WeekDay
            | T::Day
            | T::OrdinalDay
            | T::Time
            | T::Date
            | T::Datetime
            | T::Hour
            | T::Minute
            | T::Second
            | T::Millisecond
            | T::Microsecond
            | T::Nanosecond
            | T::ToString(_)
            | T::TimeStamp(_)
            | T::CastTimeUnit(_)
            | T::WithTimeUnit(_) => FunctionOptions::elementwise(),
            #[cfg(feature = "dtype-duration")]
            T::TotalDays { .. }
            | T::TotalHours { .. }
            | T::TotalMinutes { .. }
            | T::TotalSeconds { .. }
            | T::TotalMilliseconds { .. }
            | T::TotalMicroseconds { .. }
            | T::TotalNanoseconds { .. } => FunctionOptions::elementwise(),
            #[cfg(feature = "timezones")]
            T::ConvertTimeZone(_) => FunctionOptions::elementwise(),
            #[cfg(feature = "month_start")]
            T::MonthStart => FunctionOptions::elementwise(),
            #[cfg(feature = "month_end")]
            T::MonthEnd => FunctionOptions::elementwise(),
            #[cfg(feature = "timezones")]
            T::BaseUtcOffset | T::DSTOffset => FunctionOptions::elementwise(),
            T::Truncate => FunctionOptions::elementwise(),
            #[cfg(feature = "offset_by")]
            T::OffsetBy => FunctionOptions::elementwise(),
            T::Round => FunctionOptions::elementwise(),
            T::Replace => FunctionOptions::elementwise(),
            #[cfg(feature = "dtype-duration")]
            T::Duration(_) => FunctionOptions::elementwise(),
            #[cfg(feature = "timezones")]
            T::ReplaceTimeZone(_, _) => FunctionOptions::elementwise(),
            T::Combine(_) => FunctionOptions::elementwise(),
            T::DatetimeFunction { .. } => {
                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
            },
        }
    }
}

impl Display for IRTemporalFunction {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        use IRTemporalFunction::*;
        let s = match self {
            Millennium => "millennium",
            Century => "century",
            Year => "year",
            IsLeapYear => "is_leap_year",
            IsoYear => "iso_year",
            Quarter => "quarter",
            Month => "month",
            DaysInMonth => "days_in_month",
            Week => "week",
            WeekDay => "weekday",
            Day => "day",
            OrdinalDay => "ordinal_day",
            Time => "time",
            Date => "date",
            Datetime => "datetime",
            #[cfg(feature = "dtype-duration")]
            Duration(_) => "duration",
            Hour => "hour",
            Minute => "minute",
            Second => "second",
            Millisecond => "millisecond",
            Microsecond => "microsecond",
            Nanosecond => "nanosecond",
            #[cfg(feature = "dtype-duration")]
            TotalDays { .. } => "total_days",
            #[cfg(feature = "dtype-duration")]
            TotalHours { .. } => "total_hours",
            #[cfg(feature = "dtype-duration")]
            TotalMinutes { .. } => "total_minutes",
            #[cfg(feature = "dtype-duration")]
            TotalSeconds { .. } => "total_seconds",
            #[cfg(feature = "dtype-duration")]
            TotalMilliseconds { .. } => "total_milliseconds",
            #[cfg(feature = "dtype-duration")]
            TotalMicroseconds { .. } => "total_microseconds",
            #[cfg(feature = "dtype-duration")]
            TotalNanoseconds { .. } => "total_nanoseconds",
            ToString(_) => "to_string",
            #[cfg(feature = "timezones")]
            ConvertTimeZone(_) => "convert_time_zone",
            CastTimeUnit(_) => "cast_time_unit",
            WithTimeUnit(_) => "with_time_unit",
            TimeStamp(tu) => return write!(f, "dt.timestamp({tu})"),
            Truncate => "truncate",
            #[cfg(feature = "offset_by")]
            OffsetBy => "offset_by",
            #[cfg(feature = "month_start")]
            MonthStart => "month_start",
            #[cfg(feature = "month_end")]
            MonthEnd => "month_end",
            #[cfg(feature = "timezones")]
            BaseUtcOffset => "base_utc_offset",
            #[cfg(feature = "timezones")]
            DSTOffset => "dst_offset",
            Round => "round",
            Replace => "replace",
            #[cfg(feature = "timezones")]
            ReplaceTimeZone(_, _) => "replace_time_zone",
            DatetimeFunction { .. } => return write!(f, "dt.datetime"),
            Combine(_) => "combine",
        };
        write!(f, "dt.{s}")
    }
}
