//! Built-in function registry for formula expressions.

use super::FormulaError;

#[derive(Debug, Clone)]
pub struct FunctionDef {
    pub min_args: usize,
    pub max_args: usize,
    pub evaluator: fn(&[f64]) -> Result<f64, FormulaError>,
}

#[must_use]
pub fn get_function(name: &str) -> Option<FunctionDef> {
    let upper = name.to_ascii_uppercase();
    match upper.as_str() {
        "SUM" => Some(FunctionDef {
            min_args: 1,
            max_args: usize::MAX,
            evaluator: |args| Ok(args.iter().sum()),
        }),
        "AVG" => Some(FunctionDef {
            min_args: 1,
            max_args: usize::MAX,
            evaluator: |args| Ok(args.iter().sum::<f64>() / args.len() as f64),
        }),
        "AVERAGE" => Some(FunctionDef {
            min_args: 1,
            max_args: usize::MAX,
            evaluator: |args| Ok(args.iter().sum::<f64>() / args.len() as f64),
        }),
        "MEDIAN" => Some(FunctionDef {
            min_args: 1,
            max_args: usize::MAX,
            evaluator: |args| {
                let mut values = args.to_vec();
                values.sort_by(|a, b| a.total_cmp(b));
                let middle = values.len() / 2;
                if values.len() % 2 == 0 {
                    Ok((values[middle - 1] + values[middle]) / 2.0)
                } else {
                    Ok(values[middle])
                }
            },
        }),
        "STDEV" => Some(FunctionDef {
            min_args: 2,
            max_args: usize::MAX,
            evaluator: |args| {
                let mean = args.iter().sum::<f64>() / args.len() as f64;
                let variance = args.iter().map(|value| (value - mean).powi(2)).sum::<f64>()
                    / (args.len() - 1) as f64;
                Ok(variance.sqrt())
            },
        }),
        "PERCENTILE" => Some(FunctionDef {
            min_args: 2,
            max_args: usize::MAX,
            evaluator: |args| {
                let percentile = args[args.len() - 1];
                if !(0.0..=1.0).contains(&percentile) {
                    return Err(FormulaError::MathError(
                        "PERCENTILE requires a percentile between 0 and 1".to_string(),
                    ));
                }
                let mut values = args[..args.len() - 1].to_vec();
                values.sort_by(|a, b| a.total_cmp(b));
                let rank = percentile * (values.len() - 1) as f64;
                let lower = rank.floor() as usize;
                let upper = rank.ceil() as usize;
                if lower == upper {
                    Ok(values[lower])
                } else {
                    let weight = rank - lower as f64;
                    Ok(values[lower] * (1.0 - weight) + values[upper] * weight)
                }
            },
        }),
        "MIN" => Some(FunctionDef {
            min_args: 1,
            max_args: usize::MAX,
            evaluator: |args| Ok(args.iter().fold(f64::INFINITY, |a, b| a.min(*b))),
        }),
        "MAX" => Some(FunctionDef {
            min_args: 1,
            max_args: usize::MAX,
            evaluator: |args| Ok(args.iter().fold(f64::NEG_INFINITY, |a, b| a.max(*b))),
        }),
        "ABS" => Some(FunctionDef {
            min_args: 1,
            max_args: 1,
            evaluator: |args| Ok(args[0].abs()),
        }),
        "ROUND" => Some(FunctionDef {
            min_args: 1,
            max_args: 2,
            evaluator: |args| {
                let value = args[0];
                let decimals = if args.len() > 1 { args[1] } else { 0.0 };
                let multiplier = 10_f64.powi(decimals as i32);
                Ok((value * multiplier).round() / multiplier)
            },
        }),
        "IF" => Some(FunctionDef {
            min_args: 3,
            max_args: 3,
            evaluator: |args| {
                Ok(if args[0].abs() > f64::EPSILON {
                    args[1]
                } else {
                    args[2]
                })
            },
        }),
        "SQRT" => Some(FunctionDef {
            min_args: 1,
            max_args: 1,
            evaluator: |args| {
                if args[0] < 0.0 {
                    Err(FormulaError::MathError(
                        "SQRT of negative number".to_string(),
                    ))
                } else {
                    Ok(args[0].sqrt())
                }
            },
        }),
        "POW" => Some(FunctionDef {
            min_args: 2,
            max_args: 2,
            evaluator: |args| Ok(args[0].powf(args[1])),
        }),
        "COMPOUND" => Some(FunctionDef {
            min_args: 3,
            max_args: 3,
            evaluator: |args| Ok(args[0] * (1.0 + args[1]).powf(args[2])),
        }),
        "GROWTH" => Some(FunctionDef {
            min_args: 2,
            max_args: 2,
            evaluator: |args| {
                if args[0].abs() < f64::EPSILON {
                    Err(FormulaError::DivisionByZero)
                } else {
                    Ok(args[1] / args[0] - 1.0)
                }
            },
        }),
        "CAGR" => Some(FunctionDef {
            min_args: 3,
            max_args: 3,
            evaluator: |args| {
                if args[0] <= 0.0 || args[2] <= 0.0 {
                    Err(FormulaError::MathError(
                        "CAGR requires positive beginning value and periods".to_string(),
                    ))
                } else {
                    Ok((args[1] / args[0]).powf(1.0 / args[2]) - 1.0)
                }
            },
        }),
        "NPV" => Some(FunctionDef {
            min_args: 2,
            max_args: usize::MAX,
            evaluator: |args| {
                let rate = args[0];
                Ok(args[1..]
                    .iter()
                    .enumerate()
                    .map(|(index, value)| value / (1.0 + rate).powi((index + 1) as i32))
                    .sum())
            },
        }),
        "IRR" => Some(FunctionDef {
            min_args: 2,
            max_args: usize::MAX,
            evaluator: |args| {
                let mut rate: f64 = 0.1;
                for _ in 0..100 {
                    let mut value = 0.0;
                    let mut derivative = 0.0;
                    for (period, cash_flow) in args.iter().enumerate() {
                        let exponent = period as i32;
                        value += cash_flow / (1.0 + rate).powi(exponent);
                        if period > 0 {
                            derivative -=
                                period as f64 * cash_flow / (1.0 + rate).powi(exponent + 1);
                        }
                    }
                    if derivative.abs() < f64::EPSILON {
                        break;
                    }
                    let next = rate - value / derivative;
                    if (next - rate).abs() < 1e-7 {
                        return Ok(next);
                    }
                    rate = next;
                }
                Ok(rate)
            },
        }),
        "WACC" => Some(FunctionDef {
            min_args: 5,
            max_args: 5,
            evaluator: |args| {
                let equity = args[0];
                let debt = args[1];
                let cost_of_equity = args[2];
                let cost_of_debt = args[3];
                let tax_rate = args[4];
                let capital = equity + debt;
                if capital.abs() < f64::EPSILON {
                    return Err(FormulaError::DivisionByZero);
                }
                Ok((equity / capital) * cost_of_equity
                    + (debt / capital) * cost_of_debt * (1.0 - tax_rate))
            },
        }),
        "FLOOR" => Some(FunctionDef {
            min_args: 1,
            max_args: 1,
            evaluator: |args| Ok(args[0].floor()),
        }),
        "CEIL" => Some(FunctionDef {
            min_args: 1,
            max_args: 1,
            evaluator: |args| Ok(args[0].ceil()),
        }),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_sum_function() {
        let func = get_function("SUM").unwrap();
        let result = (func.evaluator)(&[1.0, 2.0, 3.0, 4.0]);
        assert_eq!(result.unwrap(), 10.0);
    }

    #[test]
    fn test_avg_function() {
        let func = get_function("AVG").unwrap();
        let result = (func.evaluator)(&[2.0, 4.0, 6.0]);
        assert_eq!(result.unwrap(), 4.0);

        let func = get_function("AVERAGE").unwrap();
        let result = (func.evaluator)(&[2.0, 4.0, 6.0]);
        assert_eq!(result.unwrap(), 4.0);
    }

    #[test]
    fn test_financial_functions() {
        let cagr = get_function("CAGR").unwrap();
        let result = (cagr.evaluator)(&[100.0, 121.0, 2.0]).unwrap();
        assert!((result - 0.1).abs() < 1e-10);

        let npv = get_function("NPV").unwrap();
        let result = (npv.evaluator)(&[0.1, 110.0, 121.0]).unwrap();
        assert!((result - 200.0).abs() < 1e-10);

        let wacc = get_function("WACC").unwrap();
        let result = (wacc.evaluator)(&[80.0, 20.0, 0.1, 0.05, 0.2]).unwrap();
        assert!((result - 0.088).abs() < 1e-10);
    }

    #[test]
    fn test_statistical_functions() {
        let median = get_function("MEDIAN").unwrap();
        assert_eq!((median.evaluator)(&[3.0, 1.0, 2.0]).unwrap(), 2.0);

        let percentile = get_function("PERCENTILE").unwrap();
        let result = (percentile.evaluator)(&[10.0, 20.0, 30.0, 0.5]).unwrap();
        assert_eq!(result, 20.0);
    }

    #[test]
    fn test_min_function() {
        let func = get_function("MIN").unwrap();
        let result = (func.evaluator)(&[5.0, 2.0, 8.0, 1.0]);
        assert_eq!(result.unwrap(), 1.0);
    }

    #[test]
    fn test_max_function() {
        let func = get_function("MAX").unwrap();
        let result = (func.evaluator)(&[5.0, 2.0, 8.0, 1.0]);
        assert_eq!(result.unwrap(), 8.0);
    }

    #[test]
    fn test_abs_function() {
        let func = get_function("ABS").unwrap();
        assert_eq!((func.evaluator)(&[5.0]).unwrap(), 5.0);
        assert_eq!((func.evaluator)(&[-5.0]).unwrap(), 5.0);
    }

    #[test]
    fn test_round_function() {
        let func = get_function("ROUND").unwrap();
        assert_eq!((func.evaluator)(&[5.5678, 2.0]).unwrap(), 5.57);
        assert_eq!((func.evaluator)(&[5.5678]).unwrap(), 6.0);
    }

    #[test]
    fn test_if_function() {
        let func = get_function("IF").unwrap();
        assert_eq!((func.evaluator)(&[1.0, 10.0, 20.0]).unwrap(), 10.0);
        assert_eq!((func.evaluator)(&[0.0, 10.0, 20.0]).unwrap(), 20.0);
        assert_eq!((func.evaluator)(&[-5.0, 10.0, 20.0]).unwrap(), 10.0);
    }

    #[test]
    fn test_sqrt_function() {
        let func = get_function("SQRT").unwrap();
        assert_eq!((func.evaluator)(&[16.0]).unwrap(), 4.0);
        assert_eq!((func.evaluator)(&[2.0]).unwrap(), 2.0_f64.sqrt());
        assert!((func.evaluator)(&[-1.0]).is_err());
    }

    #[test]
    fn test_case_insensitive_lookup() {
        assert!(get_function("sum").is_some());
        assert!(get_function("Sum").is_some());
        assert!(get_function("SUM").is_some());
        assert!(get_function("unknown").is_none());
    }
}
