static USAGE: &str = r#"
Fill empty fields in selected columns of a CSV.

This command fills empty fields in the selected column
using the last seen non-empty field in the CSV. This is
useful to forward-fill values which may only be included
the first time they are encountered.

The option `--default <value>` fills all empty values
in the selected columns with the provided default value.
When `--default` is set, it takes precedence over forward-fill
and `--first`, which become no-ops.

The option `--first` fills empty values using the first
seen non-empty value in that column, instead of the most
recent non-empty value in that column.

The option `--backfill` fills empty values at the start of
the CSV with the first valid value in that column. This
requires buffering rows with empty values in the target
column which appear before the first valid value.

The option `--groupby` groups the rows by the specified
columns before filling in the empty values. Using this
option, empty values are only filled with values which
belong to the same group of rows, as determined by the
columns selected in the `--groupby` option.

When both `--groupby` and `--backfill` are specified, and the
CSV is not sorted by the `--groupby` columns, rows may be
re-ordered during output due to the buffering of rows
collected before the first valid value.

For examples, see https://github.com/dathere/qsv/blob/master/tests/test_fill.rs.
See also https://github.com/dathere/qsv/wiki/Transform-and-Reshape#fill

Usage:
    qsv fill [options] [--] <selection> [<input>]
    qsv fill --help

fill options:
    -g, --groupby <keys>    Group by specified columns.
    -f, --first             Fill using the first valid value of a column, instead of the latest.
    -b, --backfill          Fill initial empty values with the first valid value.
    -v, --default <value>   Fill using this default value.

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 interpreted
                           as headers. (i.e., They are not searched, analyzed,
                           sliced, etc.)
    -d, --delimiter <arg>  The field delimiter for reading CSV data.
                           Must be a single character. (default: ,)
"#;

use std::{io, iter, ops};

use foldhash::{HashMap, HashMapExt, HashSet};
use serde::Deserialize;

use crate::{
    CliResult,
    config::{Config, Delimiter},
    select::{SelectColumns, Selection},
    util,
    util::ByteString,
};

type BoxedWriter = csv::Writer<Box<dyn io::Write + 'static>>;
type BoxedReader = csv::Reader<Box<dyn io::Read + Send + 'static>>;

#[derive(Deserialize)]
struct Args {
    arg_input:       Option<String>,
    arg_selection:   SelectColumns,
    flag_output:     Option<String>,
    flag_no_headers: bool,
    flag_delimiter:  Option<Delimiter>,
    flag_groupby:    Option<SelectColumns>,
    flag_first:      bool,
    flag_backfill:   bool,
    flag_default:    Option<String>,
}

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

    let rconfig = Config::new(args.arg_input.as_ref())
        .delimiter(args.flag_delimiter)
        .no_headers_flag(args.flag_no_headers)
        .select(args.arg_selection);

    let wconfig = Config::new(args.flag_output.as_ref());

    let mut rdr = rconfig.reader()?;
    let mut wtr = wconfig.writer()?;

    let headers = rdr.byte_headers()?.clone();
    let select = rconfig.selection(&headers)?;
    let groupby = match args.flag_groupby {
        Some(value) => Some(value.selection(&headers, !rconfig.no_headers)?),
        None => None,
    };

    if !rconfig.no_headers {
        rconfig.write_headers(&mut rdr, &mut wtr)?;
    }

    let filler = Filler::new(groupby, select)
        .use_first_value(args.flag_first)
        .backfill_empty_values(args.flag_backfill)
        .use_default_value(args.flag_default);
    filler.fill(&mut rdr, &mut wtr)
}

#[derive(PartialEq, Eq, Hash, Clone, Debug)]
struct ByteRecord(Vec<ByteString>);

impl<'a> From<&'a csv::ByteRecord> for ByteRecord {
    fn from(record: &'a csv::ByteRecord) -> Self {
        ByteRecord(record.iter().map(<[u8]>::to_vec).collect())
    }
}

impl iter::FromIterator<ByteString> for ByteRecord {
    fn from_iter<T: IntoIterator<Item = ByteString>>(iter: T) -> Self {
        ByteRecord(Vec::from_iter(iter))
    }
}

impl iter::IntoIterator for ByteRecord {
    type IntoIter = ::std::vec::IntoIter<ByteString>;
    type Item = ByteString;

    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

impl ops::Deref for ByteRecord {
    type Target = [ByteString];

    fn deref(&self) -> &[ByteString] {
        &self.0
    }
}

type GroupKey = Option<ByteRecord>;
type GroupBuffer = HashMap<GroupKey, Vec<ByteRecord>>;
type Grouper = HashMap<GroupKey, GroupValues>;
type GroupKeySelection = Option<Selection>;

trait GroupKeyConstructor {
    fn key(&self, record: &csv::ByteRecord) -> Result<GroupKey, String>;
}

impl GroupKeyConstructor for GroupKeySelection {
    fn key(&self, record: &csv::ByteRecord) -> Result<GroupKey, String> {
        match self {
            Some(value) => Ok(Some(value.iter().map(|&i| record[i].to_vec()).collect())),
            None => Ok(None),
        }
    }
}

#[derive(Debug)]
struct GroupValues {
    map:     HashMap<usize, ByteString>,
    default: Option<ByteString>,
}

impl GroupValues {
    fn new(default: Option<ByteString>) -> Self {
        Self {
            map: HashMap::new(),
            default,
        }
    }
}

trait GroupMemorizer {
    fn fill(&self, selection: &Selection, record: ByteRecord) -> ByteRecord;
    fn fill_into(
        &self,
        select_set: &HashSet<usize>,
        input: &csv::ByteRecord,
        output: &mut csv::ByteRecord,
    );
    fn memorize(&mut self, selection: &Selection, record: &csv::ByteRecord);
    fn memorize_first(&mut self, selection: &Selection, record: &csv::ByteRecord);
}

impl GroupMemorizer for GroupValues {
    fn memorize(&mut self, selection: &Selection, record: &csv::ByteRecord) {
        for &col in selection.iter().filter(|&col| !record[*col].is_empty()) {
            self.map.insert(col, record[col].to_vec());
        }
    }

    fn memorize_first(&mut self, selection: &Selection, record: &csv::ByteRecord) {
        for &col in selection.iter().filter(|&col| !record[*col].is_empty()) {
            self.map.entry(col).or_insert_with(|| record[col].to_vec());
        }
    }

    // Fast path: write the filled row directly into a reusable csv::ByteRecord,
    // avoiding the per-row Vec<Vec<u8>> allocations of the owned-ByteRecord path.
    // Used when buffering is impossible (no --backfill).
    fn fill_into(
        &self,
        select_set: &HashSet<usize>,
        input: &csv::ByteRecord,
        output: &mut csv::ByteRecord,
    ) {
        output.clear();
        for (i, field) in input.iter().enumerate() {
            if field.is_empty()
                && select_set.contains(&i)
                && let Some(replacement) = self
                    .default
                    .as_deref()
                    .or_else(|| self.map.get(&i).map(Vec::as_slice))
            {
                output.push_field(replacement);
                continue;
            }
            output.push_field(field);
        }
    }

    fn fill(&self, selection: &Selection, record: ByteRecord) -> ByteRecord {
        record
            .into_iter()
            .enumerate()
            .map_selected(selection, |(col, field)| {
                (
                    col,
                    if field.is_empty() {
                        self.default
                            .as_ref()
                            .or_else(|| self.map.get(&col))
                            .cloned()
                            .unwrap_or(field)
                    } else {
                        field
                    },
                )
            })
            .map(|(_, field)| field)
            .collect()
    }
}

enum MemorizeKind {
    None,
    First,
    Latest,
}

struct Filler {
    grouper:       Grouper,
    groupby:       GroupKeySelection,
    select:        Selection,
    buffer:        GroupBuffer,
    first:         bool,
    backfill:      bool,
    default_value: Option<ByteString>,
}

impl Filler {
    fn new(groupby: GroupKeySelection, select: Selection) -> Self {
        Self {
            grouper: Grouper::new(),
            groupby,
            select,
            buffer: GroupBuffer::new(),
            first: false,
            backfill: false,
            default_value: None,
        }
    }

    const fn use_first_value(mut self, first: bool) -> Self {
        self.first = first;
        self
    }

    const fn backfill_empty_values(mut self, backfill: bool) -> Self {
        self.backfill = backfill;
        self
    }

    fn use_default_value(mut self, value: Option<String>) -> Self {
        self.default_value = value.map(|v| v.as_bytes().to_vec());
        self
    }

    fn fill(mut self, rdr: &mut BoxedReader, wtr: &mut BoxedWriter) -> CliResult<()> {
        let mut record = csv::ByteRecord::new();
        let mut out_record = csv::ByteRecord::new();

        // Precompute O(1) selected-column membership lookup for the fast path.
        let select_set: HashSet<usize> = self.select.iter().copied().collect();
        let memorize_kind = match (self.default_value.is_some(), self.first) {
            (true, _) => MemorizeKind::None,
            (false, true) => MemorizeKind::First,
            (false, false) => MemorizeKind::Latest,
        };

        while rdr.read_byte_record(&mut record)? {
            // Precompute groupby key
            let key = self.groupby.key(&record)?;

            // Look up the group; only clone the key if we have to insert a new one.
            // Clone the default value lazily for the same reason.
            let group = match self.grouper.get_mut(&key) {
                Some(g) => g,
                None => self
                    .grouper
                    .entry(key.clone())
                    .or_insert_with(|| GroupValues::new(self.default_value.clone())),
            };

            match memorize_kind {
                MemorizeKind::None => {},
                MemorizeKind::First => group.memorize_first(&self.select, &record),
                MemorizeKind::Latest => group.memorize(&self.select, &record),
            }

            // Fast path: without --backfill we never buffer, so we can write the
            // filled row directly into a reusable csv::ByteRecord. This avoids
            // the Vec<Vec<u8>> allocations of the owned-ByteRecord round-trip.
            if !self.backfill {
                group.fill_into(&select_set, &record, &mut out_record);
                wtr.write_byte_record(&out_record)?;
                continue;
            }

            let row = group.fill(&self.select, ByteRecord::from(&record));

            // Handle buffering rows which still have nulls.
            if self.select.iter().any(|&i| row[i].is_empty()) {
                self.buffer.entry(key.clone()).or_default().push(row);
            } else {
                if let Some(rows) = self.buffer.remove(&key) {
                    for buffered_row in rows {
                        wtr.write_record(group.fill(&self.select, buffered_row).iter())?;
                    }
                }
                wtr.write_record(row.iter())?;
            }
        }

        // Ensure any remaining buffers are dumped at the end.
        for (key, rows) in self.buffer {
            // safety: every key inserted into `buffer` was first inserted into
            // `grouper` immediately above, so the lookup always succeeds.
            let group = self.grouper.get(&key).unwrap();
            for buffered_row in rows {
                wtr.write_record(group.fill(&self.select, buffered_row).iter())?;
            }
        }

        wtr.flush()?;
        Ok(())
    }
}

struct MapSelected<I, F> {
    selection:       Vec<usize>,
    selection_index: usize,
    index:           usize,
    iterator:        I,
    predicate:       F,
}

impl<I: iter::Iterator, F> iter::Iterator for MapSelected<I, F>
where
    F: FnMut(I::Item) -> I::Item,
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        let item = self.iterator.next()?;
        let result = match self.selection_index {
            ref mut sidx if (self.selection.get(*sidx) == Some(&self.index)) => {
                *sidx += 1;
                Some((self.predicate)(item))
            },
            _ => Some(item),
        };
        self.index += 1;
        result
    }
}

trait Selectable<B>
where
    Self: iter::Iterator<Item = B> + Sized,
{
    fn map_selected<F>(self, selector: &Selection, predicate: F) -> MapSelected<Self, F>
    where
        F: FnMut(B) -> B;
}

impl<B, C> Selectable<B> for C
where
    C: iter::Iterator<Item = B> + Sized,
{
    fn map_selected<F>(self, selector: &Selection, predicate: F) -> MapSelected<Self, F>
    where
        F: FnMut(B) -> B,
    {
        // MapSelected scans the row left-to-right and advances `selection_index`
        // only on ascending matches, so the index list must be sorted and unique
        // (user-supplied selections may be unordered or contain duplicates).
        let mut selection: Vec<usize> = selector.iter().copied().collect();
        selection.sort_unstable();
        selection.dedup();
        MapSelected {
            selection,
            selection_index: 0,
            index: 0,
            iterator: self,
            predicate,
        }
    }
}
