use std::{
    collections::BTreeMap,
    hash::{Hash, Hasher},
};

use arrow::{array::RecordBatch, compute::concat_batches, ipc::writer::StreamWriter};

use crate::{Error, Result};

#[derive(Debug, Clone)]
pub enum OutputValue {
    Frame(RecordBatch),
    Groups(BTreeMap<String, RecordBatch>),
}

#[derive(Debug, Clone)]
pub struct ProgramOutput(pub BTreeMap<String, OutputValue>);

impl Hash for ProgramOutput {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.len().hash(state);
        for (name, value) in &self.0 {
            name.hash(state);
            hash_value(value, state);
        }
    }
}

fn hash_value<H: Hasher>(value: &OutputValue, state: &mut H) {
    match value {
        OutputValue::Frame(batch) => {
            0u8.hash(state);
            hash_batch(batch, state);
        }
        OutputValue::Groups(groups) => {
            1u8.hash(state);
            groups.len().hash(state);
            for (key, batch) in groups {
                key.hash(state);
                hash_batch(batch, state);
            }
        }
    }
}

fn hash_batch<H: Hasher>(batch: &RecordBatch, state: &mut H) {
    let mut buf = Vec::new();
    match StreamWriter::try_new(&mut buf, &batch.schema()).and_then(|mut writer| {
        writer.write(batch)?;
        writer.finish()
    }) {
        Ok(()) => buf.hash(state),
        Err(_) => {
            batch.num_rows().hash(state);
            batch.num_columns().hash(state);
            batch.schema().to_string().hash(state);
        }
    }
}

pub fn fuse(outputs: &[ProgramOutput]) -> Result<ProgramOutput> {
    let mut out = BTreeMap::<String, Vec<RecordBatch>>::new();
    for output in outputs {
        for (name, value) in &output.0 {
            let OutputValue::Frame(batch) = value else {
                return Err(Error::Reader(format!(
                    "cannot fuse grouped output `{name}`"
                )));
            };
            out.entry(name.clone()).or_default().push(batch.clone());
        }
    }
    Ok(ProgramOutput(
        out.into_iter()
            .map(|(name, batches)| {
                let batch = if batches.is_empty() {
                    return Err(Error::Reader(format!(
                        "cannot fuse output `{name}` with no batches"
                    )));
                } else if batches.len() == 1 {
                    batches.into_iter().next().ok_or_else(|| {
                        Error::Reader(format!("cannot extract fused output `{name}`"))
                    })?
                } else {
                    let schema = batches[0].schema();
                    concat_batches(&schema, &batches)?
                };
                Ok((name, OutputValue::Frame(batch)))
            })
            .collect::<Result<BTreeMap<_, _>>>()?,
    ))
}
