use crate::ast::compiler::bound_nodes::{BoundKind, CompilerPhase, Node};
use crate::ast::types::Identity;
use std::collections::HashMap;
use std::rc::Rc;

pub struct CapturePass;

impl CapturePass {
    pub fn apply<P: CompilerPhase>(node: Node<P>, capture_map: &HashMap<Identity, Vec<Identity>>) -> Node<P> {
        Self::transform(node, capture_map)
    }

    fn transform<P: CompilerPhase>(mut node: Node<P>, capture_map: &HashMap<Identity, Vec<Identity>>) -> Node<P> {
        node.kind = match node.kind {
            BoundKind::Define {
                name,
                addr,
                kind,
                value,
                mut captured_by,
            } => {
                if let Some(capturers) = capture_map.get(&node.identity) {
                    captured_by.extend(capturers.iter().cloned());
                }
                BoundKind::Define {
                    name,
                    addr,
                    kind,
                    value: Rc::new(Self::transform(value.as_ref().clone(), capture_map)),
                    captured_by,
                }
            }

            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
            } => BoundKind::Lambda {
                params: Rc::new(Self::transform(params.as_ref().clone(), capture_map)),
                upvalues,
                body: Rc::new(Self::transform(body.as_ref().clone(), capture_map)),
                positional_count,
            },

            BoundKind::Call { callee, args } => BoundKind::Call {
                callee: Rc::new(Self::transform(callee.as_ref().clone(), capture_map)),
                args: Rc::new(Self::transform(args.as_ref().clone(), capture_map)),
            },

            BoundKind::Again { args } => BoundKind::Again {
                args: Rc::new(Self::transform(args.as_ref().clone(), capture_map)),
            },

            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => BoundKind::If {
                cond: Rc::new(Self::transform(cond.as_ref().clone(), capture_map)),
                then_br: Rc::new(Self::transform(then_br.as_ref().clone(), capture_map)),
                else_br: else_br
                    .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map))),
            },

            BoundKind::Destructure { pattern, value } => BoundKind::Destructure {
                pattern: Rc::new(Self::transform(pattern.as_ref().clone(), capture_map)),
                value: Rc::new(Self::transform(value.as_ref().clone(), capture_map)),
            },

            BoundKind::Pipe {
                inputs,
                lambda,
                out_type,
            } => {
                let mut t_inputs = Vec::with_capacity(inputs.len());
                for input in inputs {
                    t_inputs.push(Rc::new(Self::transform(input.as_ref().clone(), capture_map)));
                }
                BoundKind::Pipe {
                    inputs: t_inputs,
                    lambda: Rc::new(Self::transform(lambda.as_ref().clone(), capture_map)),
                    out_type,
                }
            }

            BoundKind::Block { exprs } => BoundKind::Block {
                exprs: exprs
                    .into_iter()
                    .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map)))
                    .collect(),
            },

            BoundKind::Tuple { elements } => BoundKind::Tuple {
                elements: elements
                    .into_iter()
                    .map(|e| Rc::new(Self::transform(e.as_ref().clone(), capture_map)))
                    .collect(),
            },

            BoundKind::Record { layout, values } => BoundKind::Record {
                layout,
                values: values
                    .into_iter()
                    .map(|v| Rc::new(Self::transform(v.as_ref().clone(), capture_map)))
                    .collect(),
            },

            BoundKind::Expansion {
                original_call,
                bound_expanded,
            } => BoundKind::Expansion {
                original_call,
                bound_expanded: Rc::new(Self::transform(
                    bound_expanded.as_ref().clone(),
                    capture_map,
                )),
            },

            BoundKind::GetField { rec, field } => BoundKind::GetField {
                rec: Rc::new(Self::transform(rec.as_ref().clone(), capture_map)),
                field,
            },

            other => other,
        };
        node
    }
}