static USAGE: &str = r#"
Do a spreadsheet-style pivot of a given CSV using the Pola.rs engine.   

Usage:
    qsv pivot [options] <values> <index> <columns> <agg_function>
    qsv pivot --help

pivot arguments:
    <values>               The column values to aggregate. Can be multiple columns
                           if the <columns> arguments contains multiple columns as well.
    <index>                One or multiple keys to group by.
    <columns>              Name of the column(s) whose values will be used as the
                           header of the output CSV.
    <agg_function>         The aggregate function to apply to the values. Can be
                           one of: first, sum, max, min, mean, median, last, count.

pivot options:
    --sort                 Sort the transposed columns by name. Default is by
                           order of discovery.
    --separator            The separator to use for generated column names.
                           [default: _]]            

Common options:
    -h, --help             Display this message
    -o, --output <file>    Write output to <file> instead of stdout.
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. (default: ,)
"#;

use std::{
    fs::File,
    io::{self, Write},
    path::Path,
    str,
};

use polars::{io::prelude::*, prelude::*};
use serde::Deserialize;

use crate::{config::Delimiter, util, CliError, CliResult};

#[derive(Deserialize)]
struct Args {
    arg_values:   String,
    arg_index:     String,
    arg_columns:   String,
    arg_agg_function:     String,
    flag_sort:      bool,
    flag_separator:     bool,
    flag_output:    Option<String>,
    flag_delimiter: Option<Delimiter>,
}

impl From<polars::error::PolarsError> for CliError {
    fn from(err: polars::error::PolarsError) -> CliError {
        CliError::Other(format!("Polars error: {err:?}"))
    }
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = util::get_args(USAGE, argv)?;
    let join = args.new_join()?;
    let join_shape = match (
        args.flag_left,
        args.flag_outer,
        args.flag_cross,
        args.flag_semi,
        args.flag_anti,
    ) {
        (false, false, false, false, false) => join.polars_join(JoinType::Inner),
        (true, false, false, false, false) => join.polars_join(JoinType::Left),
        (false, true, false, false, false) => join.polars_join(JoinType::Outer),
        (false, false, true, false, false) => join.polars_join(JoinType::Cross),
        (false, false, false, true, false) => join.polars_join(JoinType::Semi),
        (false, false, false, false, true) => join.polars_join(JoinType::Anti),
        _ => fail!("Please pick exactly one join operation."),
    }?;

    if !args.flag_quiet {
        eprintln!("{join_shape:?}");
    }

    Ok(())
}

struct JoinStruct {
    lf1:    LazyFrame,
    sel1:   String,
    lf2:    LazyFrame,
    sel2:   String,
    output: Option<String>,
    delim:  u8,
}

impl JoinStruct {
    fn polars_join(self, jointype: JoinType) -> CliResult<(usize, usize)> {
        let selcols1: Vec<_> = self.sel1.split(',').map(polars::lazy::dsl::col).collect();
        let selcols2: Vec<_> = self.sel2.split(',').map(polars::lazy::dsl::col).collect();

        let selcols1_len = selcols1.len();
        let selcols2_len = selcols2.len();

        if selcols1_len != selcols2_len {
            return fail_clierror!(
                "Both columns1 ({selcols1:?}) and columns2 ({selcols2:?}) must specify the same \
                 number of columns ({selcols1_len } != {selcols2_len})."
            );
        }

        let optimize_all = polars::lazy::frame::OptState {
            projection_pushdown: true,
            predicate_pushdown:  true,
            type_coercion:       true,
            simplify_expr:       true,
            file_caching:        true,
            slice_pushdown:      true,
            streaming:           true,
        };

        let mut join_results = self
            .lf1
            .with_optimizations(optimize_all)
            .join_builder()
            .with(self.lf2.with_optimizations(optimize_all))
            .left_on(selcols1)
            .right_on(selcols2)
            .how(jointype)
            .force_parallel(true)
            .finish()
            .collect()?;

        // no need to use buffered writer here, as CsvWriter already does that
        let mut out_writer = match self.output {
            Some(output_file) => {
                let path = Path::new(&output_file);
                Box::new(File::create(path).unwrap()) as Box<dyn Write>
            }
            None => Box::new(io::stdout()) as Box<dyn Write>,
        };

        // shape is the number of rows and columns
        let join_shape = join_results.shape();

        CsvWriter::new(&mut out_writer)
            .has_header(true)
            .with_delimiter(self.delim)
            .finish(&mut join_results)?;

        Ok(join_shape)
    }
}

impl Args {
    fn new_join(&self) -> CliResult<JoinStruct> {
        let delim = if let Some(delimiter) = self.flag_delimiter {
            delimiter.as_byte()
        } else {
            b','
        };

        let lf1 = LazyCsvReader::new(&self.arg_input1)
            .has_header(true)
            .with_missing_is_null(self.flag_nulls)
            .with_delimiter(delim)
            .finish()?;

        let lf2 = LazyCsvReader::new(&self.arg_input2)
            .has_header(true)
            .with_missing_is_null(self.flag_nulls)
            .with_delimiter(delim)
            .finish()?;

        Ok(JoinStruct {
            lf1,
            sel1: self.arg_columns1.clone(),
            lf2,
            sel2: self.arg_columns2.clone(),
            output: self.flag_output.clone(),
            delim,
        })
    }
}
