use std::collections::HashMap;

use arrow::{
    array::{RecordBatch, UInt32Array},
    compute::concat_batches,
};

use crate::rowkey;

pub struct GroupedRows {
    pub key_names: Vec<String>,
    pub first_indices: UInt32Array,
    pub grouped_indices: UInt32Array,
    pub group_lengths: Vec<usize>,
    pub groups: Vec<Vec<u32>>,
}

pub fn merge_pending_batches(
    label: &str,
    pending_keys: &mut Vec<RecordBatch>,
    pending_values: &mut Vec<RecordBatch>,
) -> crate::Result<Option<(RecordBatch, RecordBatch)>> {
    if pending_keys.is_empty() && pending_values.is_empty() {
        return Ok(None);
    }

    let keys = merge_batches(label, pending_keys)?;
    let values = merge_batches(label, pending_values)?;
    pending_keys.clear();
    pending_values.clear();
    Ok(Some((keys, values)))
}

pub fn group_rows(
    label: &str,
    keys: &RecordBatch,
    values: &RecordBatch,
) -> crate::Result<GroupedRows> {
    if keys.num_rows() != values.num_rows() {
        return Err(autark_error::Error::Reader(format!(
            "{label} requires matching row counts: keys={} values={}",
            keys.num_rows(),
            values.num_rows()
        )));
    }

    let (_, key_names, first_indices, groups) = (0..keys.num_rows()).try_fold(
        (
            HashMap::<String, usize>::new(),
            Vec::<String>::new(),
            Vec::<u32>::new(),
            Vec::<Vec<u32>>::new(),
        ),
        |(mut group_ix_by_key, mut key_names, mut first_indices, mut groups), row| {
            let key = rowkey::row_key(label, keys, row)?;
            if let Some(&group_ix) = group_ix_by_key.get(&key) {
                groups[group_ix].push(row as u32);
            } else {
                let group_ix = groups.len();
                group_ix_by_key.insert(key.clone(), group_ix);
                key_names.push(key);
                first_indices.push(row as u32);
                groups.push(vec![row as u32]);
            }
            Ok::<_, autark_error::Error>((group_ix_by_key, key_names, first_indices, groups))
        },
    )?;

    let group_lengths = groups.iter().map(Vec::len).collect::<Vec<_>>();
    let grouped_indices = groups
        .iter()
        .flat_map(|group| group.iter().copied())
        .collect::<Vec<_>>();

    Ok(GroupedRows {
        key_names,
        first_indices: UInt32Array::from(first_indices),
        grouped_indices: UInt32Array::from(grouped_indices),
        group_lengths,
        groups,
    })
}

fn merge_batches(label: &str, batches: &[RecordBatch]) -> crate::Result<RecordBatch> {
    let first = batches.first().ok_or_else(|| {
        autark_error::Error::Reader(format!("{label} requires at least one input batch"))
    })?;
    if batches.len() == 1 {
        Ok(first.clone())
    } else {
        Ok(concat_batches(&first.schema(), batches)?)
    }
}
