mod realize;

use std::collections::HashMap;
use std::fmt;

use mpera::{
    program::Program,
    programoutput::{OutputValue, ProgramOutput},
};
use mpera2::{pipeline::Artifact, runtime::Runtime};

use arrow::array::RecordBatch;
use arrow::datatypes::Schema as ArrowSchema;
use arrow::util::pretty::pretty_format_batches;
use autark_error::Error;
use autark_reader::OnceReader;
use autark_schema::Schema;
use autark_sinks::Sink;
use realize::stream_input_batches;

type Result<T> = std::result::Result<T, Error>;

pub struct OnceFrame {
    readers: Vec<Box<dyn OnceReader>>,
    sinks: Vec<Box<dyn Sink>>,
}

// TODO: Sep file
#[derive(Debug, Hash)]
pub struct RealizedOnceFrame {
    program_output: ProgramOutput,
}

impl RealizedOnceFrame {
    pub fn output(&self) -> &ProgramOutput {
        &self.program_output
    }
}

impl fmt::Display for RealizedOnceFrame {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (k, v) in self.program_output.0.iter() {
            match v {
                OutputValue::Frame(batch) => {
                    writeln!(
                        f,
                        "\n--------  \nKEY: {};\nVALUE: \n{}\n                ",
                        k,
                        pretty_format_batches(std::slice::from_ref(batch))
                            .map_err(|_| fmt::Error)?
                    )?;
                    writeln!(f, "--------")?;
                }
                OutputValue::Groups(groups) => {
                    writeln!(f, "\n--------  \nKEY: {};\nVALUE:", k)?;
                    for (group_key, batch) in groups {
                        writeln!(
                            f,
                            "GROUP {}:\n{}\n",
                            group_key,
                            pretty_format_batches(std::slice::from_ref(batch))
                                .map_err(|_| fmt::Error)?
                        )?;
                    }
                    writeln!(f, "--------")?;
                }
            }
        }
        Ok(())
    }
}

impl OnceFrame {
    pub fn new<R: OnceReader + 'static>(reader: R) -> OnceFrame {
        OnceFrame {
            readers: vec![Box::new(reader)],
            sinks: Vec::new(),
        }
    }

    pub fn schema(&self, index: Option<usize>) -> Result<Schema> {
        let idx = index.unwrap_or(0);
        let reader = self.readers.get(idx).ok_or(Error::EmptyReader)?;
        Ok(reader.schema()?.as_ref().clone().into())
    }

    pub fn schema_of_columns(&self, index: Option<usize>, columns: &[&str]) -> Result<Schema> {
        let idx = index.unwrap_or(0);
        let reader = self.readers.get(idx).ok_or(Error::EmptyReader)?;
        let schema = reader.schema()?;
        let fields: Result<Vec<_>> = columns
            .iter()
            .map(|name| Ok(schema.field_with_name(name)?.clone()))
            .collect();
        let fields = fields?;
        Ok(arrow::datatypes::Schema::new(fields).into())
    }

    pub fn describe(&self, index: Option<usize>) -> Result<String> {
        let schema = self.schema(index)?.into_arrow();
        let fields = schema
            .fields()
            .iter()
            .map(|field| {
                let nullability = if field.is_nullable() { "?" } else { "" };
                format!("{} {}{}", field.name(), field.data_type(), nullability)
            })
            .collect::<Vec<_>>();

        if fields.is_empty() {
            return Ok("OnceFrame {}".to_string());
        }

        Ok(format!("OnceFrame {{\n  {}\n}}", fields.join("\n  ")))
    }

    pub fn with_reader<R: OnceReader + 'static>(mut self, reader: R) -> OnceFrame {
        self.readers.push(Box::new(reader));
        self
    }

    pub fn with_sink<S: Sink + 'static>(mut self, sink: S) -> OnceFrame {
        self.sinks.push(Box::new(sink));
        self
    }

    // shall take S: Sink
    // thinking: what happens if needers more than one frame, perchance acept sequence of reader
    pub fn realize(
        &mut self,
        p: Program,
        rename: Option<Vec<HashMap<String, String>>>,
    ) -> Result<RealizedOnceFrame> {
        autark_telemetry::init_from_env();
        let artifact = Artifact::compile(p)?;
        let mut runtime = Runtime::new(artifact)?;

        for (ix, reader) in self.readers.iter_mut().enumerate() {
            stream_input_batches(
                &mut **reader,
                ix,
                rename_for_reader(rename.as_deref(), ix),
                |batch| runtime.run_batch(ix, batch),
            )?;
        }
        runtime.finalize()?;

        let output = ProgramOutput(
            runtime
                .collect_output()?
                .0
                .into_iter()
                .map(|(name, value)| (name.clone(), rename_output_value(&name, value)))
                .collect(),
        );

        self.sinks
            .iter()
            .try_for_each(|sink| sink.sink(output.clone()))?;

        Ok(RealizedOnceFrame {
            program_output: output,
        })
    }
}

fn rename_for_reader(
    rename: Option<&[HashMap<String, String>]>,
    reader_ix: usize,
) -> Option<&HashMap<String, String>> {
    rename.and_then(|maps| maps.get(reader_ix))
}

fn rename_output_value(name: &str, value: OutputValue) -> OutputValue {
    match value {
        OutputValue::Frame(batch) => OutputValue::Frame(rename_output_batch(name, batch)),
        OutputValue::Groups(groups) => OutputValue::Groups(groups),
    }
}

fn rename_output_batch(name: &str, batch: RecordBatch) -> RecordBatch {
    if batch.num_columns() != 1 {
        return batch;
    }

    let schema = batch.schema();
    let field = schema.field(0);
    let schema = ArrowSchema::new(vec![arrow::datatypes::Field::new(
        name,
        field.data_type().clone(),
        field.is_nullable(),
    )]);
    RecordBatch::try_new(std::sync::Arc::new(schema), batch.columns().to_vec()).unwrap_or(batch)
}
