use std::collections::HashMap;

use chrono::NaiveDate;

#[derive(Debug, Clone)]
pub struct TradingCalendar {
    days: Vec<NaiveDate>,
    index: HashMap<NaiveDate, usize>,
}

impl TradingCalendar {
    pub fn new(mut days: Vec<NaiveDate>) -> Self {
        days.sort_unstable();
        days.dedup();

        let index = days
            .iter()
            .copied()
            .enumerate()
            .map(|(idx, day)| (day, idx))
            .collect();

        Self { days, index }
    }

    pub fn days(&self) -> &[NaiveDate] {
        &self.days
    }

    pub fn iter(&self) -> impl Iterator<Item = NaiveDate> + '_ {
        self.days.iter().copied()
    }

    pub fn len(&self) -> usize {
        self.days.len()
    }

    pub fn is_empty(&self) -> bool {
        self.days.is_empty()
    }

    pub fn index_of(&self, date: NaiveDate) -> Option<usize> {
        self.index.get(&date).copied()
    }

    pub fn previous_day(&self, date: NaiveDate) -> Option<NaiveDate> {
        let idx = self.index_of(date)?;
        idx.checked_sub(1)
            .and_then(|prev| self.days.get(prev).copied())
    }

    pub fn previous_trading_date(&self, date: NaiveDate, n: usize) -> Option<NaiveDate> {
        if n == 0 {
            return None;
        }
        let before_count = match self.days.binary_search(&date) {
            Ok(idx) => idx,
            Err(idx) => idx,
        };
        before_count
            .checked_sub(n)
            .and_then(|idx| self.days.get(idx).copied())
    }

    pub fn next_trading_date(&self, date: NaiveDate, n: usize) -> Option<NaiveDate> {
        if n == 0 {
            return None;
        }
        let first_after = match self.days.binary_search(&date) {
            Ok(idx) => idx.saturating_add(1),
            Err(idx) => idx,
        };
        first_after
            .checked_add(n.saturating_sub(1))
            .and_then(|idx| self.days.get(idx).copied())
    }

    pub fn trading_dates(&self, start: NaiveDate, end: NaiveDate) -> Vec<NaiveDate> {
        if start > end {
            return Vec::new();
        }
        self.days
            .iter()
            .copied()
            .filter(|date| *date >= start && *date <= end)
            .collect()
    }

    pub fn trailing_days(&self, end: NaiveDate, lookback: usize) -> Vec<NaiveDate> {
        let Some(end_idx) = self.index_of(end) else {
            return Vec::new();
        };
        let start = end_idx.saturating_add(1).saturating_sub(lookback);
        self.days[start..=end_idx].to_vec()
    }
}