use std::collections::BTreeMap;

use mpera_frontend::{
    op::{Op, OpRef},
    program::Program,
};

use crate::op::{
    StatefulBatchOp,
    binary_op::BinaryOpBatchOp,
    columns::{ColumnsBatchOp, ColumnsSource},
    concat::ConcatBatchOp,
    constant::ConstantBatchOp,
    filter::FilterBatchOp,
    group_by::GroupByBatchOp,
    group_reduce::GroupReduceBatchOp,
    join::JoinBatchOp,
    order_by::OrderByBatchOp,
    reduce::ReduceBatchOp,
    rolling::RollingBatchOp,
    slice::SliceBatchOp,
    unary_op::UnaryOpBatchOp,
    where_op::WhereBatchOp,
};

pub struct Artifact {
    pub program: Program,
    pub compiled: Vec<Box<dyn StatefulBatchOp>>,
    pub outputs: BTreeMap<String, OpRef>,
}

impl Artifact {
    pub fn new(program: Program, compiled: Vec<Box<dyn StatefulBatchOp>>) -> Artifact {
        Artifact {
            program,
            compiled,
            outputs: BTreeMap::new(),
        }
    }

    pub fn compile(program: Program) -> crate::Result<Artifact> {
        let (compiled, outputs) = {
            let op_pool = program
                .op_pool
                .read()
                .map_err(|_| autark_error::Error::PoisonedLock)?;
            ensure_group_by_is_terminal(&op_pool)?;
            let reachable = reachable_oprefs(&op_pool, &compilation_roots(&program, &op_pool)?);
            let mut compiled = Vec::<Box<dyn StatefulBatchOp>>::new();
            let mut outputs = BTreeMap::new();

            for (ix, op) in op_pool.iter().enumerate() {
                let opref = OpRef(ix);
                if !reachable[ix] {
                    continue;
                }

                match op {
                    Op::Columns {
                        on,
                        index,
                        selector,
                    } => {
                        let source = match on {
                            Some(on) => ColumnsSource::Op(*on),
                            None => ColumnsSource::Input(*index),
                        };
                        compiled.push(Box::new(ColumnsBatchOp::new(
                            opref,
                            source,
                            selector.clone(),
                        )));
                    }
                    Op::Constant { value } => {
                        compiled.push(Box::new(ConstantBatchOp::new(opref, value.clone())));
                    }
                    Op::BinaryOp { kind, lhs, rhs } => {
                        compiled.push(Box::new(BinaryOpBatchOp::new(
                            opref,
                            *lhs,
                            *rhs,
                            kind.clone(),
                        )));
                    }
                    Op::Slice { on, start, end } => {
                        compiled.push(Box::new(SliceBatchOp::new(opref, *on, *start, *end)));
                    }
                    Op::Reduce { kind, on } => {
                        compiled.push(Box::new(ReduceBatchOp::new(opref, *on, kind.clone())));
                    }
                    Op::Rolling { on, n } => {
                        compiled.push(Box::new(RollingBatchOp::new(opref, *on, *n)));
                    }
                    Op::UnaryOp { kind, on } => {
                        compiled.push(Box::new(UnaryOpBatchOp::new(opref, *on, kind.clone())));
                    }
                    Op::Filter { on, mask } => {
                        compiled.push(Box::new(FilterBatchOp::new(opref, *on, *mask)));
                    }
                    Op::Where { mask, then, else_ } => {
                        compiled.push(Box::new(WhereBatchOp::new(opref, *mask, *then, *else_)));
                    }
                    Op::Concat { who } => {
                        compiled.push(Box::new(ConcatBatchOp::new(opref, who.clone())));
                    }
                    Op::GroupBy { keys, values } => {
                        compiled.push(Box::new(GroupByBatchOp::new(opref, *keys, *values)));
                    }
                    Op::GroupReduce {
                        keys,
                        values,
                        kinds,
                    } => {
                        compiled.push(Box::new(GroupReduceBatchOp::new(
                            opref,
                            *keys,
                            *values,
                            kinds.clone(),
                        )));
                    }
                    Op::Join {
                        left,
                        right,
                        left_on,
                        right_on,
                        kind,
                    } => {
                        compiled.push(Box::new(JoinBatchOp::new(
                            opref,
                            *left,
                            *right,
                            *left_on,
                            *right_on,
                            kind.clone(),
                        )));
                    }
                    Op::OrderBy {
                        what,
                        by,
                        ascending,
                    } => {
                        compiled.push(Box::new(OrderByBatchOp::new(opref, *what, *by, *ascending)));
                    }
                    Op::Output { name, value } => {
                        outputs.insert(name.clone(), *value);
                    }
                }
            }

            (compiled, outputs)
        };

        Ok(Artifact {
            program,
            compiled,
            outputs,
        })
    }
}

fn ensure_group_by_is_terminal(op_pool: &mpera_frontend::op::OpPool) -> crate::Result<()> {
    op_pool
        .iter()
        .enumerate()
        .try_for_each(|(ix, op)| -> crate::Result<()> {
            if matches!(op, Op::Output { .. }) {
                return Ok(());
            }
            let mut grouped_dependency = None;
            visit_dependencies(op, |dependency| {
                if grouped_dependency.is_none()
                    && matches!(op_pool.get(dependency), Some(Op::GroupBy { .. }))
                {
                    grouped_dependency = Some(dependency);
                }
            });
            if let Some(dependency) = grouped_dependency {
                return Err(autark_error::Error::Reader(format!(
                    "group_by output is terminal; op {} cannot consume grouped op {}",
                    ix, dependency.0
                )));
            }
            Ok(())
        })
}

fn compilation_roots(
    program: &Program,
    op_pool: &mpera_frontend::op::OpPool,
) -> crate::Result<Vec<OpRef>> {
    let output_roots = op_pool
        .iter()
        .enumerate()
        .filter_map(|(ix, op)| matches!(op, Op::Output { .. }).then_some(OpRef(ix)))
        .collect::<Vec<_>>();
    if !output_roots.is_empty() {
        return Ok(output_roots);
    }
    program
        .root
        .map(|root| vec![root])
        .ok_or(autark_error::Error::ProvidedEmptyProgram)
}

fn reachable_oprefs(op_pool: &mpera_frontend::op::OpPool, roots: &[OpRef]) -> Vec<bool> {
    let mut reachable = vec![false; op_pool.len()];
    roots
        .iter()
        .copied()
        .for_each(|opref| mark_reachable(op_pool, opref, &mut reachable));
    reachable
}

fn mark_reachable(op_pool: &mpera_frontend::op::OpPool, opref: OpRef, reachable: &mut [bool]) {
    if reachable[opref.0] {
        return;
    }
    reachable[opref.0] = true;
    if let Some(op) = op_pool.get(opref) {
        visit_dependencies(op, |dependency| {
            mark_reachable(op_pool, dependency, reachable)
        });
    }
}

fn visit_dependencies(op: &Op, mut visit: impl FnMut(OpRef)) {
    match op {
        Op::BinaryOp { lhs, rhs, .. } => {
            visit(*lhs);
            visit(*rhs);
        }
        Op::Reduce { on, .. }
        | Op::UnaryOp { on, .. }
        | Op::Slice { on, .. }
        | Op::Rolling { on, .. } => visit(*on),
        Op::Filter { on, mask } => {
            visit(*on);
            visit(*mask);
        }
        Op::Where { mask, then, else_ } => {
            visit(*mask);
            visit(*then);
            visit(*else_);
        }
        Op::Columns { on, .. } => {
            if let Some(on) = on {
                visit(*on);
            }
        }
        Op::Output { value, .. } => visit(*value),
        Op::Concat { who } => who.iter().copied().for_each(&mut visit),
        Op::GroupBy { keys, values } | Op::GroupReduce { keys, values, .. } => {
            visit(*keys);
            visit(*values);
        }
        Op::Join {
            left,
            right,
            left_on,
            right_on,
            ..
        } => {
            visit(*left);
            visit(*right);
            visit(*left_on);
            visit(*right_on);
        }
        Op::OrderBy { what, by, .. } => {
            visit(*what);
            visit(*by);
        }
        Op::Constant { .. } => {}
    }
}
