use std::fs::File;

use crate::graph::node::RunnableNode;
use crate::io::DataSource;
use async_trait::async_trait;
use itertools::Itertools;
use polars::io::SerWriter;
use polars::prelude::{col, lit, CsvWriter, Expr, LazyCsvReader, LazyFileListReader};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};

fn reduce(
    grouping_nodes: &Vec<String>,
    operations: &Vec<ReductionOperation>,
    input: &DataSource,
    output: &DataSource,
) -> anyhow::Result<()> {
    let df = LazyCsvReader::new(&input.path).finish()?;

    let mut df = df
        .group_by(
            grouping_nodes
                .iter()
                .map(|column| col(column))
                .collect_vec(),
        )
        .agg(
            &operations
                .iter()
                .map(|operation| operation.to_aggregate_function())
                .collect_vec(),
        )
        .collect()?;
    let mut file = File::create(&output.path)?;
    CsvWriter::new(&mut file).finish(&mut df)?;
    Ok(())
}

#[derive(Serialize, Deserialize, Clone, JsonSchema)]
pub enum ReductionOperationType {
    Sum,
    Multiply,
    Max,
    Min,
    Average,
    Count,
    Concat(ConcatProperties),
}

#[derive(Serialize, Deserialize, Clone, JsonSchema)]
pub struct ConcatProperties {
    pub prefix: String,
    pub suffix: String,
    pub separator: String,
}

#[derive(Serialize, Deserialize, Clone, JsonSchema)]
pub struct ReductionOperation {
    pub column_name: String,
    pub reducer: ReductionOperationType,
    // TODO: Consider adding filtering on another column, so we can get sum of males grouping on some other column
    // TODO: Also consider adding an alias here
}

impl ReductionOperation {
    fn to_aggregate_function(&self) -> Expr {
        match &self.reducer {
            ReductionOperationType::Sum => col(&self.column_name).sum(),
            // TODO: Check this, docs are a bit flaky on whether this'll work
            ReductionOperationType::Multiply => col(&self.column_name).product(),
            ReductionOperationType::Max => col(&self.column_name).max(),
            ReductionOperationType::Min => col(&self.column_name).min(),
            ReductionOperationType::Average => col(&self.column_name).mean(),
            ReductionOperationType::Count => col(&self.column_name).count(),
            ReductionOperationType::Concat(concat_properties) => {
                lit(concat_properties.prefix.clone())
                    .append(
                        col(&self.column_name)
                            .list()
                            .join(lit(concat_properties.separator.clone()), true),
                        false,
                    )
                    .append(lit(concat_properties.suffix.clone()), false)
            }
        }
    }
}

#[derive(Serialize, Deserialize, Clone, JsonSchema)]
pub struct ReductionNode {
    pub grouping_columns: Vec<String>,
    pub operations: Vec<ReductionOperation>,
    pub input_file: DataSource,
    pub output_file: DataSource,
}

pub struct ReductionNodeRunner {
    pub reduction_node: ReductionNode,
}
#[async_trait]
impl RunnableNode for ReductionNodeRunner {
    async fn run(&self) -> anyhow::Result<()> {
        reduce(
            &self.reduction_node.grouping_columns,
            &self.reduction_node.operations,
            &self.reduction_node.input_file,
            &self.reduction_node.output_file,
        )?;
        Ok(())
    }
}
