use crate::{
    PolarsPlugin,
    values::{Column, CustomValueSupport, NuDataFrame, NuExpression, PolarsPluginType},
};

use nu_plugin::{EngineInterface, EvaluatedCall, PluginCommand};
use nu_protocol::{
    Category, Example, LabeledError, PipelineData, ShellError, Signature, Span, Spanned,
    SyntaxShape, Type, Value, shell_error::generic::GenericError,
};

use polars::lazy::dsl::{
    all_horizontal, any_horizontal, max_horizontal, mean_horizontal, min_horizontal, sum_horizontal,
};
use polars::prelude::Expr;

enum HorizontalType {
    All,
    Any,
    Min,
    Max,
    Sum,
    Mean,
}

impl HorizontalType {
    fn from_str(roll_type: &str, span: Span) -> Result<Self, ShellError> {
        match roll_type {
            "all" => Ok(Self::All),
            "any" => Ok(Self::Any),
            "min" => Ok(Self::Min),
            "max" => Ok(Self::Max),
            "sum" => Ok(Self::Sum),
            "mean" => Ok(Self::Mean),
            _ => Err(ShellError::Generic(
                GenericError::new(
                    "Wrong operation",
                    "Operation not valid for cumulative",
                    span,
                )
                .with_help("Allowed values: all, any, max, min, sum, mean"),
            )),
        }
    }
}

#[derive(Clone)]
pub struct Horizontal;

impl PluginCommand for Horizontal {
    type Plugin = PolarsPlugin;

    fn name(&self) -> &str {
        "polars horizontal"
    }

    fn description(&self) -> &str {
        "Horizontal calculation across multiple columns."
    }

    fn signature(&self) -> Signature {
        Signature::build(self.name())
            .input_output_type(Type::Any, PolarsPluginType::NuExpression.into())
            .required(
                "type",
                SyntaxShape::String,
                "Horizontal operation. Values of all, any, min, max, sum, and mean are accepted.",
            )
            .rest(
                "Group-by expressions",
                SyntaxShape::Any,
                "Expression(s) that define the lazy group-by.",
            )
            .switch(
                "nulls",
                "If set, null value in the input will lead to null output.",
                Some('n'),
            )
            .category(Category::Custom("expression".into()))
    }

    fn examples(&self) -> Vec<Example<'_>> {
        vec![
            Example {
                description: "Horizontal sum across two columns (ignore nulls by default)",
                example: "[[a b]; [1 2] [2 3] [3 4] [4 5] [5 null]]
                    | polars into-df
                    | polars select (polars horizontal sum a b)
                    | polars collect",
                result: Some(
                    NuDataFrame::try_from_columns(
                        vec![Column::new(
                            "sum".to_string(),
                            vec![
                                Value::test_int(3),
                                Value::test_int(5),
                                Value::test_int(7),
                                Value::test_int(9),
                                Value::test_int(5),
                            ],
                        )],
                        None,
                        Span::test_data(),
                    )
                    .expect("simple df for test should not fail")
                    .into_value(Span::test_data()),
                ),
            },
            Example {
                description: "Horizontal sum across two columns while accounting for nulls",
                example: "[[a b]; [1 2] [2 3] [3 4] [4 5] [5 null]]
                    | polars into-df
                    | polars select (polars horizontal sum a b --nulls)
                    | polars collect",
                result: Some(
                    NuDataFrame::try_from_columns(
                        vec![Column::new(
                            "sum".to_string(),
                            vec![
                                Value::test_int(3),
                                Value::test_int(5),
                                Value::test_int(7),
                                Value::test_int(9),
                                Value::test_nothing(),
                            ],
                        )],
                        None,
                        Span::test_data(),
                    )
                    .expect("simple df for test should not fail")
                    .into_value(Span::test_data()),
                ),
            },
        ]
    }

    fn run(
        &self,
        plugin: &Self::Plugin,
        engine: &EngineInterface,
        call: &EvaluatedCall,
        _input: PipelineData,
    ) -> Result<PipelineData, LabeledError> {
        let func_type: Spanned<String> = call.req(0)?;
        let func_type = HorizontalType::from_str(&func_type.item, func_type.span)?;

        let vals: Vec<Value> = call.rest(1)?;
        let expr_value = Value::list(vals, call.head);
        let exprs = NuExpression::extract_exprs(plugin, expr_value)?;

        let ignore_nulls = !call.has_flag("nulls")?;

        command(plugin, engine, call, func_type, exprs, ignore_nulls).map_err(LabeledError::from)
    }
}

fn command(
    plugin: &PolarsPlugin,
    engine: &EngineInterface,
    call: &EvaluatedCall,
    func_type: HorizontalType,
    exprs: Vec<Expr>,
    ignore_nulls: bool,
) -> Result<PipelineData, ShellError> {
    let res: NuExpression = match func_type {
        HorizontalType::All => all_horizontal(exprs),
        HorizontalType::Any => any_horizontal(exprs),
        HorizontalType::Max => max_horizontal(exprs),
        HorizontalType::Min => min_horizontal(exprs),
        HorizontalType::Sum => sum_horizontal(exprs, ignore_nulls),
        HorizontalType::Mean => mean_horizontal(exprs, ignore_nulls),
    }
    .map_err(|e| {
        ShellError::Generic(
            GenericError::new("Cannot apply horizontal aggregation", "", call.head)
                .with_help(e.to_string()),
        )
    })?
    .into();

    res.to_pipeline_data(plugin, engine, call.head)
}

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

    #[test]
    fn test_examples() -> Result<(), ShellError> {
        test_polars_plugin_command(&Horizontal)
    }
}
