hegeltest 0.27.0

use crate::generators::{DefaultGenerator, Generator, TestCase, hashsets, integers};
use crate::test_case::invalid_argument;
use chrono::{
    DateTime, Datelike, FixedOffset, NaiveDate, NaiveDateTime, NaiveTime, NaiveWeek, TimeDelta,
    TimeZone, Timelike, Utc, Weekday, WeekdaySet,
};
use std::marker::PhantomData;

/// Convert a [`NaiveTime`] to a total nanosecond count from midnight.
///
/// chrono encodes leap seconds by letting `nanosecond()` reach into
/// `[1_000_000_000, 2_000_000_000)`; those bits are preserved here so the
/// conversion round-trips through [`total_nanos_to_time`] for the
/// end-of-day leap second. A leap second on any *other* second maps into
/// the same totals as ordinary times on the following second, so it is not
/// order-consistent as a bound — [`NaiveTimeGenerator`] rejects such bounds
/// (see [`is_mid_day_leap`]).
fn time_to_total_nanos(t: NaiveTime) -> i64 {
    let secs = i64::from(t.num_seconds_from_midnight());
    let nanos = i64::from(t.nanosecond());
    secs * 1_000_000_000 + nanos
}

/// Whether `t` uses chrono's leap-second representation on a second other
/// than 23:59:59 (the only second whose leap encoding stays
/// order-consistent under [`time_to_total_nanos`]).
fn is_mid_day_leap(t: NaiveTime) -> bool {
    t.nanosecond() >= 1_000_000_000 && t.num_seconds_from_midnight() != 86_399
}

/// Inverse of [`time_to_total_nanos`].
fn total_nanos_to_time(total: i64) -> NaiveTime {
    let (secs, nanos) = if total >= 86_400 * 1_000_000_000 {
        (86_399, (total - 86_399 * 1_000_000_000) as u32)
    } else {
        (
            (total / 1_000_000_000) as u32,
            (total % 1_000_000_000) as u32,
        )
    };
    NaiveTime::from_num_seconds_from_midnight_opt(secs, nanos).unwrap()
}

/// Default upper bound for [`naive_times`]. Excludes leap-second representation
/// — users who want to generate leap seconds must opt in via [`NaiveTimeGenerator::max_value`].
fn naive_time_default_max() -> NaiveTime {
    NaiveTime::from_hms_nano_opt(23, 59, 59, 999_999_999).unwrap()
}

/// Generator for [`chrono::WeekdaySet`] values. Created by [`weekday_sets()`].
pub struct WeekdaySetGenerator;

impl Generator<WeekdaySet> for WeekdaySetGenerator {
    fn do_draw(&self, tc: &TestCase) -> WeekdaySet {
        let days = hashsets(integers::<u8>().min_value(0).max_value(6))
            .max_size(7)
            .do_draw(tc);
        let mut set = WeekdaySet::EMPTY;
        for n in days {
            set.insert(Weekday::try_from(n).unwrap());
        }
        set
    }
}

/// Generate [`chrono::WeekdaySet`] values.
///
/// # Example
///
/// ```no_run
/// use hegel::extras::chrono as chrono_gs;
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     let s: chrono::WeekdaySet = tc.draw(chrono_gs::weekday_sets());
///     assert!(s.len() <= 7);
/// }
/// ```
pub fn weekday_sets() -> WeekdaySetGenerator {
    WeekdaySetGenerator
}

/// Generator for [`chrono::FixedOffset`] values. Created by [`fixed_offsets()`].
pub struct FixedOffsetGenerator {
    min_value: FixedOffset,
    max_value: FixedOffset,
}

impl FixedOffsetGenerator {
    /// Set the minimum offset (inclusive).
    pub fn min_value(mut self, min: FixedOffset) -> Self {
        self.min_value = min;
        self
    }

    /// Set the maximum offset (inclusive).
    pub fn max_value(mut self, max: FixedOffset) -> Self {
        self.max_value = max;
        self
    }
}

impl Generator<FixedOffset> for FixedOffsetGenerator {
    fn do_draw(&self, tc: &TestCase) -> FixedOffset {
        let min_secs = self.min_value.local_minus_utc();
        let max_secs = self.max_value.local_minus_utc();
        if min_secs > max_secs {
            invalid_argument!("Cannot have max_value < min_value");
        }
        let secs = integers::<i32>()
            .min_value(min_secs)
            .max_value(max_secs)
            .do_draw(tc);
        FixedOffset::east_opt(secs).unwrap()
    }
}

/// Generate [`chrono::FixedOffset`] values.
///
/// # Example
///
/// ```no_run
/// use chrono::FixedOffset;
/// use hegel::extras::chrono as chrono_gs;
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     let max = FixedOffset::east_opt(12 * 3600).unwrap();
///     let off = tc.draw(chrono_gs::fixed_offsets().max_value(max));
///     assert!(off.local_minus_utc() <= 12 * 3600);
/// }
/// ```
pub fn fixed_offsets() -> FixedOffsetGenerator {
    FixedOffsetGenerator {
        min_value: FixedOffset::west_opt(86_399).unwrap(),
        max_value: FixedOffset::east_opt(86_399).unwrap(),
    }
}

/// Convert a [`TimeDelta`] to a total nanosecond count.
///
/// `TimeDelta` is internally `(secs: i64, nanos: u32 in [0, 1e9))`. The total
/// nanosecond magnitude exceeds `i64`'s range by a factor of ~10^9, so we
/// widen to `i128` (which has ~10^10 of headroom past `TimeDelta::MAX`).
fn timedelta_to_nanos(d: TimeDelta) -> i128 {
    i128::from(d.num_seconds()) * 1_000_000_000 + i128::from(d.subsec_nanos())
}

/// Inverse of [`timedelta_to_nanos`].
fn nanos_to_timedelta(n: i128) -> TimeDelta {
    let secs = n.div_euclid(1_000_000_000) as i64;
    let nanos = n.rem_euclid(1_000_000_000) as u32;
    TimeDelta::new(secs, nanos).unwrap()
}

/// Generator for [`chrono::TimeDelta`] values. Created by [`time_deltas()`].
pub struct TimeDeltaGenerator {
    min_value: TimeDelta,
    max_value: TimeDelta,
}

impl TimeDeltaGenerator {
    /// Set the minimum delta (inclusive).
    pub fn min_value(mut self, min: TimeDelta) -> Self {
        self.min_value = min;
        self
    }

    /// Set the maximum delta (inclusive).
    pub fn max_value(mut self, max: TimeDelta) -> Self {
        self.max_value = max;
        self
    }
}

impl Generator<TimeDelta> for TimeDeltaGenerator {
    fn do_draw(&self, tc: &TestCase) -> TimeDelta {
        if self.min_value > self.max_value {
            invalid_argument!("Cannot have max_value < min_value");
        }
        let nanos = integers::<i128>()
            .min_value(timedelta_to_nanos(self.min_value))
            .max_value(timedelta_to_nanos(self.max_value))
            .do_draw(tc);
        nanos_to_timedelta(nanos)
    }
}

/// Generate [`chrono::TimeDelta`] values.
///
/// Defaults span the full `TimeDelta::MIN..=TimeDelta::MAX` range. Use the
/// builder methods to constrain.
///
/// # Example
///
/// ```no_run
/// use chrono::TimeDelta;
/// use hegel::extras::chrono as chrono_gs;
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     let max = TimeDelta::seconds(60);
///     let d = tc.draw(chrono_gs::time_deltas().min_value(TimeDelta::zero()).max_value(max));
///     assert!(d >= TimeDelta::zero() && d <= max);
/// }
/// ```
pub fn time_deltas() -> TimeDeltaGenerator {
    TimeDeltaGenerator {
        min_value: TimeDelta::MIN,
        max_value: TimeDelta::MAX,
    }
}

/// Generator for [`chrono::NaiveDate`] values. Created by [`naive_dates()`].
///
/// Internally the date is generated as a count of days from the Common Era
/// epoch (`NaiveDate::from_num_days_from_ce_opt`).
pub struct NaiveDateGenerator {
    min_value: NaiveDate,
    max_value: NaiveDate,
}

impl NaiveDateGenerator {
    /// Set the minimum date (inclusive).
    pub fn min_value(mut self, min: NaiveDate) -> Self {
        self.min_value = min;
        self
    }

    /// Set the maximum date (inclusive).
    pub fn max_value(mut self, max: NaiveDate) -> Self {
        self.max_value = max;
        self
    }
}

impl Generator<NaiveDate> for NaiveDateGenerator {
    fn do_draw(&self, tc: &TestCase) -> NaiveDate {
        if self.min_value > self.max_value {
            invalid_argument!("Cannot have max_value < min_value");
        }
        let d = tc.generate_date(hegel_date(self.min_value), hegel_date(self.max_value));
        NaiveDate::from_ymd_opt(d.year, u32::from(d.month), u32::from(d.day)).unwrap()
    }
}

/// Convert a [`NaiveDate`] to the engine's date struct. Every `NaiveDate`
/// fits: chrono spans years ±262143 and the engine accepts ±999999.
fn hegel_date(d: NaiveDate) -> hegel_c::hegel_date_t {
    hegel_c::hegel_date_t {
        year: d.year(),
        month: d.month() as u8,
        day: d.day() as u8,
    }
}

/// Generate [`chrono::NaiveDate`] values.
///
/// # Example
///
/// ```no_run
/// use chrono::NaiveDate;
/// use hegel::extras::chrono as chrono_gs;
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     let min = NaiveDate::from_ymd_opt(2024, 1, 1).unwrap();
///     let max = NaiveDate::from_ymd_opt(2024, 12, 31).unwrap();
///     let d = tc.draw(chrono_gs::naive_dates().min_value(min).max_value(max));
///     assert_eq!(d.iso_week().year(), 2024);
/// }
/// ```
pub fn naive_dates() -> NaiveDateGenerator {
    NaiveDateGenerator {
        min_value: NaiveDate::MIN,
        max_value: NaiveDate::MAX,
    }
}

/// Generator for [`chrono::NaiveTime`] values. Created by [`naive_times()`].
pub struct NaiveTimeGenerator {
    min_value: NaiveTime,
    max_value: NaiveTime,
}

impl NaiveTimeGenerator {
    /// Set the minimum time (inclusive).
    pub fn min_value(mut self, min: NaiveTime) -> Self {
        self.min_value = min;
        self
    }

    /// Set the maximum time (inclusive).
    pub fn max_value(mut self, max: NaiveTime) -> Self {
        self.max_value = max;
        self
    }
}

impl Generator<NaiveTime> for NaiveTimeGenerator {
    fn do_draw(&self, tc: &TestCase) -> NaiveTime {
        if is_mid_day_leap(self.min_value) || is_mid_day_leap(self.max_value) {
            invalid_argument!(
                "naive_times bounds using a leap second are only supported at 23:59:59"
            );
        }
        let min_nanos = time_to_total_nanos(self.min_value);
        let max_nanos = time_to_total_nanos(self.max_value);
        if min_nanos > max_nanos {
            invalid_argument!("Cannot have max_value < min_value");
        }
        let n = integers::<i64>()
            .min_value(min_nanos)
            .max_value(max_nanos)
            .do_draw(tc);
        total_nanos_to_time(n)
    }
}

/// Generate [`chrono::NaiveTime`] values.
///
/// # Example
///
/// ```no_run
/// use chrono::NaiveTime;
/// use hegel::extras::chrono as chrono_gs;
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     let max = NaiveTime::from_hms_opt(12, 0, 0).unwrap();
///     let t = tc.draw(chrono_gs::naive_times().max_value(max));
///     assert!(t <= max);
/// }
/// ```
pub fn naive_times() -> NaiveTimeGenerator {
    NaiveTimeGenerator {
        min_value: NaiveTime::MIN,
        max_value: naive_time_default_max(),
    }
}

/// Generator for [`chrono::NaiveDateTime`] values. Created by [`naive_datetimes()`].
pub struct NaiveDateTimeGenerator {
    min_value: NaiveDateTime,
    max_value: NaiveDateTime,
}

impl NaiveDateTimeGenerator {
    /// Set the minimum datetime (inclusive).
    pub fn min_value(mut self, min: NaiveDateTime) -> Self {
        self.min_value = min;
        self
    }

    /// Set the maximum datetime (inclusive).
    pub fn max_value(mut self, max: NaiveDateTime) -> Self {
        self.max_value = max;
        self
    }
}

impl Generator<NaiveDateTime> for NaiveDateTimeGenerator {
    fn do_draw(&self, tc: &TestCase) -> NaiveDateTime {
        if self.min_value.time().nanosecond() >= 1_000_000_000
            || self.max_value.time().nanosecond() >= 1_000_000_000
        {
            invalid_argument!(
                "naive_datetimes bounds cannot use a leap second: the epoch-nanosecond \
                 encoding maps a leap second into the following second"
            );
        }
        if self.min_value > self.max_value {
            invalid_argument!("Cannot have max_value < min_value");
        }
        let n = integers::<i128>()
            .min_value(datetime_to_nanos(&self.min_value.and_utc()))
            .max_value(datetime_to_nanos(&self.max_value.and_utc()))
            .do_draw(tc);
        nanos_to_utc_datetime(n).naive_utc()
    }
}

/// Generate [`chrono::NaiveDateTime`] values.
///
/// Defaults span chrono's full `DateTime<Utc>` window —
/// [`DateTime::<Utc>::MIN_UTC`] through [`DateTime::<Utc>::MAX_UTC`].
///
/// # Example
///
/// ```no_run
/// use chrono::{NaiveDate, NaiveDateTime, NaiveTime};
/// use hegel::extras::chrono as chrono_gs;
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     let min = NaiveDateTime::new(
///         NaiveDate::from_ymd_opt(2024, 1, 1).unwrap(),
///         NaiveTime::from_hms_opt(0, 0, 0).unwrap(),
///     );
///     let max = NaiveDateTime::new(
///         NaiveDate::from_ymd_opt(2024, 12, 31).unwrap(),
///         NaiveTime::from_hms_nano_opt(23, 59, 59, 999_999_999).unwrap(),
///     );
///     let dt = tc.draw(chrono_gs::naive_datetimes().min_value(min).max_value(max));
///     assert!(dt >= min && dt <= max);
/// }
/// ```
pub fn naive_datetimes() -> NaiveDateTimeGenerator {
    NaiveDateTimeGenerator {
        min_value: DateTime::<Utc>::MIN_UTC.naive_utc(),
        max_value: DateTime::<Utc>::MAX_UTC.naive_utc(),
    }
}

/// Generator for [`chrono::NaiveWeek`] values. Created by [`naive_weeks()`].
pub struct NaiveWeekGenerator<S = <Weekday as DefaultGenerator>::Generator> {
    date_gen: NaiveDateGenerator,
    start_gen: S,
}

impl<S> NaiveWeekGenerator<S> {
    /// Set the minimum source date (inclusive).
    pub fn min_date(mut self, min: NaiveDate) -> Self {
        self.date_gen = self.date_gen.min_value(min);
        self
    }

    /// Set the maximum source date (inclusive).
    pub fn max_date(mut self, max: NaiveDate) -> Self {
        self.date_gen = self.date_gen.max_value(max);
        self
    }

    /// Replace the start-of-week generator.
    pub fn weekday_starts<S2>(self, start_gen: S2) -> NaiveWeekGenerator<S2>
    where
        S2: Generator<Weekday>,
    {
        NaiveWeekGenerator {
            date_gen: self.date_gen,
            start_gen,
        }
    }
}

impl<S: Generator<Weekday>> Generator<NaiveWeek> for NaiveWeekGenerator<S> {
    fn do_draw(&self, tc: &TestCase) -> NaiveWeek {
        let (date, start) = crate::generators::tuples2(&self.date_gen, &self.start_gen).do_draw(tc);
        date.week(start)
    }
}

/// Generate [`chrono::NaiveWeek`] values.
///
/// Each draw produces a 7-day window around a randomly chosen [`NaiveDate`],
/// with the start-of-week [`Weekday`] also drawn at random (full 7-day range
/// by default). Use [`weekday_starts`](NaiveWeekGenerator::weekday_starts) to substitute a
/// different start-day strategy — e.g. `gs::just(Weekday::Sun)` for Sunday-only.
///
/// # Example
///
/// ```no_run
/// use chrono::Weekday;
/// use hegel::extras::chrono as chrono_gs;
/// use hegel::generators as gs;
///
/// #[hegel::test]
/// fn my_test(tc: hegel::TestCase) {
///     // Default: random start day per draw
///     let _ = tc.draw(chrono_gs::naive_weeks());
///
///     // All Sunday-start weeks
///     let _ = tc.draw(chrono_gs::naive_weeks().weekday_starts(gs::just(Weekday::Sun)));
/// }
/// ```
pub fn naive_weeks() -> NaiveWeekGenerator {
    // Pull the default date range in by a week on each side: a NaiveWeek
    // whose window pokes past NaiveDate::MIN/MAX panics in chrono's own
    // first_day/last_day/days accessors, so default-generated weeks would
    // otherwise be landmines. An explicit min_date/max_date can still opt
    // back into the extremes.
    NaiveWeekGenerator {
        date_gen: naive_dates()
            .min_value(
                NaiveDate::MIN
                    .checked_add_days(chrono::Days::new(7))
                    .unwrap(),
            )
            .max_value(
                NaiveDate::MAX
                    .checked_sub_days(chrono::Days::new(7))
                    .unwrap(),
            ),
        start_gen: Weekday::default_generator(),
    }
}

/// Convert a [`DateTime`] to a total nanosecond count from the Unix epoch.
fn datetime_to_nanos<Tz: TimeZone>(dt: &DateTime<Tz>) -> i128 {
    i128::from(dt.timestamp()) * 1_000_000_000 + i128::from(dt.timestamp_subsec_nanos())
}

/// Inverse of [`datetime_to_nanos`], producing a `DateTime<Utc>`.
fn nanos_to_utc_datetime(n: i128) -> DateTime<Utc> {
    let secs = n.div_euclid(1_000_000_000) as i64;
    let nsecs = n.rem_euclid(1_000_000_000) as u32;
    DateTime::<Utc>::from_timestamp(secs, nsecs).unwrap()
}

/// Generator for [`chrono::DateTime`] values. Created by [`datetimes()`].
pub struct DateTimeGenerator<G = FixedOffsetGenerator, Tz: TimeZone = FixedOffset> {
    tz_gen: G,
    min_value: NaiveDateTime,
    max_value: NaiveDateTime,
    _phantom: PhantomData<fn() -> Tz>,
}

impl<G, Tz: TimeZone> DateTimeGenerator<G, Tz> {
    /// Set the minimum wall-clock datetime (inclusive).
    pub fn min_value(mut self, min: NaiveDateTime) -> Self {
        self.min_value = min;
        self
    }

    /// Set the maximum wall-clock datetime (inclusive).
    pub fn max_value(mut self, max: NaiveDateTime) -> Self {
        self.max_value = max;
        self
    }

    /// Use the specified timezone generator.
    ///
    /// The returned generator has output type `DateTime<Tz2>` where `Tz2`
    /// is the timezone type produced by `tz_gen`. Wall-clock bounds are
    /// timezone-agnostic and carry across.
    pub fn timezones<G2, Tz2>(self, tz_gen: G2) -> DateTimeGenerator<G2, Tz2>
    where
        G2: Generator<Tz2>,
        Tz2: TimeZone,
    {
        DateTimeGenerator {
            tz_gen,
            min_value: self.min_value,
            max_value: self.max_value,
            _phantom: PhantomData,
        }
    }
}

impl<G, Tz> Generator<DateTime<Tz>> for DateTimeGenerator<G, Tz>
where
    G: Generator<Tz>,
    Tz: TimeZone + Send + Sync + 'static,
{
    fn do_draw(&self, tc: &TestCase) -> DateTime<Tz> {
        let naive = naive_datetimes()
            .min_value(self.min_value)
            .max_value(self.max_value)
            .do_draw(tc);
        let tz = self.tz_gen.do_draw(tc);
        match tz.from_local_datetime(&naive).earliest() {
            Some(dt) => dt,
            None => {
                tc.assume(false);
                unreachable!()
            }
        }
    }
}

/// Generate [`chrono::DateTime`] values.
///
/// ```no_run
/// use chrono::{FixedOffset, Utc};
/// use hegel::extras::chrono as chrono_gs;
/// use hegel::generators as gs;
///
/// // Default: DateTime<FixedOffset> with varying offsets
/// let g = chrono_gs::datetimes();
///
/// // All UTC
/// let g = chrono_gs::datetimes().timezones(gs::just(Utc));
///
/// // All in a single fixed offset
/// let off = FixedOffset::east_opt(9 * 3600).unwrap();
/// let g = chrono_gs::datetimes().timezones(gs::just(off));
///
/// // Constrain offsets to a sub-range
/// let g = chrono_gs::datetimes()
///     .timezones(chrono_gs::fixed_offsets().min_value(off));
/// ```
pub fn datetimes() -> DateTimeGenerator<FixedOffsetGenerator, FixedOffset> {
    DateTimeGenerator {
        tz_gen: fixed_offsets(),
        min_value: DateTime::<Utc>::MIN_UTC.naive_utc(),
        max_value: DateTime::<Utc>::MAX_UTC.naive_utc(),
        _phantom: PhantomData,
    }
}