use std::sync::Arc;

use arrow::{
    array::{Array, ArrayRef, RecordBatch, UInt32Array},
    compute::take,
    datatypes::{Field, Schema},
};
use arrow_select::concat::concat;
use mpera_frontend::op::{OpRef, ReduceKind};
use rayon::prelude::*;

use crate::{
    luminal_interop,
    op::group_common::{group_rows, merge_pending_batches},
    op::{OpOut, StatefulBatchOp},
    runtimedb::RuntimeDB,
    typecheck::reduce_output_dtype_for_array,
};

pub struct GroupReduceBatchOp {
    pub opref: OpRef,
    pub keys: OpRef,
    pub values: OpRef,
    pub kinds: Vec<ReduceKind>,
    pending_keys: Vec<RecordBatch>,
    pending_values: Vec<RecordBatch>,
}

impl GroupReduceBatchOp {
    pub fn new(opref: OpRef, keys: OpRef, values: OpRef, kinds: Vec<ReduceKind>) -> Self {
        Self {
            opref,
            keys,
            values,
            kinds,
            pending_keys: Vec::new(),
            pending_values: Vec::new(),
        }
    }
}

impl StatefulBatchOp for GroupReduceBatchOp {
    fn opref(&self) -> OpRef {
        self.opref
    }

    fn consume(&mut self, rtdb: &RuntimeDB) -> OpOut {
        let Some(keys) = rtdb.read(self.keys) else {
            return Ok(None);
        };
        let Some(values) = rtdb.read(self.values) else {
            return Ok(None);
        };

        self.pending_keys.push(keys.clone());
        self.pending_values.push(values.clone());
        Ok(None)
    }

    fn finalize(&mut self) -> crate::Result<Vec<RecordBatch>> {
        let Some((keys, values)) = merge_pending_batches(
            "group_reduce",
            &mut self.pending_keys,
            &mut self.pending_values,
        )?
        else {
            return Ok(Vec::new());
        };
        Ok(vec![group_reduce(&keys, &values, &self.kinds)?])
    }
}

fn group_reduce(
    keys: &RecordBatch,
    values: &RecordBatch,
    kinds: &[ReduceKind],
) -> crate::Result<RecordBatch> {
    let grouped = group_rows("group_reduce", keys, values)?;
    aggregate_groups(keys, values, &grouped.first_indices, &grouped.groups, kinds)
}

fn take_keys(keys: &RecordBatch, indices: &UInt32Array) -> crate::Result<RecordBatch> {
    let columns = keys
        .columns()
        .iter()
        .map(|column| Ok(take(column.as_ref(), indices, None)?))
        .collect::<crate::Result<Vec<_>>>()?;

    Ok(RecordBatch::try_new(keys.schema(), columns)?)
}

fn aggregate_groups(
    keys: &RecordBatch,
    values: &RecordBatch,
    first_indices: &UInt32Array,
    groups: &[Vec<u32>],
    kinds: &[ReduceKind],
) -> crate::Result<RecordBatch> {
    let key_batch = take_keys(keys, first_indices)?;
    let key_fields = key_batch
        .schema()
        .fields()
        .iter()
        .cloned()
        .collect::<Vec<_>>();
    let values_schema = values.schema();

    let aggregated = (0..values.num_columns())
        .map(|value_ix| -> crate::Result<Vec<(Arc<Field>, ArrayRef)>> {
            let value = values.column(value_ix);
            let field = values_schema.field(value_ix);

            kinds
                .iter()
                .map(|kind| -> crate::Result<(Arc<Field>, ArrayRef)> {
                    let reduced = groups
                        .par_iter()
                        .map(|group| reduce_group(value.as_ref(), group, kind))
                        .collect::<Vec<_>>()
                        .into_iter()
                        .collect::<crate::Result<Vec<_>>>()?;
                    let reduced_refs = reduced
                        .iter()
                        .map(|array| array.as_ref())
                        .collect::<Vec<_>>();
                    Ok((
                        Arc::new(Field::new(
                            &reduced_field_name(field.name(), kind),
                            reduce_output_dtype_for_array(kind, value.as_ref())?,
                            true,
                        )),
                        concat(&reduced_refs)?,
                    ))
                })
                .collect()
        })
        .collect::<crate::Result<Vec<_>>>()?;

    let fields = key_fields
        .into_iter()
        .chain(
            aggregated
                .iter()
                .flat_map(|entries| entries.iter().map(|(field, _)| field.clone())),
        )
        .collect::<Vec<_>>();
    let columns = key_batch
        .columns()
        .iter()
        .cloned()
        .chain(
            aggregated
                .into_iter()
                .flat_map(|entries| entries.into_iter().map(|(_, column)| column)),
        )
        .collect::<Vec<_>>();

    Ok(RecordBatch::try_new(
        Arc::new(Schema::new(fields)),
        columns,
    )?)
}

fn reduce_group(array: &dyn Array, group: &[u32], kind: &ReduceKind) -> crate::Result<ArrayRef> {
    let taken = take(array, &UInt32Array::from(group.to_vec()), None)?;
    luminal_interop::reduce_array(kind, taken.as_ref())
}

fn reduced_field_name(name: &str, kind: &ReduceKind) -> String {
    format!("{name}_{}", reduce_kind_name(kind))
}

fn reduce_kind_name(kind: &ReduceKind) -> &'static str {
    match kind {
        ReduceKind::Count => "count",
        ReduceKind::Sum => "sum",
        ReduceKind::Product => "product",
        ReduceKind::Mean => "mean",
        ReduceKind::Stdev => "stdev",
        ReduceKind::Min => "min",
        ReduceKind::Max => "max",
        ReduceKind::Last => "last",
        ReduceKind::Percentile { .. } => "percentile",
    }
}
