use std::iter::once;
use std::num::NonZeroUsize;

use ordered_float::NotNan;
use simd_csv::ByteRecord;

use crate::CliResult;
use crate::cmd::parallel::Args as ParallelArgs;
use crate::collections::{ClusteredInsertHashmap, DynamicOrd, TopKHeapMapWithTies};
use crate::config::{Config, Delimiter};
use crate::select::SelectedColumns;
use crate::util;

#[derive(PartialEq, Eq, PartialOrd, Ord, Clone)]
pub enum Value {
    Float(NotNan<f64>),
    String(Vec<u8>),
}

impl Value {
    pub fn new_float(cell: &[u8]) -> Option<Self> {
        fast_float::parse::<f64, &[u8]>(cell)
            .ok()
            .and_then(|f| NotNan::new(f).ok())
            .map(Self::Float)
    }

    pub fn new_string(cell: &[u8]) -> Option<Self> {
        Some(Self::String(cell.to_vec()))
    }
}

static USAGE: &str = "
Find top k values in selected column and return the associated CSV rows.

Runs in O(n * log k) time, n being the number of rows in target CSV file, and
consuming only O(k) memory, which is of course better than piping `xan sort`
into `xan head`.

Note that rows having empty values or values that cannot be parsed as numbers
in selected columns will be ignored along the way.

This command can also return the first k values or last k values in lexicographic
order using the -L/--lexicographic flag (note that the logic of the command is
tailored for numerical values and is therefore the reverse of `xan sort` in this
regard).

Examples:

Top 10 values in \"score\" column:

    $ xan top score file.csv

Top 50 values:

    $ xan top -l 50 score file.csv

Smallest 10 values:

    $ xan top -R score file.csv

Top 10 values with potential ties:

    $ xan top -T score file.csv

Top 10 values per distinct value of the \"category\" column:

    $ xan top -g category score file.csv

The same with a preprended \"rank\" column:

    $ xan top -g category -r rank score file.csv

Last 10 names in lexicographic order:

    $ xan top -L name file.csv

First 10 names in lexicographic order:

    $ xan top -LR name file.csv

Usage:
    xan top <column> [options] [<input>]
    xan top --help

top options:
    -l, --limit <n>          Number of top items to return. Cannot be < 1.
                             [default: 10]
    -R, --reverse            Reverse order.
    -L, --lexicographic      Rank values lexicographically instead of considering
                             them as numbers.
    -g, --groupby <cols>     Return top n values per group, represented
                             by the values in given columns.
    -S, --sorted             When used with -g/--groupby, indicates that the input
                             is sorted on the group columns so the command can run
                             faster and use memory proportional on -l/--limit rather
                             than group cardinality. This does not indicate the file
                             is sorted on scored column, because using `xan slice` on
                             such a file would effectively do the same thing as `xan top`
                             in this particular case.
    -r, --rank <col>         Name of a rank column to prepend.
    -T, --ties               Keep all rows tied for last. Will therefore
                             consume O(k + t) memory, t being the number of ties.
    -p, --parallel           Whether to use parallelization to speed up computation.
                             Will automatically select a suitable number of threads to use
                             based on your number of cores. Use -t, --threads if you want to
                             indicate the number of threads yourself.
    -t, --threads <threads>  Parellize computations using this many threads. Use -p, --parallel
                             if you want the number of threads to be automatically chosen instead.

Common options:
    -h, --help               Display this message
    -o, --output <file>      Write output to <file> instead of stdout.
    -n, --no-headers         When set, the first row will not be evaled
                             as headers.
    -d, --delimiter <arg>    The field delimiter for reading CSV data.
                             Must be a single character.
";

#[derive(Deserialize)]
struct Args {
    arg_input: Option<String>,
    arg_column: SelectedColumns,
    flag_no_headers: bool,
    flag_output: Option<String>,
    flag_delimiter: Option<Delimiter>,
    flag_limit: NonZeroUsize,
    flag_reverse: bool,
    flag_groupby: Option<SelectedColumns>,
    flag_sorted: bool,
    flag_rank: Option<String>,
    flag_ties: bool,
    flag_lexicographic: bool,
    flag_parallel: bool,
    flag_threads: Option<NonZeroUsize>,
}

impl Args {
    fn new_value(&self, cell: &[u8]) -> Option<Value> {
        if self.flag_lexicographic {
            Value::new_string(cell)
        } else {
            Value::new_float(cell)
        }
    }
}

pub fn run(argv: &[&str]) -> CliResult<()> {
    let args: Args = util::get_args(USAGE, argv)?;

    if args.flag_parallel || args.flag_threads.is_some() {
        if args.flag_groupby.is_some() {
            Err("-p/--parallel or -t/--threads cannot be used with -g/--groupby!")?;
        }

        let mut parallel_args = ParallelArgs::single_file(&args.arg_input, args.flag_threads)?;

        parallel_args.cmd_top = true;
        parallel_args.arg_column = Some(args.arg_column);
        parallel_args.flag_limit = args.flag_limit.get();
        parallel_args.flag_rank = args.flag_rank;
        parallel_args.flag_ties = args.flag_ties;
        parallel_args.flag_reverse = args.flag_reverse;
        parallel_args.flag_lexicographic = args.flag_lexicographic;

        parallel_args.flag_no_headers = args.flag_no_headers;
        parallel_args.flag_output = args.flag_output;
        parallel_args.flag_delimiter = args.flag_delimiter;

        return parallel_args.run();
    }

    let rconf = Config::new(&args.arg_input)
        .delimiter(args.flag_delimiter)
        .no_headers(args.flag_no_headers)
        .select(args.arg_column.clone());

    let mut rdr = rconf.simd_reader()?;
    let headers = rdr.byte_headers()?;
    let score_column = rconf.single_selection(headers)?;

    let groupby_sel_opt = args
        .flag_groupby
        .as_ref()
        .map(|cols| cols.selection(headers, !rconf.no_headers))
        .transpose()?;

    let mut wtr = Config::new(&args.flag_output).simd_writer()?;

    if !rconf.no_headers {
        if let Some(name) = &args.flag_rank {
            wtr.write_record(once(name.as_bytes()).chain(headers.iter()))?;
        } else {
            wtr.write_byte_record(headers)?;
        }
    }

    if let Some(sel) = &groupby_sel_opt {
        let mut record = ByteRecord::new();

        if args.flag_sorted {
            let mut heap = TopKHeapMapWithTies::<DynamicOrd<Value>, ByteRecord>::with_capacity(
                usize::from(args.flag_limit),
                args.flag_ties,
            );
            let mut buffer: Vec<ByteRecord> = Vec::with_capacity(args.flag_limit.get());

            let mut flush = |heap: &mut TopKHeapMapWithTies<DynamicOrd<Value>, ByteRecord>,
                             buffer: &mut Vec<ByteRecord>|
             -> CliResult<()> {
                buffer.clear();

                while let Some((_, record)) = heap.pop() {
                    buffer.push(record);
                }

                buffer.reverse();

                for (i, record) in buffer.iter().enumerate() {
                    if args.flag_rank.is_some() {
                        wtr.write_record(
                            once((i + 1).to_string().as_bytes()).chain(record.iter()),
                        )?;
                    } else {
                        wtr.write_byte_record(record)?;
                    }
                }

                Ok(())
            };

            let mut current_group_opt: Option<ByteRecord> = None;

            while rdr.read_byte_record(&mut record)? {
                if let Some(score) = args.new_value(&record[score_column]) {
                    let group: ByteRecord = sel.select(&record).collect();

                    match &current_group_opt {
                        Some(current_group) if current_group != &group => {
                            flush(&mut heap, &mut buffer)?;
                        }
                        _ => (),
                    };

                    heap.push_with(DynamicOrd::new(score, args.flag_reverse), || record.clone());
                    current_group_opt = Some(group);
                }
            }

            flush(&mut heap, &mut buffer)?;
        } else {
            let mut groups: ClusteredInsertHashmap<
                ByteRecord,
                TopKHeapMapWithTies<DynamicOrd<Value>, ByteRecord>,
            > = ClusteredInsertHashmap::new();

            while rdr.read_byte_record(&mut record)? {
                if let Some(score) = args.new_value(&record[score_column]) {
                    let group = sel.select(&record).collect();

                    let heap = groups.insert_with(group, || {
                        TopKHeapMapWithTies::with_capacity(
                            usize::from(args.flag_limit),
                            args.flag_ties,
                        )
                    });

                    heap.push_with(DynamicOrd::new(score, args.flag_reverse), || record.clone());
                }
            }

            for heap in groups.into_values() {
                for (i, (_, record)) in heap.into_sorted_vec().into_iter().enumerate() {
                    if args.flag_rank.is_some() {
                        wtr.write_record(
                            once((i + 1).to_string().as_bytes()).chain(record.iter()),
                        )?;
                    } else {
                        wtr.write_byte_record(&record)?;
                    }
                }
            }
        }
    } else {
        let mut record = ByteRecord::new();
        let mut heap = TopKHeapMapWithTies::<DynamicOrd<Value>, ByteRecord>::with_capacity(
            usize::from(args.flag_limit),
            args.flag_ties,
        );

        while rdr.read_byte_record(&mut record)? {
            if let Some(score) = args.new_value(&record[score_column]) {
                heap.push_with(DynamicOrd::new(score, args.flag_reverse), || record.clone());
            }
        }

        for (i, (_, record)) in heap.into_sorted_vec().into_iter().enumerate() {
            if args.flag_rank.is_some() {
                wtr.write_record(once((i + 1).to_string().as_bytes()).chain(record.iter()))?;
            } else {
                wtr.write_byte_record(&record)?;
            }
        }
    }

    Ok(wtr.flush()?)
}
