use chrono::NaiveDate;
use holiday_engine::{get_china_holidays, is_weekend, HolidayInfo};
use serde::Serialize;
use shift_algorithm::{get_shift_type_for_date, ShiftCycleConfig, ShiftType};
use std::collections::{HashMap, HashSet};
#[derive(Debug, Clone)]
struct DayStatus {
date: NaiveDate,
is_rest: bool,
is_holiday: bool,
is_weekend: bool,
is_adjusted_work_day: bool,
holiday_name: Option<String>,
}
impl DayStatus {
fn is_off(&self) -> bool {
self.is_rest
|| (self.is_holiday && !self.is_adjusted_work_day)
|| (self.is_weekend && !self.is_adjusted_work_day)
}
}
#[derive(Debug, Clone, Serialize)]
pub struct LeaveStrategy {
pub leave_days: u32,
pub total_break_days: u32,
pub leave_dates: Vec<NaiveDate>,
pub break_start: NaiveDate,
pub break_end: NaiveDate,
pub holiday_overlap: u32,
pub weekend_overlap: u32,
pub overlapping_holiday_names: Vec<String>,
pub efficiency: f64,
pub score: f64,
}
fn build_daily_status(
start_date: NaiveDate,
days: u32,
team_phase_offset: u32,
config: &ShiftCycleConfig,
holidays: &HashMap<NaiveDate, HolidayInfo>,
) -> Vec<DayStatus> {
(0..days)
.map(|offset| {
let date = start_date + chrono::Duration::days(offset as i64);
let shift_type = get_shift_type_for_date(date, config, team_phase_offset);
let is_rest = matches!(shift_type, ShiftType::Rest | ShiftType::Study);
let holiday_info = holidays.get(&date);
let is_holiday = holiday_info.is_some_and(|h| h.is_holiday);
let is_adjusted_work_day = holiday_info.is_some_and(|h| !h.is_holiday);
let holiday_name = if is_holiday {
holiday_info.map(|h| h.name.to_string())
} else {
None
};
DayStatus { date, is_rest, is_holiday, is_weekend: is_weekend(date), is_adjusted_work_day, holiday_name }
})
.collect()
}
pub fn find_best_leave_plans(
today: NaiveDate,
days_to_analyze: u32,
config: &ShiftCycleConfig,
team_phase_offset: u32,
holidays: Option<&HashMap<NaiveDate, HolidayInfo>>,
max_leave_days: u32,
) -> Vec<LeaveStrategy> {
if days_to_analyze < 1 || max_leave_days < 1 {
return vec![];
}
let hols = holidays.cloned().unwrap_or_else(get_china_holidays);
let status = build_daily_status(today, days_to_analyze, team_phase_offset, config, &hols);
let n = status.len();
let mut rest_before = vec![0u32; n];
let mut rest_after = vec![0u32; n];
for i in 1..n {
rest_before[i] = if status[i - 1].is_off() { rest_before[i - 1] + 1 } else { 0 };
}
for i in (0..n - 1).rev() {
rest_after[i] = if status[i + 1].is_off() { rest_after[i + 1] + 1 } else { 0 };
}
let mut strategies: Vec<LeaveStrategy> = Vec::new();
let min_leave_days = if max_leave_days == 1 { 1 } else { 2 };
for leave_days in min_leave_days..=max_leave_days {
for start_idx in 0..=(n as i32 - leave_days as i32) {
let start_idx = start_idx as usize;
let has_shift_rest = (0..leave_days).any(|j| status[start_idx + j as usize].is_rest);
if has_shift_rest { continue; }
let left_rest = rest_before[start_idx];
let right_rest = rest_after[start_idx + leave_days as usize - 1];
let total_break = left_rest + leave_days + right_rest;
if total_break <= leave_days { continue; }
let gap_start = start_idx as i32 - left_rest as i32;
let gap_end = start_idx as i32 + leave_days as i32 - 1 + right_rest as i32;
let break_start_date = status[gap_start as usize].date;
let break_end_date = status[gap_end as usize].date;
let mut holiday_overlap = 0u32;
let mut weekend_overlap = 0u32;
let mut holiday_names: HashSet<String> = HashSet::new();
for idx in gap_start..=gap_end {
let ds = &status[idx as usize];
if ds.is_holiday { holiday_overlap += 1; if let Some(ref name) = ds.holiday_name { holiday_names.insert(name.clone()); } }
if ds.is_weekend && !ds.is_adjusted_work_day { weekend_overlap += 1; }
}
let leave_date_list: Vec<NaiveDate> = (0..leave_days)
.map(|j| status[start_idx + j as usize].date)
.collect();
let efficiency = total_break as f64 / leave_days as f64;
strategies.push(LeaveStrategy {
leave_days, total_break_days: total_break,
leave_dates: leave_date_list, break_start: break_start_date,
break_end: break_end_date, holiday_overlap, weekend_overlap,
overlapping_holiday_names: holiday_names.into_iter().collect(),
efficiency, score: 0.0,
});
}
}
let mut deduped: HashMap<String, LeaveStrategy> = HashMap::new();
for s in strategies {
let key = format!("{}_{}", s.break_start, s.break_end);
match deduped.get(&key) {
Some(existing) if existing.leave_days <= s.leave_days => {}
_ => { deduped.insert(key, s); }
}
}
if deduped.is_empty() { return vec![]; }
let deduped_list: Vec<LeaveStrategy> = deduped.into_values().collect();
let max_efficiency = deduped_list.iter().map(|s| s.efficiency).fold(0.0f64, f64::max);
let max_break = deduped_list.iter().map(|s| s.total_break_days).max().unwrap_or(1);
let max_family_bonus = deduped_list.iter().map(|s| s.holiday_overlap * 2 + s.weekend_overlap).max().unwrap_or(1).max(1);
let mut scored: Vec<LeaveStrategy> = deduped_list.into_iter().map(|s| {
let eff_score = if max_efficiency > 0.0 { s.efficiency / max_efficiency } else { 0.0 };
let len_score = if max_break > 0 { s.total_break_days as f64 / max_break as f64 } else { 0.0 };
let family_bonus = (s.holiday_overlap * 2 + s.weekend_overlap) as f64;
let fam_score = family_bonus / max_family_bonus as f64;
let score = 0.50 * eff_score + 0.25 * len_score + 0.25 * fam_score;
LeaveStrategy { score, ..s }
}).collect();
scored.sort_by(|a, b| {
b.score.partial_cmp(&a.score)
.unwrap_or(std::cmp::Ordering::Equal)
.then_with(|| a.break_start.cmp(&b.break_start))
});
scored
}
#[cfg(test)]
mod tests {
use super::*;
use shift_algorithm::cycle::default_config;
#[test]
fn returns_empty_for_zero_analyze_days() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 0, &config, 0, None, 5);
assert!(plans.is_empty());
}
#[test]
fn returns_empty_for_zero_max_leave() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 0);
assert!(plans.is_empty());
}
#[test]
fn produces_strategies_with_default_config() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 5);
assert!(!plans.is_empty());
}
#[test]
fn strategies_are_sorted_by_score_desc() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 5);
for w in plans.windows(2) {
assert!(w[0].score >= w[1].score);
}
}
#[test]
fn total_break_exceeds_leave_days() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 5);
for s in &plans {
assert!(s.total_break_days > s.leave_days);
}
}
#[test]
fn leave_dates_count_matches() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 5);
for s in &plans {
assert_eq!(s.leave_dates.len() as u32, s.leave_days);
}
}
#[test]
fn no_duplicate_break_ranges() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 180, &config, 0, None, 5);
let mut seen: HashSet<(NaiveDate, NaiveDate)> = HashSet::new();
for s in &plans {
let key = (s.break_start, s.break_end);
assert!(seen.insert(key));
}
}
#[test]
fn respects_max_leave_days() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 3);
for s in &plans {
assert!(s.leave_days <= 3);
}
}
#[test]
fn dates_within_analysis_window() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let window = 60u32;
let end = today + chrono::Duration::days(window as i64);
let plans = find_best_leave_plans(today, window, &config, 0, None, 5);
for s in &plans {
assert!(s.break_start >= today);
assert!(s.break_end < end);
}
}
#[test]
fn includes_national_day_strategies() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 9, 1).unwrap();
let plans = find_best_leave_plans(today, 60, &config, 0, None, 5);
let has_national_day = plans.iter().any(|s| {
s.overlapping_holiday_names.iter().any(|n| n.contains("国庆"))
});
assert!(has_national_day);
}
#[test]
fn different_team_produces_different_strategies() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans_team1 = find_best_leave_plans(today, 90, &config, 0, None, 5);
let plans_team2 = find_best_leave_plans(today, 90, &config, 7, None, 5);
let keys1: HashSet<String> = plans_team1.iter().map(|s| format!("{}_{}", s.break_start, s.break_end)).collect();
let keys2: HashSet<String> = plans_team2.iter().map(|s| format!("{}_{}", s.break_start, s.break_end)).collect();
assert_ne!(keys1, keys2);
}
#[test]
fn custom_cycle_produces_strategies() {
use ShiftType::*;
let config = ShiftCycleConfig {
cycle: vec![Morning, Afternoon, Rest, Night, Rest, Morning, Afternoon],
cycle_length: 7,
reference_date: shift_algorithm::cycle::default_reference_date(),
total_teams: 2,
};
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 5);
assert!(!plans.is_empty());
}
#[test]
fn efficiency_is_at_least_one() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 5);
for s in &plans { assert!(s.efficiency >= 1.0); }
}
#[test]
fn score_is_between_zero_and_one() {
let config = default_config();
let today = NaiveDate::from_ymd_opt(2026, 5, 22).unwrap();
let plans = find_best_leave_plans(today, 90, &config, 0, None, 5);
for s in &plans { assert!((0.0..=1.01).contains(&s.score)); }
}
}