use std::sync::Arc;

use arrow::{array::RecordBatch, compute::take, datatypes::Schema};
use mpera_frontend::op::OpRef;

use crate::{
    op::group_common::{group_rows as partition_rows, merge_pending_batches},
    op::{OpOut, StatefulBatchOp},
    runtimedb::RuntimeDB,
};

const GROUP_KEYS_META: &str = "autark:group_keys";
const GROUP_LENGTHS_META: &str = "autark:group_lengths";

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

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

impl StatefulBatchOp for GroupByBatchOp {
    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_by", &mut self.pending_keys, &mut self.pending_values)?
        else {
            return Ok(Vec::new());
        };
        Ok(vec![group_keys(&keys, &values)?])
    }
}

pub fn split_grouped_output(
    batch: &RecordBatch,
) -> crate::Result<std::collections::HashMap<String, RecordBatch>> {
    let schema = batch.schema();
    let metadata = schema.metadata();
    let keys_json = metadata.get(GROUP_KEYS_META).ok_or_else(|| {
        autark_error::Error::Reader("group_by output missing group key metadata".to_string())
    })?;
    let lengths_json = metadata.get(GROUP_LENGTHS_META).ok_or_else(|| {
        autark_error::Error::Reader("group_by output missing group length metadata".to_string())
    })?;
    let keys: Vec<String> = serde_json::from_str(keys_json)?;
    let lengths: Vec<usize> = serde_json::from_str(lengths_json)?;
    if keys.len() != lengths.len() {
        return Err(autark_error::Error::Reader(format!(
            "group_by metadata mismatch: keys={} lengths={}",
            keys.len(),
            lengths.len()
        )));
    }

    let base_schema = Arc::new(Schema::new(
        schema
            .fields()
            .iter()
            .map(|field| field.as_ref().clone())
            .collect::<Vec<_>>(),
    ));
    let mut offset = 0usize;
    keys.into_iter()
        .zip(lengths)
        .map(|(key, len)| {
            let batch = batch.slice(offset, len);
            offset += len;
            let batch = RecordBatch::try_new(base_schema.clone(), batch.columns().to_vec())?;
            Ok((key, batch))
        })
        .collect()
}

fn group_keys(keys: &RecordBatch, values: &RecordBatch) -> crate::Result<RecordBatch> {
    let grouped = partition_rows("group_by", keys, values)?;
    let columns = values
        .columns()
        .iter()
        .map(|column| Ok(take(column.as_ref(), &grouped.grouped_indices, None)?))
        .collect::<crate::Result<Vec<_>>>()?;

    let mut metadata = values.schema().metadata().clone();
    metadata.insert(
        GROUP_KEYS_META.to_string(),
        serde_json::to_string(&grouped.key_names)?,
    );
    metadata.insert(
        GROUP_LENGTHS_META.to_string(),
        serde_json::to_string(&grouped.group_lengths)?,
    );
    let schema = Arc::new(Schema::new_with_metadata(
        values
            .schema()
            .fields()
            .iter()
            .map(|field| field.as_ref().clone())
            .collect::<Vec<_>>(),
        metadata,
    ));
    Ok(RecordBatch::try_new(schema, columns)?)
}
