hamelin_eval 0.10.13

//! Timestamp truncation operations for forward and reverse evaluation
//!
//! This module provides helpers for truncating timestamps to various time units
//! and reversing those truncations for constraint solving.

use chrono::{DateTime, Datelike, Duration, TimeZone as ChronoTimeZone, Timelike, Utc, Weekday};

use crate::eval::error::EvalResult;

/// Unix epoch date, used as the anchor point for multi-day truncation
/// (local-day-count modular arithmetic aligned to local midnight).
const EPOCH_DATE: chrono::NaiveDate = match chrono::NaiveDate::from_ymd_opt(1970, 1, 1) {
    Some(d) => d,
    None => panic!("1970-01-01 is a valid date"),
};

/// First Monday of the Unix era, used as the anchor point for multi-week truncation
/// to match the SQL/DataFusion translation.
const REFERENCE_MONDAY: chrono::NaiveDate = match chrono::NaiveDate::from_ymd_opt(1970, 1, 5) {
    Some(d) => d,
    None => panic!("1970-01-05 is a valid date"),
};
use crate::value::{TimeZone, TimestampValue};
use hamelin_lib::tree::ast::expression::TruncUnit;

/// Forward truncation: truncate a timestamp to the specified unit
///
/// The truncation is performed in the timestamp's timezone. For example,
/// truncating to @day in US/Pacific will truncate to midnight Pacific time,
/// not midnight UTC.
///
/// Returns an error if the timestamp has TimeZone::Any (invalid for forward evaluation).
///
/// ## DST Handling
///
/// During DST transitions, some times don't exist (spring forward gap) or occur twice
/// (fall back overlap). This function uses `.earliest()` for consistent behavior:
///
/// - **Gap (spring forward)**: If the truncated time doesn't exist (e.g., 2:00 AM on
///   spring-forward day), uses the first valid time after the gap (e.g., 3:00 AM).
/// - **Overlap (fall back)**: If the truncated time occurs twice (e.g., 1:00 AM on
///   fall-back day), uses the first occurrence (before clocks fall back).
///
/// This ensures truncation is deterministic and always succeeds, even during DST transitions.
pub fn truncate_timestamp(
    ts: &TimestampValue,
    unit: &TruncUnit,
    multiplier: u32,
) -> EvalResult<TimestampValue> {
    if multiplier == 0 {
        return Err(crate::eval::error::EvalError::execution(
            "Truncation multiplier must be at least 1".to_string(),
        ));
    }

    if ts.timezone().is_any() {
        return Err(crate::eval::error::EvalError::execution(
            "Cannot truncate timestamp with unconstrained timezone (TimeZone::Any)".to_string(),
        ));
    }

    // Truncation depends on the timezone representation, so we need to work in that timezone
    let truncated_instant = match ts.timezone() {
        TimeZone::Named(tz) => {
            // Convert to named timezone, truncate, convert back to UTC
            let ts_in_tz = ts.instant().with_timezone(tz);
            truncate_generic(&ts_in_tz, unit, multiplier, tz)
        }
        TimeZone::FixedOffset(offset) => {
            let ts_in_offset = ts.instant().with_timezone(offset);
            truncate_generic(&ts_in_offset, unit, multiplier, offset)
        }
        TimeZone::Any => unreachable!(), // Already checked above
    };

    Ok(TimestampValue::new(
        truncated_instant,
        ts.timezone().clone(),
    ))
}

/// Truncate a timestamp in any timezone (generic implementation).
///
/// Note: `ts` has already been converted to the target timezone `tz`, so we can use
/// its year(), month(), day(), hour(), minute(), second() methods to get components
/// in the target timezone.
///
/// When `multiplier > 1` for sub-day units (s/m/h), epoch-modular arithmetic is used:
/// `floor(epoch_s / (N * unit_s)) * (N * unit_s)`, producing continuous buckets that
/// roll across parent-unit boundaries (matching Splunk's behavior).
///
/// For multi-day truncation, local-day-count arithmetic is used instead:
/// `floor(local_day_since_epoch / N) * N`, aligned to local midnight (matching Splunk).
fn truncate_generic<Tz: ChronoTimeZone>(
    ts: &DateTime<Tz>,
    unit: &TruncUnit,
    multiplier: u32,
    tz: &Tz,
) -> DateTime<Utc> {
    let truncated = match unit {
        TruncUnit::Second => {
            if multiplier > 1 {
                let step = multiplier as i64;
                let epoch_secs = ts.timestamp();
                let snapped = epoch_secs.div_euclid(step) * step;
                return DateTime::from_timestamp(snapped, 0)
                    .unwrap_or_else(|| ts.with_timezone(&Utc));
            }
            let result = tz.with_ymd_and_hms(
                ts.year(),
                ts.month(),
                ts.day(),
                ts.hour(),
                ts.minute(),
                ts.second(),
            );
            result.earliest().unwrap_or_else(|| ts.clone())
        }

        TruncUnit::Minute => {
            if multiplier > 1 {
                let step = multiplier as i64 * 60;
                let epoch_secs = ts.timestamp();
                let snapped = epoch_secs.div_euclid(step) * step;
                return DateTime::from_timestamp(snapped, 0)
                    .unwrap_or_else(|| ts.with_timezone(&Utc));
            }
            let result =
                tz.with_ymd_and_hms(ts.year(), ts.month(), ts.day(), ts.hour(), ts.minute(), 0);
            result.earliest().unwrap_or_else(|| ts.clone())
        }

        TruncUnit::Hour => {
            if multiplier > 1 {
                let step = multiplier as i64 * 3600;
                let epoch_secs = ts.timestamp();
                let snapped = epoch_secs.div_euclid(step) * step;
                return DateTime::from_timestamp(snapped, 0)
                    .unwrap_or_else(|| ts.with_timezone(&Utc));
            }
            let result = tz.with_ymd_and_hms(ts.year(), ts.month(), ts.day(), ts.hour(), 0, 0);
            result.earliest().unwrap_or_else(|| ts.clone())
        }

        TruncUnit::Day => {
            if multiplier > 1 {
                let local_date = ts.date_naive();
                let local_day = local_date.signed_duration_since(EPOCH_DATE).num_days();
                let snapped_day = local_day.div_euclid(multiplier as i64) * multiplier as i64;
                let snapped_date = EPOCH_DATE + Duration::days(snapped_day);
                snapped_date
                    .and_hms_opt(0, 0, 0)
                    .and_then(|naive| tz.from_local_datetime(&naive).earliest())
                    .unwrap_or_else(|| ts.clone())
            } else {
                let date = ts.date_naive();
                date.and_hms_opt(0, 0, 0)
                    .and_then(|naive| tz.from_local_datetime(&naive).earliest())
                    .unwrap_or_else(|| ts.clone())
            }
        }

        TruncUnit::Week => {
            let days_since_monday: i64 = match ts.weekday() {
                Weekday::Mon => 0,
                Weekday::Tue => 1,
                Weekday::Wed => 2,
                Weekday::Thu => 3,
                Weekday::Fri => 4,
                Weekday::Sat => 5,
                Weekday::Sun => 6,
            };
            let ts_monday = ts.clone() - Duration::days(days_since_monday);
            if multiplier > 1 {
                // Snap to N-week boundaries relative to Monday 1970-01-05 (first Monday
                // of Unix era), matching the SQL/DataFusion translation anchor point.
                let reference_monday = REFERENCE_MONDAY;
                let monday_date = ts_monday.date_naive();
                let days_from_ref = monday_date
                    .signed_duration_since(reference_monday)
                    .num_days();
                let week_days = multiplier as i64 * 7;
                let snapped_offset = days_from_ref.div_euclid(week_days) * week_days;
                let snapped_date = reference_monday + Duration::days(snapped_offset);
                snapped_date
                    .and_hms_opt(0, 0, 0)
                    .and_then(|naive| tz.from_local_datetime(&naive).earliest())
                    .unwrap_or_else(|| ts.clone())
            } else {
                let date = ts_monday.date_naive();
                date.and_hms_opt(0, 0, 0)
                    .and_then(|naive| tz.from_local_datetime(&naive).earliest())
                    .unwrap_or_else(|| ts.clone())
            }
        }

        TruncUnit::Month => {
            let snapped_month = ((ts.month() - 1) / multiplier) * multiplier + 1;
            let result = tz.with_ymd_and_hms(ts.year(), snapped_month, 1, 0, 0, 0);
            result.earliest().unwrap_or_else(|| ts.clone())
        }

        TruncUnit::Quarter => {
            let quarter_months = multiplier * 3;
            let snapped_month = ((ts.month() - 1) / quarter_months) * quarter_months + 1;
            let result = tz.with_ymd_and_hms(ts.year(), snapped_month, 1, 0, 0, 0);
            result.earliest().unwrap_or_else(|| ts.clone())
        }

        TruncUnit::Year => {
            let m = multiplier as i64;
            let snapped_year = ((ts.year() as i64).div_euclid(m) * m) as i32;
            let result = tz.with_ymd_and_hms(snapped_year, 1, 1, 0, 0, 0);
            result.earliest().unwrap_or_else(|| ts.clone())
        }
    };

    // Convert back to UTC
    truncated.with_timezone(&Utc)
}

/// Calculate the next boundary after truncation
///
/// Given a truncated timestamp, returns the timestamp that represents the
/// start of the next period. The boundary is calculated in the timestamp's timezone.
pub fn next_truncation_boundary(
    truncated_ts: &TimestampValue,
    unit: &TruncUnit,
    multiplier: u32,
) -> EvalResult<TimestampValue> {
    if multiplier == 0 {
        return Err(crate::eval::error::EvalError::execution(
            "Truncation multiplier must be at least 1".to_string(),
        ));
    }

    if truncated_ts.timezone().is_any() {
        return Err(crate::eval::error::EvalError::execution(
            "Cannot calculate next boundary for timestamp with unconstrained timezone (TimeZone::Any)".to_string(),
        ));
    }

    let next_instant = match truncated_ts.timezone() {
        TimeZone::Named(tz) => {
            let ts_in_tz = truncated_ts.instant().with_timezone(tz);
            next_boundary_generic(&ts_in_tz, unit, multiplier, tz)
        }
        TimeZone::FixedOffset(offset) => {
            let ts_in_offset = truncated_ts.instant().with_timezone(offset);
            next_boundary_generic(&ts_in_offset, unit, multiplier, offset)
        }
        TimeZone::Any => unreachable!(),
    };

    Ok(TimestampValue::new(
        next_instant,
        truncated_ts.timezone().clone(),
    ))
}

/// Calculate the next boundary in any timezone (generic implementation).
///
/// Steps forward by `multiplier` units from the current truncated boundary.
///
/// For [`TruncUnit::Month`] / [`TruncUnit::Quarter`], this matches the same period grid as
/// [`truncate_generic`]: when the step crosses past December, the next boundary is January of
/// the following year (not `(month + step) % 12`), so short final buckets roll into the next year.
fn next_boundary_generic<Tz: ChronoTimeZone>(
    truncated_ts: &DateTime<Tz>,
    unit: &TruncUnit,
    multiplier: u32,
    tz: &Tz,
) -> DateTime<Utc> {
    let m = multiplier as i64;

    let next = match unit {
        TruncUnit::Second => truncated_ts.clone() + Duration::seconds(m),
        TruncUnit::Minute => truncated_ts.clone() + Duration::minutes(m),
        TruncUnit::Hour => truncated_ts.clone() + Duration::hours(m),

        TruncUnit::Day => {
            let next_date = truncated_ts.date_naive() + chrono::Days::new(multiplier as u64);
            next_date
                .and_hms_opt(0, 0, 0)
                .and_then(|naive| tz.from_local_datetime(&naive).earliest())
                .unwrap_or_else(|| truncated_ts.clone() + Duration::days(m))
        }

        TruncUnit::Week => {
            let next_date = truncated_ts.date_naive() + chrono::Days::new(multiplier as u64 * 7);
            next_date
                .and_hms_opt(0, 0, 0)
                .and_then(|naive| tz.from_local_datetime(&naive).earliest())
                .unwrap_or_else(|| truncated_ts.clone() + Duration::weeks(m))
        }

        TruncUnit::Month => {
            let next_month_candidate = truncated_ts.month() + multiplier;
            let (target_year, target_month) = if next_month_candidate <= 12 {
                (truncated_ts.year(), next_month_candidate)
            } else {
                (truncated_ts.year() + 1, 1)
            };
            tz.with_ymd_and_hms(target_year, target_month, 1, 0, 0, 0)
                .earliest()
                .unwrap_or_else(|| truncated_ts.clone() + Duration::days(30 * m))
        }

        TruncUnit::Quarter => {
            let step_months = multiplier * 3;
            let next_month_candidate = truncated_ts.month() + step_months;
            let (target_year, target_month) = if next_month_candidate <= 12 {
                (truncated_ts.year(), next_month_candidate)
            } else {
                (truncated_ts.year() + 1, 1)
            };
            tz.with_ymd_and_hms(target_year, target_month, 1, 0, 0, 0)
                .earliest()
                .unwrap_or_else(|| truncated_ts.clone() + Duration::days(90 * m))
        }

        TruncUnit::Year => {
            let next_year = truncated_ts.year() + multiplier as i32;
            tz.with_ymd_and_hms(next_year, 1, 1, 0, 0, 0)
                .earliest()
                .unwrap_or_else(|| truncated_ts.clone() + Duration::days(365 * m))
        }
    };

    // Convert back to UTC
    next.with_timezone(&Utc)
}

#[cfg(test)]
mod tests {
    use super::*;
    use chrono::TimeZone as ChronoTimeZoneTrait;
    use chrono::Timelike;

    #[test]
    fn test_truncate_to_second() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        let ts_with_nanos = ts.with_nanosecond(123456789).unwrap();
        let ts_value = TimestampValue::utc(ts_with_nanos);

        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Second, 1).unwrap();
        assert_eq!(truncated.instant(), &ts);
        assert!(truncated.timezone().is_utc());
    }

    #[test]
    fn test_truncate_to_minute() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Minute, 1).unwrap();
        assert_eq!(truncated.instant(), &expected);
    }

    #[test]
    fn test_truncate_to_hour() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 3, 15, 14, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Hour, 1).unwrap();
        assert_eq!(truncated.instant(), &expected);
    }

    #[test]
    fn test_truncate_to_day() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 3, 15, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Day, 1).unwrap();
        assert_eq!(truncated.instant(), &expected);
    }

    #[test]
    fn test_truncate_to_week() {
        // Friday March 15, 2024
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        // Should truncate to Monday March 11, 2024
        let expected = Utc.with_ymd_and_hms(2024, 3, 11, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Week, 1).unwrap();
        assert_eq!(truncated.instant(), &expected);
    }

    #[test]
    fn test_truncate_to_month() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Month, 1).unwrap();
        assert_eq!(truncated.instant(), &expected);
    }

    #[test]
    fn test_truncate_to_quarter() {
        // Q1 test
        let ts_q1 = Utc.with_ymd_and_hms(2024, 2, 15, 14, 30, 45).unwrap();
        let expected_q1 = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        let ts_value_q1 = TimestampValue::utc(ts_q1);
        assert_eq!(
            truncate_timestamp(&ts_value_q1, &TruncUnit::Quarter, 1)
                .unwrap()
                .instant(),
            &expected_q1
        );

        // Q2 test
        let ts_q2 = Utc.with_ymd_and_hms(2024, 5, 15, 14, 30, 45).unwrap();
        let expected_q2 = Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap();
        let ts_value_q2 = TimestampValue::utc(ts_q2);
        assert_eq!(
            truncate_timestamp(&ts_value_q2, &TruncUnit::Quarter, 1)
                .unwrap()
                .instant(),
            &expected_q2
        );

        // Q3 test
        let ts_q3 = Utc.with_ymd_and_hms(2024, 8, 15, 14, 30, 45).unwrap();
        let expected_q3 = Utc.with_ymd_and_hms(2024, 7, 1, 0, 0, 0).unwrap();
        let ts_value_q3 = TimestampValue::utc(ts_q3);
        assert_eq!(
            truncate_timestamp(&ts_value_q3, &TruncUnit::Quarter, 1)
                .unwrap()
                .instant(),
            &expected_q3
        );

        // Q4 test
        let ts_q4 = Utc.with_ymd_and_hms(2024, 11, 15, 14, 30, 45).unwrap();
        let expected_q4 = Utc.with_ymd_and_hms(2024, 10, 1, 0, 0, 0).unwrap();
        let ts_value_q4 = TimestampValue::utc(ts_q4);
        assert_eq!(
            truncate_timestamp(&ts_value_q4, &TruncUnit::Quarter, 1)
                .unwrap()
                .instant(),
            &expected_q4
        );
    }

    #[test]
    fn test_truncate_to_year() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Year, 1).unwrap();
        assert_eq!(truncated.instant(), &expected);
    }

    #[test]
    fn test_truncate_with_multiplier_2h() {
        let ts_value = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 15, 15, 30, 0).unwrap());
        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Hour, 2).unwrap();
        assert_eq!(
            truncated.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 15, 14, 0, 0).unwrap()
        );

        let ts_value2 = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 15, 13, 45, 0).unwrap());
        let truncated2 = truncate_timestamp(&ts_value2, &TruncUnit::Hour, 2).unwrap();
        assert_eq!(
            truncated2.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 15, 12, 0, 0).unwrap()
        );
    }

    #[test]
    fn test_truncate_with_multiplier_15m() {
        let ts_value = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 15, 14, 37, 0).unwrap());
        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Minute, 15).unwrap();
        assert_eq!(
            truncated.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 0).unwrap()
        );
    }

    #[test]
    fn test_truncate_with_multiplier_30s() {
        let ts_value = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap());
        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Second, 30).unwrap();
        assert_eq!(
            truncated.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 30).unwrap()
        );
    }

    #[test]
    fn test_truncate_7h_crosses_midnight() {
        // Epoch-modular: 7h grid boundaries are continuous, not reset at midnight.
        // On 2024-03-15, the 7h grid lands at 04:00, 11:00, 18:00 (from epoch).
        let ts_value = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 15, 4, 30, 0).unwrap());
        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Hour, 7).unwrap();
        assert_eq!(
            truncated.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 15, 4, 0, 0).unwrap()
        );

        // 01:00 Mar 15 snaps back to 21:00 Mar 14 (crosses midnight)
        let ts_value2 = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 15, 1, 0, 0).unwrap());
        let truncated2 = truncate_timestamp(&ts_value2, &TruncUnit::Hour, 7).unwrap();
        assert_eq!(
            truncated2.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 14, 21, 0, 0).unwrap()
        );
    }

    #[test]
    fn test_truncate_3d_crosses_month() {
        // Epoch-modular: 3-day grid is continuous from epoch, crossing month boundaries.
        // The grid near March 2024: ..., Feb 26, Feb 29, Mar 3, Mar 6, ...
        // Mar 1 snaps to Feb 29 (crosses month boundary).
        let ts_value = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 1, 12, 0, 0).unwrap());
        let truncated = truncate_timestamp(&ts_value, &TruncUnit::Day, 3).unwrap();
        assert_eq!(
            truncated.instant(),
            &Utc.with_ymd_and_hms(2024, 2, 29, 0, 0, 0).unwrap()
        );

        // Mar 3 starts a new 3-day bucket
        let ts_value2 = TimestampValue::utc(Utc.with_ymd_and_hms(2024, 3, 3, 6, 0, 0).unwrap());
        let truncated2 = truncate_timestamp(&ts_value2, &TruncUnit::Day, 3).unwrap();
        assert_eq!(
            truncated2.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 3, 0, 0, 0).unwrap()
        );
    }

    #[test]
    fn test_next_boundary_7h() {
        // Next boundary after a 7h truncation point at 04:00 → 11:00
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 4, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Hour, 7).unwrap();
        assert_eq!(
            next.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 15, 11, 0, 0).unwrap()
        );
    }

    #[test]
    fn test_next_boundary_3d() {
        // Next boundary after a 3-day truncation point: Feb 29 + 3d = Mar 3
        let ts = Utc.with_ymd_and_hms(2024, 2, 29, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Day, 3).unwrap();
        assert_eq!(
            next.instant(),
            &Utc.with_ymd_and_hms(2024, 3, 3, 0, 0, 0).unwrap()
        );
    }

    #[test]
    fn test_next_boundary_second() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 45).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 3, 15, 14, 30, 46).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Second, 1).unwrap();
        assert_eq!(next.instant(), &expected);
    }

    #[test]
    fn test_next_boundary_month() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 1, 0, 0, 0).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Month, 1).unwrap();
        assert_eq!(next.instant(), &expected);

        // Test year boundary
        let ts_dec = Utc.with_ymd_and_hms(2024, 12, 1, 0, 0, 0).unwrap();
        let expected_jan = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let ts_value_dec = TimestampValue::utc(ts_dec);

        let next_jan = next_truncation_boundary(&ts_value_dec, &TruncUnit::Month, 1).unwrap();
        assert_eq!(next_jan.instant(), &expected_jan);
    }

    #[test]
    fn test_next_boundary_quarter() {
        // Q1 -> Q2
        let ts_q1 = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        let expected_q2 = Utc.with_ymd_and_hms(2024, 4, 1, 0, 0, 0).unwrap();
        let ts_value_q1 = TimestampValue::utc(ts_q1);
        assert_eq!(
            next_truncation_boundary(&ts_value_q1, &TruncUnit::Quarter, 1)
                .unwrap()
                .instant(),
            &expected_q2
        );

        // Q4 -> Q1 next year
        let ts_q4 = Utc.with_ymd_and_hms(2024, 10, 1, 0, 0, 0).unwrap();
        let expected_q1_next = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let ts_value_q4 = TimestampValue::utc(ts_q4);
        assert_eq!(
            next_truncation_boundary(&ts_value_q4, &TruncUnit::Quarter, 1)
                .unwrap()
                .instant(),
            &expected_q1_next
        );
    }

    #[test]
    fn test_next_boundary_year() {
        let ts = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        let expected = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Year, 1).unwrap();
        assert_eq!(next.instant(), &expected);
    }

    #[test]
    fn test_next_boundary_with_multiplier_2h() {
        let ts = Utc.with_ymd_and_hms(2024, 3, 15, 14, 0, 0).unwrap();
        let expected = Utc.with_ymd_and_hms(2024, 3, 15, 16, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(ts);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Hour, 2).unwrap();
        assert_eq!(next.instant(), &expected);
    }

    #[test]
    fn test_next_boundary_5mon_short_last_bucket() {
        // @5mon boundaries: Jan, Jun, Nov. The Nov→Jan bucket is only 2 months.
        let nov = Utc.with_ymd_and_hms(2024, 11, 1, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(nov);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Month, 5).unwrap();
        assert_eq!(
            next.instant(),
            &Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap()
        );

        // Mid-year boundary should still advance by the full 5 months
        let jan = Utc.with_ymd_and_hms(2024, 1, 1, 0, 0, 0).unwrap();
        let ts_jan = TimestampValue::utc(jan);
        let next_jan = next_truncation_boundary(&ts_jan, &TruncUnit::Month, 5).unwrap();
        assert_eq!(
            next_jan.instant(),
            &Utc.with_ymd_and_hms(2024, 6, 1, 0, 0, 0).unwrap()
        );
    }

    #[test]
    fn test_next_boundary_3q_short_last_bucket() {
        // @3q = @9mon boundaries: Jan, Oct. The Oct→Jan bucket is only 3 months.
        let oct = Utc.with_ymd_and_hms(2024, 10, 1, 0, 0, 0).unwrap();
        let ts_value = TimestampValue::utc(oct);

        let next = next_truncation_boundary(&ts_value, &TruncUnit::Quarter, 3).unwrap();
        assert_eq!(
            next.instant(),
            &Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap()
        );
    }
}