hamelin_eval 0.7.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;
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) -> EvalResult<TimestampValue> {
    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, tz)
        }
        TimeZone::FixedOffset(offset) => {
            // Convert to fixed offset, truncate, convert back to UTC
            let ts_in_offset = ts.instant().with_timezone(offset);
            truncate_generic(&ts_in_offset, unit, 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.
fn truncate_generic<Tz: ChronoTimeZone>(
    ts: &DateTime<Tz>,
    unit: &TruncUnit,
    tz: &Tz,
) -> DateTime<Utc> {
    let truncated = match unit {
        TruncUnit::Second => {
            // Truncate by zeroing out nanoseconds
            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 => {
            // Truncate by zeroing out seconds (and nanoseconds)
            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 => {
            // Truncate by zeroing out minutes, seconds, and nanoseconds
            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 => {
            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 => {
            // Truncate to the beginning of the week (Monday)
            let days_since_monday = 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);
            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 result = tz.with_ymd_and_hms(ts.year(), ts.month(), 1, 0, 0, 0);
            result.earliest().unwrap_or_else(|| ts.clone())
        }

        TruncUnit::Quarter => {
            let quarter_month = match ts.month() {
                1..=3 => 1,    // Q1
                4..=6 => 4,    // Q2
                7..=9 => 7,    // Q3
                10..=12 => 10, // Q4
                _ => 1,
            };
            let result = tz.with_ymd_and_hms(ts.year(), quarter_month, 1, 0, 0, 0);
            result.earliest().unwrap_or_else(|| ts.clone())
        }

        TruncUnit::Year => {
            let result = tz.with_ymd_and_hms(ts.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,
) -> EvalResult<TimestampValue> {
    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, tz)
        }
        TimeZone::FixedOffset(offset) => {
            let ts_in_offset = truncated_ts.instant().with_timezone(offset);
            next_boundary_generic(&ts_in_offset, unit, offset)
        }
        TimeZone::Any => unreachable!(),
    };

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

/// Calculate the next boundary in any timezone (generic implementation)
///
/// Note: `truncated_ts` has already been converted to the target timezone `tz`.
/// For calendar-based units (Day, Week, Month, Quarter, Year), we use calendar arithmetic
/// to handle DST transitions correctly.
fn next_boundary_generic<Tz: ChronoTimeZone>(
    truncated_ts: &DateTime<Tz>,
    unit: &TruncUnit,
    tz: &Tz,
) -> DateTime<Utc> {
    let next = match unit {
        // Fixed-duration units: these are safe to use Duration
        TruncUnit::Second => truncated_ts.clone() + Duration::seconds(1),
        TruncUnit::Minute => truncated_ts.clone() + Duration::minutes(1),
        TruncUnit::Hour => truncated_ts.clone() + Duration::hours(1),

        // Calendar-based units: use calendar arithmetic to handle DST correctly
        TruncUnit::Day => {
            // Add 1 calendar day (not 24 hours, to handle DST)
            let next_date = truncated_ts.date_naive() + chrono::Days::new(1);
            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(1))
        }

        TruncUnit::Week => {
            // Add 7 calendar days (not 7*24 hours, to handle DST)
            let next_date = truncated_ts.date_naive() + chrono::Days::new(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(1))
        }

        TruncUnit::Month => {
            let year = truncated_ts.year();
            let month = truncated_ts.month();
            let (next_year, next_month) = if month == 12 {
                (year + 1, 1)
            } else {
                (year, month + 1)
            };
            tz.with_ymd_and_hms(next_year, next_month, 1, 0, 0, 0)
                .earliest()
                .unwrap_or_else(|| truncated_ts.clone() + Duration::days(30))
        }

        TruncUnit::Quarter => {
            let year = truncated_ts.year();
            let month = truncated_ts.month();
            let (next_year, next_month) = match month {
                1..=3 => (year, 4),       // Q1 -> Q2
                4..=6 => (year, 7),       // Q2 -> Q3
                7..=9 => (year, 10),      // Q3 -> Q4
                10..=12 => (year + 1, 1), // Q4 -> Q1 next year
                _ => (year, month + 3),
            };
            tz.with_ymd_and_hms(next_year, next_month, 1, 0, 0, 0)
                .earliest()
                .unwrap_or_else(|| truncated_ts.clone() + Duration::days(90))
        }

        TruncUnit::Year => tz
            .with_ymd_and_hms(truncated_ts.year() + 1, 1, 1, 0, 0, 0)
            .earliest()
            .unwrap_or_else(|| truncated_ts.clone() + Duration::days(365)),
    };

    // 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).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).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).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).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).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).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)
                .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)
                .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)
                .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)
                .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).unwrap();
        assert_eq!(truncated.instant(), &expected);
    }

    #[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).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).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).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)
                .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)
                .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).unwrap();
        assert_eq!(next.instant(), &expected);
    }
}