use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, NodeMetrics};
use crate::ast::nodes::Node;
use crate::ast::types::{Purity, StaticType, Value, RecordLayout};
use std::cell::RefCell;
use std::rc::Rc;

pub struct Folder<'a> {
    pub globals: &'a Option<Rc<RefCell<Vec<Value>>>>,
}

impl<'a> Folder<'a> {
    pub fn new(globals: &'a Option<Rc<RefCell<Vec<Value>>>>) -> Self {
        Self { globals }
    }

    pub fn make_constant_node(&self, val: Value, template: &AnalyzedNode) -> AnalyzedNode {
        let ty = val.static_type();
        let typed_original = Rc::new(Node {
            identity: template.identity.clone(),
            kind: BoundKind::Constant(val.clone()),
            ty: ty.clone(),
        });
        Node {
            identity: template.identity.clone(),
            kind: BoundKind::Constant(val),
            ty: NodeMetrics {
                original: typed_original,
                purity: Purity::Pure,
                is_recursive: false,
            },
        }
    }

    pub fn make_nop_node(&self, template: &AnalyzedNode) -> AnalyzedNode {
        let typed_original = Rc::new(Node {
            identity: template.identity.clone(),
            kind: BoundKind::Nop,
            ty: StaticType::Void,
        });
        Node {
            identity: template.identity.clone(),
            kind: BoundKind::Nop,
            ty: NodeMetrics {
                original: typed_original,
                purity: Purity::Pure,
                is_recursive: false,
            },
        }
    }

    pub fn try_fold_record(
        &self,
        fields: &[(AnalyzedNode, AnalyzedNode)],
        template: &AnalyzedNode,
    ) -> Option<AnalyzedNode> {
        let mut layout_fields = Vec::with_capacity(fields.len());
        let mut values = Vec::with_capacity(fields.len());

        for (k_node, v_node) in fields {
            if let BoundKind::Constant(Value::Keyword(kw)) = &k_node.kind {
                if let BoundKind::Constant(val) = &v_node.kind {
                    layout_fields.push((*kw, val.static_type()));
                    values.push(val.clone());
                } else {
                    return None;
                }
            } else {
                return None;
            }
        }

        let layout = RecordLayout::get_or_create(layout_fields);
        let record_val = Value::Record(layout, Rc::new(values));
        Some(self.make_constant_node(record_val, template))
    }

    pub fn try_fold_pure(
        &self,
        callee: &AnalyzedNode,
        arg_nodes: &[AnalyzedNode],
    ) -> Option<AnalyzedNode> {
        if callee.ty.purity < Purity::Pure {
            return None;
        }

        let mut arg_values = Vec::with_capacity(arg_nodes.len());
        for node in arg_nodes {
            if let BoundKind::Constant(val) = &node.kind {
                arg_values.push(val.clone());
            } else {
                return None;
            }
        }
        let func_val = match &callee.kind {
            BoundKind::Get {
                addr: Address::Global(idx),
                ..
            } => self.globals.as_ref()?.borrow().get(idx.0 as usize)?.clone(),
            BoundKind::Constant(val) => val.clone(),
            _ => return None,
        };
        let result = match func_val {
            Value::Function(f) => (f.func)(arg_values),
            _ => return None,
        };
        Some(self.make_constant_node(result, callee))
    }
}