tiny_counter/store/

limiter.rs

/// Rate limiting with fluent constraint API.
use std::fmt;
use std::sync::Arc;

use chrono::{DateTime, Datelike, Duration, Local, Timelike, Utc, Weekday};

use crate::{Error, LimitExceeded, Result, TimeUnit, TimeWindow};

use super::inner::EventStoreInner;
use super::query::Query;
use super::EventId;

/// Fluent builder for rate limiting with multiple constraints.
///
/// Created by `EventStore::limit()`, this allows chaining multiple constraints
/// that must all pass before an event can be recorded.
pub struct Limiter {
    store: Arc<EventStoreInner>,
    constraints: Vec<Constraint>,
}

/// A rate limiting constraint.
#[derive(Debug, Clone)]
pub enum Constraint {
    /// Event must not exceed count within time window.
    AtMost {
        event_id: String,
        limit: u32,
        window_count: usize,
        time_unit: TimeUnit,
    },
    /// Event must have been recorded at least count times within time window.
    AtLeast {
        event_id: String,
        count: u32,
        window_count: usize,
        time_unit: TimeUnit,
    },
    /// Event must wait duration since last occurrence.
    Cooldown {
        event_id: String,
        duration: Duration,
    },
    /// Prerequisite event must have occurred within duration.
    Within {
        prerequisite_event: String,
        duration: Duration,
    },
    /// Current time must be within schedule.
    During { schedule: Schedule },
    /// Current time must be outside schedule.
    OutsideOf { schedule: Schedule },
}

/// A time-based schedule for constraint checking.
#[derive(Clone)]
pub enum Schedule {
    /// Business hours (e.g., 9am-5pm).
    /// When `use_local_tz` is true, hours are interpreted in the system's local timezone.
    /// When false, hours are interpreted in UTC.
    Hours {
        start_hour: u8,
        end_hour: u8,
        use_local_tz: bool,
    },

    /// Monday through Friday.
    /// When `use_local_tz` is true, weekday is determined in the system's local timezone.
    /// When false, weekday is determined in UTC.
    Weekdays { use_local_tz: bool },

    /// Saturday and Sunday.
    /// When `use_local_tz` is true, weekday is determined in the system's local timezone.
    /// When false, weekday is determined in UTC.
    Weekends { use_local_tz: bool },

    /// Custom schedule using a closure.
    Custom(Arc<dyn Fn(DateTime<Utc>) -> bool + Send + Sync>),
}

impl fmt::Debug for Schedule {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Schedule::Hours {
                start_hour,
                end_hour,
                use_local_tz,
            } => {
                let tz_str = if *use_local_tz { "Local" } else { "UTC" };
                write!(
                    f,
                    "Hours {{ start_hour: {}, end_hour: {}, timezone: {} }}",
                    start_hour, end_hour, tz_str
                )
            }
            Schedule::Weekdays { use_local_tz } => {
                let tz_str = if *use_local_tz { "Local" } else { "UTC" };
                write!(f, "Weekdays({})", tz_str)
            }
            Schedule::Weekends { use_local_tz } => {
                let tz_str = if *use_local_tz { "Local" } else { "UTC" };
                write!(f, "Weekends({})", tz_str)
            }
            Schedule::Custom(_) => write!(f, "Custom(<closure>)"),
        }
    }
}

impl Schedule {
    /// Create a business hours schedule (UTC).
    ///
    /// For local timezone-aware schedules, use `hours_local_tz()`.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - `start_hour` is greater than 23
    /// - `end_hour` is greater than 23
    /// - `start_hour` is greater than or equal to `end_hour`
    ///
    /// # Examples
    ///
    /// ```
    /// use tiny_counter::Schedule;
    ///
    /// let schedule = Schedule::hours(9, 17).unwrap();  // 9am-5pm UTC
    /// let schedule = Schedule::hours(6, 22).unwrap();  // 6am-10pm UTC
    /// let schedule = Schedule::hours_local_tz(9, 17).unwrap();  // 9am-5pm local time
    /// ```
    pub fn hours(start_hour: u8, end_hour: u8) -> Result<Self> {
        Self::validate_hours(start_hour, end_hour)?;
        Ok(Schedule::Hours {
            start_hour,
            end_hour,
            use_local_tz: false,
        })
    }

    /// Create a business hours schedule using the system's local timezone.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - `start_hour` is greater than 23
    /// - `end_hour` is greater than 23
    /// - `start_hour` is greater than or equal to `end_hour`
    ///
    /// # Examples
    ///
    /// ```
    /// use tiny_counter::Schedule;
    ///
    /// let schedule = Schedule::hours_local_tz(9, 17).unwrap();  // 9am-5pm in system's local timezone
    /// ```
    pub fn hours_local_tz(start_hour: u8, end_hour: u8) -> Result<Self> {
        Self::validate_hours(start_hour, end_hour)?;
        Ok(Schedule::Hours {
            start_hour,
            end_hour,
            use_local_tz: true,
        })
    }

    /// Create a weekdays schedule (Monday-Friday, UTC).
    pub fn weekdays() -> Self {
        Schedule::Weekdays {
            use_local_tz: false,
        }
    }

    /// Create a weekdays schedule using the system's local timezone.
    pub fn weekdays_local_tz() -> Self {
        Schedule::Weekdays { use_local_tz: true }
    }

    /// Create a weekends schedule (Saturday-Sunday, UTC).
    pub fn weekends() -> Self {
        Schedule::Weekends {
            use_local_tz: false,
        }
    }

    /// Create a weekends schedule using the system's local timezone.
    pub fn weekends_local_tz() -> Self {
        Schedule::Weekends { use_local_tz: true }
    }

    fn validate_hours(start_hour: u8, end_hour: u8) -> Result<()> {
        if start_hour > 23 {
            Err(Error::InvalidHour(start_hour))
        } else if end_hour > 23 {
            Err(Error::InvalidHour(end_hour))
        } else if start_hour >= end_hour {
            Err(Error::InvalidRange(start_hour, end_hour))
        } else {
            Ok(())
        }
    }
}

/// Usage information for a rate limit constraint.
pub struct LimitUsage {
    pub count: u32,
    pub limit: u32,
    pub remaining: u32,
}

/// A reservation for a rate-limited event.
///
/// Created by `Limiter::reserve()`, this represents a reserved slot that must be
/// explicitly committed or will automatically cancel on drop.
///
/// # Examples
///
/// ```
/// use tiny_counter::{EventStore, TimeUnit};
///
/// let store = EventStore::new();
///
/// // Reserve a slot
/// let reservation = store.limit()
///     .at_most("api_call", 10, TimeUnit::Hours)
///     .reserve("api_call")
///     .unwrap();
///
/// // Do some work...
/// // If successful, commit
/// reservation.commit();
///
/// // If reservation is dropped without commit, it auto-cancels
/// ```
pub struct Reservation {
    store: Arc<EventStoreInner>,
    event_id: String,
    committed: bool,
}

impl Reservation {
    /// Commit the reservation, recording the event.
    ///
    /// This decrements the pending counter and records the event atomically.
    pub fn commit(mut self) {
        self.release_and_record();
        self.committed = true;
    }

    /// Cancel the reservation without recording the event.
    ///
    /// This is equivalent to dropping the reservation without committing.
    pub fn cancel(mut self) {
        self.release_pending();
        self.committed = true;
    }

    /// Release the pending reservation and record the event.
    fn release_and_record(&self) {
        if let Some(counter_arc) = self.store.events.get(&self.event_id) {
            let mut counter = counter_arc.lock().unwrap();
            counter.decrement_pending();
            let now = self.store.clock_now();
            counter.advance_if_needed(now);
            counter.record(1);
            counter.mark_dirty();
        }
    }

    /// Release the pending reservation without recording.
    fn release_pending(&self) {
        if let Some(counter_arc) = self.store.events.get(&self.event_id) {
            let mut counter = counter_arc.lock().unwrap();
            counter.decrement_pending();
            counter.mark_dirty();
        }
    }
}

impl Drop for Reservation {
    fn drop(&mut self) {
        if !self.committed {
            // Auto-cancel if not committed
            self.release_pending();
        }
    }
}

impl Limiter {
    /// Creates a new Limiter builder.
    pub(crate) fn new(store: Arc<EventStoreInner>) -> Self {
        Self {
            store,
            constraints: Vec::new(),
        }
    }

    /// Add an "at most" constraint: event cannot exceed limit within time window.
    ///
    /// Accepts flexible time window specifications:
    /// - `TimeUnit` - defaults to 1 unit (e.g., `TimeUnit::Days` means 1 day)
    /// - `(usize, TimeUnit)` - explicit count (e.g., `(7, TimeUnit::Days)` means 7 days)
    /// - `Duration` - converted to best matching TimeUnit
    ///
    /// # Examples
    ///
    /// ```
    /// use tiny_counter::{EventStore, TimeUnit};
    /// use chrono::Duration;
    ///
    /// let mut store = EventStore::new();
    ///
    /// // Single time unit (backward compatible)
    /// let result = store
    ///     .limit()
    ///     .at_most("api_call", 10, TimeUnit::Minutes)
    ///     .check_and_record("api_call");
    /// assert!(result.is_ok());
    ///
    /// // Multiple time units with tuple
    /// let result = store
    ///     .limit()
    ///     .at_most("api_call", 100, (7, TimeUnit::Days))
    ///     .check("api_call");
    /// assert!(result.is_ok());
    ///
    /// // Duration syntax
    /// let result = store
    ///     .limit()
    ///     .at_most("api_call", 100, Duration::days(7))
    ///     .check("api_call");
    /// assert!(result.is_ok());
    /// ```
    pub fn at_most(
        mut self,
        event_id: impl EventId,
        limit: u32,
        window: impl Into<TimeWindow>,
    ) -> Self {
        let window = window.into();
        self.constraints.push(Constraint::AtMost {
            event_id: event_id.as_ref().to_string(),
            limit,
            window_count: window.count,
            time_unit: window.time_unit,
        });
        self
    }

    /// Add an "at least" constraint: prerequisite must occur at least count times within time window.
    ///
    /// Accepts flexible time window specifications:
    /// - `TimeUnit` - defaults to 1 unit (e.g., `TimeUnit::Days` means 1 day)
    /// - `(usize, TimeUnit)` - explicit count (e.g., `(7, TimeUnit::Days)` means 7 days)
    /// - `Duration` - converted to best matching TimeUnit
    pub fn at_least(
        mut self,
        event_id: impl EventId,
        count: u32,
        window: impl Into<TimeWindow>,
    ) -> Self {
        let window = window.into();
        self.constraints.push(Constraint::AtLeast {
            event_id: event_id.as_ref().to_string(),
            count,
            window_count: window.count,
            time_unit: window.time_unit,
        });
        self
    }

    /// Add a cooldown constraint: event must wait duration since last occurrence.
    pub fn cooldown(mut self, event_id: impl EventId, duration: Duration) -> Self {
        self.constraints.push(Constraint::Cooldown {
            event_id: event_id.as_ref().to_string(),
            duration,
        });
        self
    }

    /// Add a "within" constraint: prerequisite must have occurred within duration.
    pub fn within(mut self, prerequisite_event: impl EventId, duration: Duration) -> Self {
        self.constraints.push(Constraint::Within {
            prerequisite_event: prerequisite_event.as_ref().to_string(),
            duration,
        });
        self
    }

    /// Add a "during" constraint: event must occur within schedule.
    pub fn during(mut self, schedule: Schedule) -> Self {
        self.constraints.push(Constraint::During { schedule });
        self
    }

    /// Add an "outside of" constraint: event must occur outside schedule.
    pub fn outside_of(mut self, schedule: Schedule) -> Self {
        self.constraints.push(Constraint::OutsideOf { schedule });
        self
    }

    /// Check if all constraints pass without recording the event.
    ///
    /// # Examples
    ///
    /// ```
    /// use tiny_counter::{EventStore, TimeUnit};
    ///
    /// let store = EventStore::new();
    ///
    /// let result = store
    ///     .limit()
    ///     .at_most("action", 5, TimeUnit::Minutes)
    ///     .check("action");
    ///
    /// assert!(result.is_ok());
    /// ```
    pub fn check(self, event_id: impl EventId) -> Result<()> {
        let event_id_str = event_id.as_ref();
        self.evaluate_constraints(event_id_str)
    }

    /// Check if all constraints pass, and if so, record the event.
    ///
    /// This operation is atomic: the event is recorded only if all constraints pass.
    ///
    /// # Examples
    ///
    /// ```
    /// use tiny_counter::{EventStore, TimeUnit};
    ///
    /// let store = EventStore::new();
    ///
    /// // First call succeeds and records the event
    /// let result = store
    ///     .limit()
    ///     .at_most("limited_action", 1, TimeUnit::Minutes)
    ///     .check_and_record("limited_action");
    /// assert!(result.is_ok());
    ///
    /// // Second call fails because limit is reached
    /// let result = store
    ///     .limit()
    ///     .at_most("limited_action", 1, TimeUnit::Minutes)
    ///     .check_and_record("limited_action");
    /// assert!(result.is_err());
    /// ```
    pub fn check_and_record(self, event_id: impl EventId) -> Result<()> {
        let event_id_str = event_id.as_ref();
        self.evaluate_constraints(event_id_str)?;

        // Inline record logic: get or create counter and record
        let now = self.store.clock.now();
        let counter = self.store.get_counter_for_record(event_id_str);

        let mut counter = counter.lock().unwrap();
        counter.advance_if_needed(now);
        counter.record(1);
        counter.mark_dirty();

        Ok(())
    }

    /// Reserve a slot if all constraints pass, returning a Reservation.
    ///
    /// This atomically checks rate limits including pending reservations,
    /// preventing race conditions. The reservation must be explicitly committed
    /// to record the event, or it will auto-cancel on drop.
    ///
    /// # Examples
    ///
    /// ```
    /// use tiny_counter::{EventStore, TimeUnit};
    ///
    /// let store = EventStore::new();
    ///
    /// // Reserve a slot (atomically checks and reserves)
    /// let reservation = store.limit()
    ///     .at_most("api_call", 10, TimeUnit::Hours)
    ///     .reserve("api_call").unwrap();
    ///
    /// // Do work (e.g., make API call)
    /// // ...
    ///
    /// // Commit if successful
    /// reservation.commit();
    /// ```
    ///
    /// # Transactional Pattern
    ///
    /// ```
    /// use tiny_counter::{EventStore, TimeUnit};
    ///
    /// fn make_api_call() -> Result<(), Box<dyn std::error::Error>> {
    ///     // ... API call logic
    ///     Ok(())
    /// }
    ///
    /// let store = EventStore::new();
    ///
    /// let reservation = store.limit()
    ///     .at_most("api_call", 10, TimeUnit::Hours)
    ///     .reserve("api_call").unwrap();
    ///
    /// match make_api_call() {
    ///     Ok(_) => reservation.commit(),  // Record on success
    ///     Err(_) => {
    ///         // Reservation auto-cancels on drop
    ///         // Or explicitly: reservation.cancel();
    ///     }
    /// }
    /// # Ok::<(), Box<dyn std::error::Error>>(())
    /// ```
    pub fn reserve(self, event_id: impl EventId) -> Result<Reservation> {
        let clock_now = self.store.clock_now();
        let event_id_str = event_id.as_ref();

        // Get or create counter with write lock
        let counter_arc = self.store.get_counter_for_record(event_id_str);

        let mut counter = counter_arc.lock().unwrap();
        counter.advance_if_needed(clock_now);

        // Check all constraints including pending
        for constraint in &self.constraints {
            match constraint {
                Constraint::AtMost {
                    event_id: _constraint_event,
                    limit,
                    window_count,
                    time_unit,
                } => {
                    // Use total_with_pending to include reservations
                    let total = counter.total_with_pending(*time_unit, 0..*window_count);
                    if total >= *limit {
                        return Err(Error::LimitExceeded(LimitExceeded {
                            event_id: event_id_str.to_string(),
                            constraint: constraint.clone(),
                            retry_after: Some(time_unit.duration()),
                        }));
                    }
                }
                // Other constraint types checked without pending
                _ => {
                    // Release lock temporarily to evaluate
                    drop(counter);
                    self.evaluate_constraint(event_id_str, constraint)?;
                    counter = counter_arc.lock().unwrap();
                }
            }
        }

        // All constraints pass - reserve the slot
        counter.increment_pending();
        counter.mark_dirty();

        Ok(Reservation {
            store: self.store.clone(),
            event_id: event_id_str.to_string(),
            committed: false,
        })
    }

    /// Check if all constraints would pass (convenience method).
    pub fn allowed(self, event_id: impl EventId) -> bool {
        self.check(event_id).is_ok()
    }

    /// Get usage information for an "at most" constraint.
    ///
    /// Returns usage data for the first AtMost constraint that matches the event_id.
    pub fn usage(self, event_id: impl EventId) -> Result<LimitUsage> {
        let event_id_str = event_id.as_ref();

        // Find the first AtMost constraint matching this event
        for constraint in &self.constraints {
            if let Constraint::AtMost {
                event_id: constraint_event,
                limit,
                window_count,
                time_unit,
            } = constraint
            {
                if constraint_event == event_id_str {
                    let current = self
                        .query_sum_window(constraint_event, *window_count, *time_unit)
                        .unwrap_or(0);
                    let remaining = limit.saturating_sub(current);
                    return Ok(LimitUsage {
                        count: current,
                        limit: *limit,
                        remaining,
                    });
                }
            }
        }

        // No AtMost constraint found, return zero usage
        Ok(LimitUsage {
            count: 0,
            limit: 0,
            remaining: 0,
        })
    }

    /// Evaluates all constraints, short-circuiting on first failure.
    fn evaluate_constraints(&self, event_id: &str) -> Result<()> {
        for constraint in &self.constraints {
            self.evaluate_constraint(event_id, constraint)?;
        }
        Ok(())
    }

    /// Evaluates a single constraint.
    fn evaluate_constraint(&self, event_id: &str, constraint: &Constraint) -> Result<()> {
        match constraint {
            Constraint::AtMost {
                event_id: constraint_event,
                limit,
                window_count,
                time_unit,
            } => {
                let current = self
                    .query_sum_window(constraint_event, *window_count, *time_unit)
                    .unwrap_or(0);
                if current >= *limit {
                    return Err(Error::LimitExceeded(LimitExceeded {
                        event_id: event_id.to_string(),
                        constraint: constraint.clone(),
                        retry_after: Some(time_unit.duration()),
                    }));
                }
            }
            Constraint::AtLeast {
                event_id: constraint_event,
                count,
                window_count,
                time_unit,
            } => {
                let current = self
                    .query_sum_window(constraint_event, *window_count, *time_unit)
                    .unwrap_or(0);
                if current < *count {
                    return Err(Error::LimitExceeded(LimitExceeded {
                        event_id: event_id.to_string(),
                        constraint: constraint.clone(),
                        retry_after: None,
                    }));
                }
            }
            Constraint::Cooldown {
                event_id: constraint_event,
                duration,
            } => {
                if let Some(elapsed) = self.query_last_seen(constraint_event) {
                    if elapsed < *duration {
                        let remaining = *duration - elapsed;
                        return Err(Error::LimitExceeded(LimitExceeded {
                            event_id: event_id.to_string(),
                            constraint: constraint.clone(),
                            retry_after: Some(remaining),
                        }));
                    }
                }
                // Never seen before = cooldown is satisfied
            }
            Constraint::Within {
                prerequisite_event,
                duration,
            } => {
                match self.query_last_seen(prerequisite_event) {
                    None => {
                        // Never seen
                        return Err(Error::LimitExceeded(LimitExceeded {
                            event_id: event_id.to_string(),
                            constraint: constraint.clone(),
                            retry_after: None,
                        }));
                    }
                    Some(elapsed) => {
                        if elapsed > *duration {
                            // Too long ago
                            return Err(Error::LimitExceeded(LimitExceeded {
                                event_id: event_id.to_string(),
                                constraint: constraint.clone(),
                                retry_after: None,
                            }));
                        }
                    }
                }
            }
            Constraint::During { schedule } => {
                if !schedule.allows(&self.store) {
                    return Err(Error::LimitExceeded(LimitExceeded {
                        event_id: event_id.to_string(),
                        constraint: constraint.clone(),
                        retry_after: schedule.next_allowed(&self.store),
                    }));
                }
            }
            Constraint::OutsideOf { schedule } => {
                if schedule.allows(&self.store) {
                    return Err(Error::LimitExceeded(LimitExceeded {
                        event_id: event_id.to_string(),
                        constraint: constraint.clone(),
                        retry_after: schedule.block_end(&self.store),
                    }));
                }
            }
        }
        Ok(())
    }

    /// Query the count for an event within a time window.
    fn query_sum_window(
        &self,
        event_id: &str,
        window_count: usize,
        time_unit: TimeUnit,
    ) -> Option<u32> {
        let query = Query::new(self.store.clone(), event_id.to_string());
        query.last(window_count, time_unit).sum()
    }

    /// Query the time since last seen.
    fn query_last_seen(&self, event_id: &str) -> Option<Duration> {
        Query::new(self.store.clone(), event_id.to_string()).last_seen()
    }
}

impl Schedule {
    /// Check if the schedule allows the current time.
    fn allows(&self, store: &EventStoreInner) -> bool {
        let now = store.clock_now();
        match self {
            Schedule::Hours {
                start_hour,
                end_hour,
                use_local_tz,
            } => {
                let hour = if *use_local_tz {
                    now.with_timezone(&Local).hour() as u8
                } else {
                    now.hour() as u8
                };
                hour >= *start_hour && hour < *end_hour
            }
            Schedule::Weekdays { use_local_tz } => {
                let weekday = if *use_local_tz {
                    now.with_timezone(&Local).weekday()
                } else {
                    now.weekday()
                };
                !matches!(weekday, Weekday::Sat | Weekday::Sun)
            }
            Schedule::Weekends { use_local_tz } => {
                let weekday = if *use_local_tz {
                    now.with_timezone(&Local).weekday()
                } else {
                    now.weekday()
                };
                matches!(weekday, Weekday::Sat | Weekday::Sun)
            }
            Schedule::Custom(f) => f(now),
        }
    }

    /// Calculate when the schedule next allows events.
    fn next_allowed(&self, store: &EventStoreInner) -> Option<Duration> {
        let now = store.clock_now();
        match self {
            Schedule::Hours {
                start_hour,
                end_hour: _,
                use_local_tz,
            } => {
                let current_hour = if *use_local_tz {
                    now.with_timezone(&Local).hour() as u8
                } else {
                    now.hour() as u8
                };

                if current_hour < *start_hour {
                    // Before business hours today
                    let hours_until = (*start_hour - current_hour) as i64;
                    Some(Duration::hours(hours_until))
                } else {
                    // After business hours, wait until tomorrow
                    let hours_until = 24 - current_hour + *start_hour;
                    Some(Duration::hours(hours_until as i64))
                }
            }
            Schedule::Weekdays { use_local_tz } => {
                let weekday = if *use_local_tz {
                    now.with_timezone(&Local).weekday()
                } else {
                    now.weekday()
                };
                // If weekend, wait until Monday
                match weekday {
                    Weekday::Sat => Some(Duration::days(2)),
                    Weekday::Sun => Some(Duration::days(1)),
                    _ => None,
                }
            }
            Schedule::Weekends { use_local_tz } => {
                let weekday = if *use_local_tz {
                    now.with_timezone(&Local).weekday()
                } else {
                    now.weekday()
                };
                // If weekday, wait until Saturday
                let days_until_sat = match weekday {
                    Weekday::Mon => 5,
                    Weekday::Tue => 4,
                    Weekday::Wed => 3,
                    Weekday::Thu => 2,
                    Weekday::Fri => 1,
                    _ => 0,
                };
                if days_until_sat > 0 {
                    Some(Duration::days(days_until_sat))
                } else {
                    None
                }
            }
            Schedule::Custom(_) => None, // Can't calculate for custom schedules
        }
    }

    /// Calculate when the schedule block ends.
    fn block_end(&self, store: &EventStoreInner) -> Option<Duration> {
        let now = store.clock_now();
        match self {
            Schedule::Hours {
                start_hour: _,
                end_hour,
                use_local_tz,
            } => {
                let current_hour = if *use_local_tz {
                    now.with_timezone(&Local).hour() as u8
                } else {
                    now.hour() as u8
                };

                if current_hour < *end_hour {
                    let hours_until = (*end_hour - current_hour) as i64;
                    Some(Duration::hours(hours_until))
                } else {
                    None
                }
            }
            Schedule::Weekdays { use_local_tz } => {
                let weekday = if *use_local_tz {
                    now.with_timezone(&Local).weekday()
                } else {
                    now.weekday()
                };
                // Blocked during weekdays, ends on Saturday
                match weekday {
                    Weekday::Mon => Some(Duration::days(5)),
                    Weekday::Tue => Some(Duration::days(4)),
                    Weekday::Wed => Some(Duration::days(3)),
                    Weekday::Thu => Some(Duration::days(2)),
                    Weekday::Fri => Some(Duration::days(1)),
                    _ => None,
                }
            }
            Schedule::Weekends { use_local_tz } => {
                let weekday = if *use_local_tz {
                    now.with_timezone(&Local).weekday()
                } else {
                    now.weekday()
                };
                // Blocked during weekends, ends on Monday
                match weekday {
                    Weekday::Sat => Some(Duration::days(2)),
                    Weekday::Sun => Some(Duration::days(1)),
                    _ => None,
                }
            }
            Schedule::Custom(_) => None, // Can't calculate for custom schedules
        }
    }
}

#[cfg(test)]
mod tests {
    use std::thread;

    use chrono::{Duration, TimeZone, Weekday};

    use crate::{EventStore, TestClock, TimeUnit};

    use super::*;

    #[test]
    fn test_at_most_constraint_passes() {
        let store = EventStore::new();
        store.record_count("test_event", 5);

        let result = store
            .limit()
            .at_most("test_event", 10, TimeUnit::Days)
            .check("action");

        assert!(result.is_ok());
    }

    #[test]
    fn test_at_most_constraint_fails() {
        let store = EventStore::new();
        store.record_count("test_event", 10);

        let result = store
            .limit()
            .at_most("test_event", 5, TimeUnit::Days)
            .check("action");

        assert!(result.is_err());
        let Error::LimitExceeded(err) = result.unwrap_err() else {
            panic!("Expected LimitExceeded error");
        };
        assert_eq!(err.event_id, "action");
        assert!(err.retry_after.is_some());
    }

    #[test]
    fn test_at_least_constraint_passes() {
        let store = EventStore::new();
        store.record_count("prerequisite", 5);

        let result = store
            .limit()
            .at_least("prerequisite", 3, TimeUnit::Days)
            .check("action");

        assert!(result.is_ok());
    }

    #[test]
    fn test_at_least_constraint_fails() {
        let store = EventStore::new();
        store.record_count("prerequisite", 2);

        let result = store
            .limit()
            .at_least("prerequisite", 5, TimeUnit::Days)
            .check("action");

        assert!(result.is_err());
        let Error::LimitExceeded(err) = result.unwrap_err() else {
            panic!("Expected LimitExceeded error");
        };
        assert!(err.retry_after.is_none());
    }

    #[test]
    fn test_cooldown_constraint_passes() {
        let store = EventStore::new();
        store.record_ago("test_event", Duration::hours(2));

        let result = store
            .limit()
            .cooldown("test_event", Duration::hours(1))
            .check("action");

        assert!(result.is_ok());
    }

    #[test]
    fn test_cooldown_constraint_fails() {
        let store = EventStore::new();
        store.record("test_event");

        let result = store
            .limit()
            .cooldown("test_event", Duration::hours(1))
            .check("action");

        assert!(result.is_err());
        let Error::LimitExceeded(err) = result.unwrap_err() else {
            panic!("Expected LimitExceeded error");
        };
        assert!(err.retry_after.is_some());
    }

    #[test]
    fn test_within_constraint_passes() {
        let store = EventStore::new();
        store.record("prerequisite");

        let result = store
            .limit()
            .within("prerequisite", Duration::hours(1))
            .check("action");

        assert!(result.is_ok());
    }

    #[test]
    fn test_within_constraint_fails_never_seen() {
        let store = EventStore::new();

        let result = store
            .limit()
            .within("prerequisite", Duration::hours(1))
            .check("action");

        assert!(result.is_err());
    }

    #[test]
    fn test_within_constraint_fails_too_long_ago() {
        let store = EventStore::new();
        store.record_ago("prerequisite", Duration::hours(2));

        let result = store
            .limit()
            .within("prerequisite", Duration::hours(1))
            .check("action");

        assert!(result.is_err());
    }

    #[test]
    fn test_multiple_constraints_all_pass() {
        let store = EventStore::new();
        store.record_count("prerequisite", 5);
        store.record_count("test_event", 2);

        let result = store
            .limit()
            .at_least("prerequisite", 3, TimeUnit::Days)
            .at_most("test_event", 10, TimeUnit::Days)
            .check("action");

        assert!(result.is_ok());
    }

    #[test]
    fn test_multiple_constraints_first_fails() {
        let store = EventStore::new();
        store.record_count("prerequisite", 1);
        store.record_count("test_event", 2);

        let result = store
            .limit()
            .at_least("prerequisite", 3, TimeUnit::Days)
            .at_most("test_event", 10, TimeUnit::Days)
            .check("action");

        assert!(result.is_err());
    }

    #[test]
    fn test_check_and_record_success() {
        let store = EventStore::new();

        let result = store
            .limit()
            .at_most("test_event", 5, TimeUnit::Days)
            .check_and_record("test_event");

        assert!(result.is_ok());

        // Verify event was recorded
        let count = store.query("test_event").last_days(1).sum();
        assert_eq!(count, Some(1));
    }

    #[test]
    fn test_check_and_record_failure_does_not_record() {
        let store = EventStore::new();
        store.record_count("test_event", 10);

        let result = store
            .limit()
            .at_most("test_event", 5, TimeUnit::Days)
            .check_and_record("test_event");

        assert!(result.is_err());

        // Verify count is still 10 (not 11)
        let count = store.query("test_event").last_days(1).sum();
        assert_eq!(count, Some(10));
    }

    #[test]
    fn test_allowed_convenience_method() {
        let store = EventStore::new();
        store.record_count("test_event", 3);

        let allowed = store
            .limit()
            .at_most("test_event", 5, TimeUnit::Days)
            .allowed("action");
        assert!(allowed);

        // After adding 2 more (total 5), should hit limit of 5
        store.record_count("test_event", 2);
        let not_allowed = store
            .limit()
            .at_most("test_event", 5, TimeUnit::Days)
            .allowed("action");
        assert!(!not_allowed);
    }

    #[test]
    fn test_usage_returns_correct_counts() {
        let store = EventStore::new();
        store.record_count("test_event", 3);

        let usage = store
            .limit()
            .at_most("test_event", 10, TimeUnit::Days)
            .usage("test_event")
            .unwrap();

        assert_eq!(usage.count, 3);
        assert_eq!(usage.limit, 10);
        assert_eq!(usage.remaining, 7);
    }

    #[test]
    fn test_business_hours_schedule_during_hours() {
        // Set to Tuesday at 10 AM (during business hours)
        let tuesday_10am = Utc.with_ymd_and_hms(2025, 1, 7, 10, 0, 0).unwrap();
        let clock = TestClock::build_for_testing_at(tuesday_10am);

        let store = EventStore::builder()
            .with_clock(Arc::new(clock.clone()))
            .build()
            .unwrap();

        // Should succeed during business hours (9am-5pm)
        let result = store
            .limit()
            .during(Schedule::hours(9, 17).unwrap())
            .check_and_record("action");
        assert!(result.is_ok());

        // Set to Tuesday at 7 PM (outside business hours)
        let tuesday_7pm = Utc.with_ymd_and_hms(2025, 1, 7, 19, 0, 0).unwrap();
        clock.set(tuesday_7pm);

        // Should fail outside business hours
        let result = store
            .limit()
            .during(Schedule::hours(9, 17).unwrap())
            .check("action");
        assert!(result.is_err());
    }

    #[test]
    fn test_weekdays_schedule() {
        // Set to Tuesday (weekday)
        let tuesday = Utc.with_ymd_and_hms(2025, 1, 7, 12, 0, 0).unwrap();
        assert_eq!(tuesday.weekday(), Weekday::Tue);
        let clock = TestClock::build_for_testing_at(tuesday);

        let store = EventStore::builder()
            .with_clock(Arc::new(clock.clone()))
            .build()
            .unwrap();

        // Should succeed on weekday
        let result = store
            .limit()
            .during(Schedule::weekdays())
            .check_and_record("action");
        assert!(result.is_ok());

        // Set to Saturday (weekend)
        let saturday = Utc.with_ymd_and_hms(2025, 1, 11, 12, 0, 0).unwrap();
        assert_eq!(saturday.weekday(), Weekday::Sat);
        clock.set(saturday);

        // Should fail on weekend
        let result = store.limit().during(Schedule::weekdays()).check("action");
        assert!(result.is_err());
    }

    #[test]
    fn test_custom_schedule() {
        let store = EventStore::new();
        let schedule = Schedule::Custom(Arc::new(|_time| true)); // Always allowed

        let result = store.limit().during(schedule).check("action");
        assert!(result.is_ok());
    }

    #[test]
    fn test_schedule_hours_constructor() {
        let schedule = Schedule::hours(9, 17).unwrap();
        match schedule {
            Schedule::Hours {
                start_hour,
                end_hour,
                ..
            } => {
                assert_eq!(start_hour, 9);
                assert_eq!(end_hour, 17);
            }
            _ => panic!("Expected Hours variant"),
        }
    }

    #[test]
    fn test_schedule_hours_validates_start_hour() {
        let result = Schedule::hours(24, 17);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::InvalidHour(hour) => assert_eq!(hour, 24),
            _ => panic!("Expected InvalidHour error"),
        }
    }

    #[test]
    fn test_schedule_hours_validates_end_hour() {
        let result = Schedule::hours(9, 25);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::InvalidHour(hour) => assert_eq!(hour, 25),
            _ => panic!("Expected InvalidHour error"),
        }
    }

    #[test]
    fn test_schedule_hours_validates_range() {
        let result = Schedule::hours(17, 9);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::InvalidRange(start, end) => {
                assert_eq!(start, 17);
                assert_eq!(end, 9);
            }
            _ => panic!("Expected InvalidRange error"),
        }
    }

    #[test]
    fn test_schedule_hours_rejects_equal_hours() {
        let result = Schedule::hours(9, 9);
        assert!(result.is_err());
        match result.unwrap_err() {
            Error::InvalidRange(start, end) => {
                assert_eq!(start, 9);
                assert_eq!(end, 9);
            }
            _ => panic!("Expected InvalidRange error"),
        }
    }

    #[test]
    fn test_schedule_hours_accepts_valid_boundary_cases() {
        // 0-23 should work
        let schedule = Schedule::hours(0, 23).unwrap();
        match schedule {
            Schedule::Hours {
                start_hour,
                end_hour,
                ..
            } => {
                assert_eq!(start_hour, 0);
                assert_eq!(end_hour, 23);
            }
            _ => panic!("Expected Hours variant"),
        }

        // Single hour difference should work
        let schedule = Schedule::hours(0, 1).unwrap();
        match schedule {
            Schedule::Hours {
                start_hour,
                end_hour,
                ..
            } => {
                assert_eq!(start_hour, 0);
                assert_eq!(end_hour, 1);
            }
            _ => panic!("Expected Hours variant"),
        }
    }

    #[test]
    fn test_retry_after_calculation_for_at_most() {
        let store = EventStore::new();
        store.record_count("test_event", 10);

        let result = store
            .limit()
            .at_most("test_event", 5, TimeUnit::Hours)
            .check("action");

        assert!(result.is_err());
        let Error::LimitExceeded(err) = result.unwrap_err() else {
            panic!("Expected LimitExceeded error");
        };
        assert_eq!(err.retry_after, Some(Duration::hours(1)));
    }

    #[test]
    fn test_cooldown_never_seen_passes() {
        let store = EventStore::new();

        let result = store
            .limit()
            .cooldown("never_seen", Duration::hours(1))
            .check("action");

        assert!(result.is_ok());
    }

    #[test]
    fn test_flexible_time_windows_with_tuple() {
        let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let clock = TestClock::build_for_testing_at(start_time);

        let store = EventStore::builder()
            .with_clock(Arc::new(clock.clone()))
            .build()
            .unwrap();

        // Record 10 events over 3 days
        for _i in 0..10 {
            store.record("api");
            clock.advance(Duration::hours(7));
        }

        // Should pass - 10 recorded, limit 11 means we can still record more
        let result = store
            .limit()
            .at_most("api", 11, (7, TimeUnit::Days))
            .check("api");
        assert!(result.is_ok());

        // Should fail - 10 >= 10 in last 7 days
        let result = store
            .limit()
            .at_most("api", 10, (7, TimeUnit::Days))
            .check("api");
        assert!(result.is_err());
    }

    #[test]
    fn test_flexible_time_windows_with_duration() {
        let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let clock = TestClock::build_for_testing_at(start_time);

        let store = EventStore::builder()
            .with_clock(Arc::new(clock.clone()))
            .build()
            .unwrap();

        // Record 10 events over 3 days
        for _i in 0..10 {
            store.record("api");
            clock.advance(Duration::hours(7));
        }

        // Duration syntax - should pass with limit 11
        let result = store
            .limit()
            .at_most("api", 11, Duration::days(7))
            .check("api");
        assert!(result.is_ok());

        // Duration syntax - should fail with limit 10
        let result = store
            .limit()
            .at_most("api", 10, Duration::days(7))
            .check("api");
        assert!(result.is_err());
    }

    #[test]
    fn test_flexible_time_windows_backward_compatible() {
        let store = EventStore::new();
        store.record_count("api", 5);

        // Old syntax should still work - TimeUnit means 1 unit
        let result = store
            .limit()
            .at_most("api", 10, TimeUnit::Days)
            .check("api");
        assert!(result.is_ok());
    }

    #[test]
    fn test_flexible_time_windows_30_minutes() {
        let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let clock = TestClock::build_for_testing_at(start_time);

        let store = EventStore::builder()
            .with_clock(Arc::new(clock.clone()))
            .build()
            .unwrap();

        // Record 10 events over 20 minutes
        for _i in 0..10 {
            store.record("api");
            clock.advance(Duration::minutes(2));
        }

        // Should pass - 10 recorded, limit 11
        let result = store
            .limit()
            .at_most("api", 11, (30, TimeUnit::Minutes))
            .check("api");
        assert!(result.is_ok());

        // Should fail - 10 >= 10 in last 30 minutes
        let result = store
            .limit()
            .at_most("api", 10, (30, TimeUnit::Minutes))
            .check("api");
        assert!(result.is_err());
    }

    #[test]
    fn test_at_least_with_flexible_windows() {
        let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let clock = TestClock::build_for_testing_at(start_time);

        let store = EventStore::builder()
            .with_clock(Arc::new(clock.clone()))
            .build()
            .unwrap();

        // Record 10 events over 3 days
        for _i in 0..10 {
            store.record("prerequisite");
            clock.advance(Duration::hours(7));
        }

        // Should pass - at least 5 in last 7 days
        let result = store
            .limit()
            .at_least("prerequisite", 5, (7, TimeUnit::Days))
            .check("action");
        assert!(result.is_ok());

        // Should fail - need at least 15 in last 7 days
        let result = store
            .limit()
            .at_least("prerequisite", 15, (7, TimeUnit::Days))
            .check("action");
        assert!(result.is_err());
    }

    #[test]
    fn test_duration_conversion_to_time_window() {
        let start_time = Utc.with_ymd_and_hms(2025, 1, 1, 0, 0, 0).unwrap();
        let clock = TestClock::build_for_testing_at(start_time);

        let store = EventStore::builder()
            .with_clock(Arc::new(clock.clone()))
            .build()
            .unwrap();

        // Record events
        for _i in 0..5 {
            store.record("api");
            clock.advance(Duration::hours(1));
        }

        // Test Duration::hours conversion
        let result = store
            .limit()
            .at_most("api", 10, Duration::hours(12))
            .check("api");
        assert!(result.is_ok());

        // Test Duration::minutes conversion
        store.record_count("fast", 20);
        let result = store
            .limit()
            .at_most("fast", 25, Duration::minutes(30))
            .check("fast");
        assert!(result.is_ok());
    }

    #[test]
    fn test_reservation_prevents_race_condition() {
        let store = EventStore::new();

        // Reserve first slot
        let res1 = store
            .limit()
            .at_most("api", 1, TimeUnit::Hours)
            .reserve("api")
            .unwrap();

        // Second reservation should fail (pending is counted)
        let res2 = store
            .limit()
            .at_most("api", 1, TimeUnit::Hours)
            .reserve("api");
        assert!(res2.is_err());

        // Commit first reservation
        res1.commit();

        // Now recorded
        assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
    }

    #[test]
    fn test_reservation_auto_cancel_on_drop() {
        let store = EventStore::new();

        {
            let _res = store
                .limit()
                .at_most("api", 10, TimeUnit::Hours)
                .reserve("api")
                .unwrap();
            // Drop without commit
        }

        // Nothing recorded (but counter was created, so it returns Some(0))
        let sum = store.query("api").last_hours(1).sum();
        assert!(sum.is_none() || sum == Some(0));

        // Can reserve again (pending was released)
        let res = store
            .limit()
            .at_most("api", 1, TimeUnit::Hours)
            .reserve("api");
        assert!(res.is_ok());
    }

    #[test]
    fn test_reservation_explicit_cancel() {
        let store = EventStore::new();

        let res = store
            .limit()
            .at_most("api", 10, TimeUnit::Hours)
            .reserve("api")
            .unwrap();

        // Explicitly cancel
        res.cancel();

        // Nothing recorded (but counter was created, so it returns Some(0))
        let sum = store.query("api").last_hours(1).sum();
        assert!(sum.is_none() || sum == Some(0));

        // Can reserve again
        let res = store
            .limit()
            .at_most("api", 1, TimeUnit::Hours)
            .reserve("api");
        assert!(res.is_ok());
    }

    #[test]
    fn test_reservation_commit_records_event() {
        let store = EventStore::new();

        let res = store
            .limit()
            .at_most("api", 10, TimeUnit::Hours)
            .reserve("api")
            .unwrap();

        // Commit the reservation
        res.commit();

        // Event is recorded
        assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
    }

    #[test]
    fn test_reservation_concurrent_limits() {
        let store = Arc::new(EventStore::new());

        // Try to reserve 20 times with limit of 10
        let handles: Vec<_> = (0..20)
            .map(|_| {
                let store = store.clone();
                thread::spawn(move || {
                    store
                        .limit()
                        .at_most("api", 10, TimeUnit::Hours)
                        .reserve("api")
                })
            })
            .collect();

        let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect();

        // Exactly 10 should succeed
        let successes: Vec<_> = results.into_iter().filter(|r| r.is_ok()).collect();
        assert_eq!(successes.len(), 10);

        // Commit all successful reservations
        for result in successes {
            result.unwrap().commit();
        }

        // Exactly 10 recorded
        assert_eq!(store.query("api").last_hours(1).sum(), Some(10));
    }

    #[test]
    fn test_reservation_transactional_pattern() {
        fn simulated_work(should_succeed: bool) -> std::result::Result<(), &'static str> {
            if should_succeed {
                Ok(())
            } else {
                Err("Work failed")
            }
        }

        let store = EventStore::new();

        // Success case
        let res = store
            .limit()
            .at_most("api", 10, TimeUnit::Hours)
            .reserve("api")
            .unwrap();

        match simulated_work(true) {
            Ok(_) => res.commit(),
            Err(_) => res.cancel(),
        }

        // Event recorded
        assert_eq!(store.query("api").last_hours(1).sum(), Some(1));

        // Failure case
        let res = store
            .limit()
            .at_most("api", 10, TimeUnit::Hours)
            .reserve("api")
            .unwrap();

        match simulated_work(false) {
            Ok(_) => res.commit(),
            Err(_) => res.cancel(),
        }

        // Still only 1 recorded (second was cancelled)
        assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
    }

    #[test]
    fn test_reservation_with_multiple_constraints() {
        let store = EventStore::new();

        // Set up prerequisite
        store.record_count("prerequisite", 5);

        // Reserve with multiple constraints
        let res = store
            .limit()
            .at_most("api", 10, TimeUnit::Hours)
            .at_least("prerequisite", 3, TimeUnit::Days)
            .reserve("api")
            .unwrap();

        res.commit();

        // Verify event recorded
        assert_eq!(store.query("api").last_hours(1).sum(), Some(1));
    }
}