use arrow::array::RecordBatch;
use mpera_frontend::programoutput::{OutputValue, ProgramOutput};

use crate::{op::group_by::split_grouped_output, pipeline::Artifact, runtimedb::RuntimeDB};

pub struct Runtime {
    pub artifact: Artifact,
    pub rtdb: RuntimeDB,
}

impl Runtime {
    pub fn new(artifact: Artifact) -> crate::Result<Self> {
        let op_count = artifact
            .program
            .len()
            .map_err(|_| autark_error::Error::PoisonedLock)?;
        Ok(Self {
            artifact,
            rtdb: RuntimeDB::new(op_count),
        })
    }

    pub fn run_batch(&mut self, frame_ix: usize, batch: RecordBatch) -> crate::Result<()> {
        self.rtdb.set_input(frame_ix, batch);
        for op in &mut self.artifact.compiled {
            match op.consume(&self.rtdb)? {
                Some(out) => self.rtdb.write(op.opref(), out),
                None => self.rtdb.block(op.opref()),
            }
        }
        Ok(())
    }

    pub fn finalize(&mut self) -> crate::Result<()> {
        for op in &mut self.artifact.compiled {
            let out = op.finalize()?;
            if !out.is_empty() {
                self.rtdb.write(op.opref(), out);
            }
        }
        Ok(())
    }

    pub fn collect_output(&self) -> crate::Result<ProgramOutput> {
        let op_pool = self
            .artifact
            .program
            .op_pool
            .read()
            .map_err(|_| autark_error::Error::PoisonedLock)?;
        let outputs = self
            .artifact
            .outputs
            .iter()
            .map(|(name, opref)| {
                let batches = self.rtdb.read_batches(*opref).ok_or_else(|| {
                    autark_error::Error::Reader(format!("output `{name}` is blocking"))
                })?;
                let batch = merge_output_batches(name, batches)?;
                let value = match op_pool.get(*opref) {
                    Some(mpera_frontend::op::Op::GroupBy { .. }) => {
                        OutputValue::Groups(split_grouped_output(&batch)?.into_iter().collect())
                    }
                    _ => OutputValue::Frame(batch),
                };
                Ok((name.clone(), value))
            })
            .collect::<crate::Result<_>>()?;
        Ok(ProgramOutput(outputs))
    }
}

fn merge_output_batches(name: &str, batches: &[RecordBatch]) -> crate::Result<RecordBatch> {
    let first = batches.first().ok_or_else(|| {
        autark_error::Error::Reader(format!("output `{name}` produced no record batches"))
    })?;
    if batches.len() == 1 {
        Ok(first.clone())
    } else {
        Ok(arrow::compute::concat_batches(&first.schema(), batches)?)
    }
}
