use chrono::{Datelike, Duration, NaiveDate, Weekday};
use std::error::Error;

#[derive(Debug, Clone, Copy)]
pub enum BDateFreq {
    Daily,
    WeeklyMonday,
    MonthStart,
    QuarterStart,
    YearStart,
    MonthEnd,
    QuarterEnd,
    WeeklyFriday,
    YearEnd,
}

#[derive(Debug, Clone, Copy)]
pub enum AggregationType {
    Start, // Indicates picking the first date in a group.
    End,   // Indicates picking the last date in a group.
}

impl BDateFreq {
    pub fn from_string(freq: String) -> Result<Self, Box<dyn Error>> {
        // use from_str to convert the string to a BDateFreq enum
        Self::from_str(&freq)
    }
    pub fn from_str(freq: &str) -> Result<Self, Box<dyn Error>> {
        match freq {
            "D" => Ok(BDateFreq::Daily),
            "W" => Ok(BDateFreq::WeeklyMonday),
            "M" => Ok(BDateFreq::MonthStart),
            "Q" => Ok(BDateFreq::QuarterStart),
            "A" => Ok(BDateFreq::YearStart),
            "ME" => Ok(BDateFreq::MonthEnd),
            "QE" => Ok(BDateFreq::QuarterEnd),
            "WF" => Ok(BDateFreq::WeeklyFriday),
            "YE" => Ok(BDateFreq::YearEnd),
            _ => Err("Invalid frequency specified".into()),
        }
    }
    pub fn agg_type(&self) -> AggregationType {
        match self {
            BDateFreq::Daily
            | BDateFreq::WeeklyMonday
            | BDateFreq::MonthStart
            | BDateFreq::QuarterStart
            | BDateFreq::YearStart => AggregationType::Start,

            BDateFreq::WeeklyFriday
            | BDateFreq::MonthEnd
            | BDateFreq::QuarterEnd
            | BDateFreq::YearEnd => AggregationType::End,
        }
    }
}

/// Returns only the business dates (Mon-Fri) between start_date and end_date
/// that match the desired frequency.
pub fn get_bdates_list_with_freq(
    start_date_str: &str,
    end_date_str: &str,
    freq: BDateFreq,
) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
    let start_date = NaiveDate::parse_from_str(start_date_str, "%Y-%m-%d")?;
    let end_date = NaiveDate::parse_from_str(end_date_str, "%Y-%m-%d")?;

    if start_date > end_date {
        return Ok(Vec::new());
    }

    let mut dates = match freq {
        BDateFreq::Daily => collect_daily(start_date, end_date),

        BDateFreq::WeeklyMonday => collect_weekly(start_date, end_date, Weekday::Mon),
        BDateFreq::WeeklyFriday => collect_weekly(start_date, end_date, Weekday::Fri),

        BDateFreq::MonthStart => collect_monthly(start_date, end_date, /*start=*/ true),
        BDateFreq::MonthEnd => collect_monthly(start_date, end_date, /*start=*/ false),

        BDateFreq::QuarterStart => collect_quarterly(start_date, end_date, /*start=*/ true),
        BDateFreq::QuarterEnd => collect_quarterly(start_date, end_date, /*start=*/ false),

        BDateFreq::YearStart => collect_yearly(start_date, end_date, /*start=*/ true),
        BDateFreq::YearEnd => collect_yearly(start_date, end_date, /*start=*/ false),
    };

    // Filter out any weekend days that might slip in edge cases (e.g. if the
    // computed "start of month" fell on Sat/Sun).
    dates.retain(|d| d.weekday() != Weekday::Sat && d.weekday() != Weekday::Sun);

    Ok(dates)
}

/* ------------------------------ Helpers ------------------------------ */

/// Return all business days, day-by-day.
fn collect_daily(start_date: NaiveDate, end_date: NaiveDate) -> Vec<NaiveDate> {
    let mut result = Vec::new();
    let mut current = start_date;
    while current <= end_date {
        if is_weekday(current) {
            result.push(current);
        }
        current = current.succ_opt().unwrap();
    }
    result
}

/// Return the specified weekday (e.g. Monday, Friday) in each week of the range.
fn collect_weekly(
    start_date: NaiveDate,
    end_date: NaiveDate,
    target_weekday: Weekday,
) -> Vec<NaiveDate> {
    let mut result = Vec::new();

    // Find the first `target_weekday` on or after `start_date`.
    // If `start_date` is already e.g. Monday, we can use it as is.
    // Otherwise, jump ahead until we get that weekday.
    let mut current = move_to_weekday_on_or_after(start_date, target_weekday);

    // Step in 7-day increments (full weeks).
    while current <= end_date {
        result.push(current);
        current = current + Duration::days(7);
    }
    result
}

/// Return either first or last business day in each month of the range.
fn collect_monthly(
    start_date: NaiveDate,
    end_date: NaiveDate,
    want_first_day: bool,
) -> Vec<NaiveDate> {
    let mut result = Vec::new();

    // We'll iterate month by month, from (start_year, start_month) up to
    // (end_year, end_month).
    let mut year = start_date.year();
    let mut month = start_date.month();

    // A small helper that updates year/month by +1 month.
    let next_month = |(yr, mo): (i32, u32)| -> (i32, u32) {
        if mo == 12 {
            (yr + 1, 1)
        } else {
            (yr, mo + 1)
        }
    };

    // Move `(year, month)` backward if necessary so that `(year, month)`
    // definitely covers the entire period from `start_date` onward.
    // Actually, it’s simpler to start from the actual (year, month) of start_date
    // and go up. We'll just skip if the computed "day" < start_date.

    // Continue while we haven't passed (end_year, end_month).
    while year < end_date.year() || (year == end_date.year() && month <= end_date.month()) {
        // Compute the date that represents either first or last business day
        // for this (year, month).
        let candidate = if want_first_day {
            first_business_day_of_month(year, month)
        } else {
            last_business_day_of_month(year, month)
        };
        if candidate >= start_date && candidate <= end_date {
            result.push(candidate);
        }

        // Move to the next month.
        let (ny, nm) = next_month((year, month));
        year = ny;
        month = nm;
    }

    result
}

/// Return either the first or last business day in each quarter of the range.
fn collect_quarterly(
    start_date: NaiveDate,
    end_date: NaiveDate,
    want_first_day: bool,
) -> Vec<NaiveDate> {
    let mut result = Vec::new();

    // We'll figure out which quarter `start_date` is in, then jump quarter-by-quarter.
    // Quarters are: Q1 = months 1–3, Q2 = 4–6, Q3 = 7–9, Q4 = 10–12.
    // Start by computing the (year, quarter_index) for start_date.
    let mut year = start_date.year();
    let mut q = month_to_quarter(start_date.month());

    while quarter_to_first_date(year, q) > end_date {
        // If even the earliest day in that quarter is > end_date, we’re done.
        return result;
    }

    // Move backward if the quarter’s last day < start_date, etc.
    // But simpler: we’ll do a loop that increments quarter by quarter, and
    // pick the appropriate date each time. We break when we pass end_date.

    loop {
        // For the current year+quarter, compute the date that’s either the first or last
        // business day of that quarter:
        let candidate = if want_first_day {
            first_business_day_of_quarter(year, q)
        } else {
            last_business_day_of_quarter(year, q)
        };

        if candidate > end_date {
            break;
        }
        if candidate >= start_date {
            result.push(candidate);
        }

        // Move to next quarter.
        if q == 4 {
            year += 1;
            q = 1;
        } else {
            q += 1;
        }
    }

    result
}

/// Return either the first or last business day in each year of the range.
fn collect_yearly(
    start_date: NaiveDate,
    end_date: NaiveDate,
    want_first_day: bool,
) -> Vec<NaiveDate> {
    let mut result = Vec::new();
    let mut year = start_date.year();

    // Step year-by-year from `start_date.year()` up to `end_date.year()`.
    while year <= end_date.year() {
        let candidate = if want_first_day {
            first_business_day_of_year(year)
        } else {
            last_business_day_of_year(year)
        };
        if candidate >= start_date && candidate <= end_date {
            result.push(candidate);
        }
        year += 1;
    }
    result
}

/* ---------------------- Low-Level Utility Functions ---------------------- */

/// Is this a weekday (Mon-Fri)?
fn is_weekday(date: NaiveDate) -> bool {
    match date.weekday() {
        Weekday::Sat | Weekday::Sun => false,
        _ => true,
    }
}

/// Given a date and a `target_weekday`, returns the date that is the first
/// `target_weekday` on or after the given date.
fn move_to_weekday_on_or_after(date: NaiveDate, target: Weekday) -> NaiveDate {
    let mut current = date;
    while current.weekday() != target {
        current = current.succ_opt().unwrap();
    }
    current
}

/// Return the earliest business day of (year, month).
fn first_business_day_of_month(year: i32, month: u32) -> NaiveDate {
    // Start with the 1st of the month.
    let mut d = NaiveDate::from_ymd_opt(year, month, 1).expect("invalid year-month");
    // If it’s Sat/Sun, move forward until we get a weekday.
    while !is_weekday(d) {
        d = d.succ_opt().unwrap();
    }
    d
}

/// Return the latest business day of (year, month).
fn last_business_day_of_month(year: i32, month: u32) -> NaiveDate {
    let last_dom = days_in_month(year, month);
    let mut d = NaiveDate::from_ymd_opt(year, month, last_dom).expect("invalid year-month");
    // If it’s Sat/Sun, move backward until we get a weekday.
    while !is_weekday(d) {
        d = d.pred_opt().unwrap();
    }
    d
}

/// Number of days in a month (not considering leap years *beyond* chrono's normal handling).
fn days_in_month(year: i32, month: u32) -> u32 {
    // Chrono can handle this if we do a little trick:
    // Construct the 1st of the next month, then subtract 1 day.
    // For example:
    // if month == 12 => next = (year+1, 1, 1)
    // else => next = (year, month+1, 1)
    let (ny, nm) = if month == 12 {
        (year + 1, 1)
    } else {
        (year, month + 1)
    };
    let first_of_next = NaiveDate::from_ymd_opt(ny, nm, 1).unwrap();
    let last_of_this = first_of_next.pred_opt().unwrap();
    last_of_this.day()
}

/// Convert a month (1..12) to a quarter (1..4).
fn month_to_quarter(m: u32) -> u32 {
    (m - 1) / 3 + 1
}

/// Returns 1st day of a given (year, quarter).
fn quarter_to_first_date(year: i32, quarter: u32) -> NaiveDate {
    let month = match quarter {
        1 => 1,
        2 => 4,
        3 => 7,
        4 => 10,
        _ => panic!("invalid quarter"),
    };
    NaiveDate::from_ymd_opt(year, month, 1).unwrap()
}

/// Return the earliest business day in (year, quarter).
fn first_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate {
    let mut d = quarter_to_first_date(year, quarter);
    while !is_weekday(d) {
        d = d.succ_opt().unwrap();
    }
    d
}

/// Return the last business day in (year, quarter).
fn last_business_day_of_quarter(year: i32, quarter: u32) -> NaiveDate {
    // The last month in the quarter is quarter_to_first_date(...) + 2 months
    // Then we find the last day of that month.
    let start = quarter_to_first_date(year, quarter);
    let last_month = start.month() + 2; // e.g. Q1 => month=1 => +2=3 => March
    last_business_day_of_month(year, last_month)
}

/// Returns Jan 1st of a given year (adjust if weekend).
fn first_business_day_of_year(year: i32) -> NaiveDate {
    let mut d = NaiveDate::from_ymd_opt(year, 1, 1).unwrap();
    while !is_weekday(d) {
        d = d.succ_opt().unwrap();
    }
    d
}

/// Returns Dec 31st of a given year (adjust if weekend).
fn last_business_day_of_year(year: i32) -> NaiveDate {
    let mut d = NaiveDate::from_ymd_opt(year, 12, 31).unwrap();
    while !is_weekday(d) {
        d = d.pred_opt().unwrap();
    }
    d
}
