use crate::ast::compiler::bound_nodes::{BoundKind, BoundNode};
use crate::ast::types::Identity;
use std::collections::HashMap;

pub struct CapturePass;

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

    fn transform(mut node: BoundNode, capture_map: &HashMap<Identity, Vec<Identity>>) -> BoundNode {
        match node.kind {
            BoundKind::Define {
                name,
                addr,
                kind,
                value,
                ..
            } => {
                let captured_by = capture_map.get(&node.identity).cloned().unwrap_or_default();
                node.kind = BoundKind::Define {
                    name,
                    addr,
                    kind,
                    value: Box::new(Self::transform(*value, capture_map)),
                    captured_by,
                };
            }

            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => {
                node.kind = BoundKind::If {
                    cond: Box::new(Self::transform(*cond, capture_map)),
                    then_br: Box::new(Self::transform(*then_br, capture_map)),
                    else_br: else_br.map(|e| Box::new(Self::transform(*e, capture_map))),
                };
            }

            BoundKind::Set { addr, value } => {
                node.kind = BoundKind::Set {
                    addr,
                    value: Box::new(Self::transform(*value, capture_map)),
                };
            }

            BoundKind::FieldAccessor(_) => {}

            BoundKind::GetField { rec, field } => {
                node.kind = BoundKind::GetField {
                    rec: Box::new(Self::transform(*rec, capture_map)),
                    field,
                };
            }

            BoundKind::Destructure { pattern, value } => {
                node.kind = BoundKind::Destructure {
                    pattern: Box::new(Self::transform(*pattern, capture_map)),
                    value: Box::new(Self::transform(*value, capture_map)),
                };
            }

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

            BoundKind::Call { callee, args } => {
                node.kind = BoundKind::Call {
                    callee: Box::new(Self::transform(*callee, capture_map)),
                    args: Box::new(Self::transform(*args, capture_map)),
                };
            }

            BoundKind::Again { args } => {
                node.kind = BoundKind::Again {
                    args: Box::new(Self::transform(*args, capture_map)),
                };
            }
            BoundKind::Pipe { inputs, lambda, out_type } => {
                let mut t_inputs = Vec::with_capacity(inputs.len());
                for input in inputs {
                    t_inputs.push(Self::transform(input, capture_map));
                }
                node.kind = BoundKind::Pipe {
                    inputs: t_inputs,
                    lambda: Box::new(Self::transform(*lambda, capture_map)),
                    out_type: out_type.clone(),
                };
            }
            BoundKind::Block { exprs } => {
                node.kind = BoundKind::Block {
                    exprs: exprs
                        .into_iter()
                        .map(|e| Self::transform(e, capture_map))
                        .collect(),
                };
            }

            BoundKind::Tuple { elements } => {
                node.kind = BoundKind::Tuple {
                    elements: elements
                        .into_iter()
                        .map(|e| Self::transform(e, capture_map))
                        .collect(),
                };
            }

            BoundKind::Record { layout, values } => {
                node.kind = BoundKind::Record {
                    layout,
                    values: values
                        .into_iter()
                        .map(|v| Self::transform(v, capture_map))
                        .collect(),
                };
            }

            BoundKind::Expansion {
                original_call,
                bound_expanded,
            } => {
                node.kind = BoundKind::Expansion {
                    original_call,
                    bound_expanded: Box::new(Self::transform(*bound_expanded, capture_map)),
                };
            }

            BoundKind::Nop
            | BoundKind::Constant(_)
            | BoundKind::Get { .. }
            | BoundKind::Extension(_) => {}
        }
        node
    }
}