xan 0.56.0

use rand::seq::SliceRandom;
use std::io::{SeekFrom, Write};

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

static USAGE: &str = "
Shuffle the rows of a given CSV file. This requires loading the whole
file in memory. If memory is scarce and target file is seekable (not stdin,
nor unindexed compressed file), you can also use the -e/--external flag that
only requires memory proportional to the number of rows of the file.

Usage:
    xan shuffle [options] [<input>]
    xan shuffle --help

shuffle options:
    --seed <number>  RNG seed.
    -e, --external   Shuffle the file without buffering it into memory. Only
                     works if target is seekable (no stdin etc.).

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 included in
                           the count.
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character.
";

#[derive(Deserialize)]
struct Args {
    arg_input: Option<String>,
    flag_output: Option<String>,
    flag_no_headers: bool,
    flag_delimiter: Option<Delimiter>,
    flag_seed: Option<usize>,
    flag_external: bool,
}

fn run_external(args: Args) -> CliResult<()> {
    let rconf = Config::new(&args.arg_input)
        .delimiter(args.flag_delimiter)
        .no_headers(args.flag_no_headers);
    let wconf = Config::new(&args.flag_output);

    // Seedmemorying rng
    let mut rng = util::acquire_rng(args.flag_seed);

    let mut header_len: Option<usize> = None;
    let mut positions: Vec<(u64, usize)> = Vec::new();
    let mut last_pos: u64 = 0;

    let mut output_wtr = wconf.buf_io_writer()?;

    {
        let mut rdr = rconf.simd_reader()?;

        if !rconf.no_headers {
            let header = rdr.byte_headers()?;

            if !header.is_empty() {
                last_pos = rdr.position();
                header_len = Some(last_pos as usize);
            }
        }

        let mut record = simd_csv::ByteRecord::new();

        while rdr.read_byte_record(&mut record)? {
            let pos = rdr.position();
            positions.push((last_pos, (pos - last_pos) as usize));
            last_pos = pos;
        }

        positions.shuffle(&mut rng);
    }

    let mut input_rdr = rconf.seekable_io_reader()?;
    let mut reading_buffer: Vec<u8> = Vec::new();

    if let Some(l) = header_len {
        reading_buffer.reserve(l);
        reading_buffer.extend((reading_buffer.len()..l).map(|_| 0));

        input_rdr.read_exact(&mut reading_buffer[0..l])?;
        output_wtr.write_all(&reading_buffer[0..l])?;
    }

    for (byte_offset, l) in positions {
        input_rdr.seek(SeekFrom::Start(byte_offset))?;

        reading_buffer.reserve(l);
        reading_buffer.extend((reading_buffer.len()..l).map(|_| 0));

        input_rdr.read_exact(&mut reading_buffer[0..l])?;
        output_wtr.write_all(&reading_buffer[0..l])?;
    }

    Ok(output_wtr.flush()?)
}

fn run_in_memory(args: Args) -> CliResult<()> {
    let rconf = Config::new(&args.arg_input)
        .delimiter(args.flag_delimiter)
        .no_headers(args.flag_no_headers);
    let wconf = Config::new(&args.flag_output);

    // Seeding rng
    let mut rng = util::acquire_rng(args.flag_seed);

    let mut rdr = rconf.simd_reader()?;
    let mut wtr = wconf.simd_writer()?;

    if !rconf.no_headers {
        wtr.write_byte_record(rdr.byte_headers()?)?;
    }

    let mut rows: Vec<simd_csv::ByteRecord> = Vec::new();

    for record in rdr.into_byte_records() {
        rows.push(record?);
    }

    rows.shuffle(&mut rng);

    for record in rows {
        wtr.write_byte_record(&record)?;
    }

    Ok(wtr.flush()?)
}

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

    if args.flag_external {
        run_external(args)
    } else {
        run_in_memory(args)
    }
}