use crate::ast::nodes::{
    Address, AnalyzedNode, IdentifierBinding, Node, NodeKind, NodeMetrics,
};
use crate::ast::types::{Purity, RecordLayout, StaticType, Value};
use crate::ast::vm::GlobalStore;
use std::rc::Rc;

pub struct Folder<'a> {
    pub globals: &'a Option<GlobalStore>,
}

impl<'a> Folder<'a> {
    pub fn new(globals: &'a Option<GlobalStore>) -> 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: NodeKind::Constant(val.clone()),
            comments: template.comments.clone(),
            ty: ty.clone(),
        });
        Node {
            identity: template.identity.clone(),
            kind: NodeKind::Constant(val),
            comments: template.comments.clone(),
            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: NodeKind::Nop,
            comments: template.comments.clone(),
            ty: StaticType::Void,
        });
        Node {
            identity: template.identity.clone(),
            kind: NodeKind::Nop,
            comments: template.comments.clone(),
            ty: NodeMetrics {
                original: typed_original,
                purity: Purity::Pure,
                is_recursive: false,
            },
        }
    }

    pub fn try_fold_record(
        &self,
        layout: &std::sync::Arc<RecordLayout>,
        values: &[Rc<AnalyzedNode>],
        template: &AnalyzedNode,
    ) -> Option<AnalyzedNode> {
        let mut constant_values = Vec::with_capacity(values.len());

        for v_node in values {
            if let NodeKind::Constant(val) = &v_node.kind {
                constant_values.push(val.clone());
            } else {
                return None;
            }
        }

        // Recompute layout from concrete value types to eliminate stale TypeVars.
        let new_fields: Vec<_> = layout
            .fields
            .iter()
            .zip(values.iter())
            .map(|((keyword, _), v_node)| (*keyword, v_node.ty.original.ty.clone()))
            .collect();
        let new_layout = RecordLayout::get_or_create(new_fields);

        let record_val = Value::Record(new_layout, Rc::new(constant_values));
        Some(self.make_constant_node(record_val, template))
    }

    pub fn try_fold_pure(
        &self,
        callee: &AnalyzedNode,
        arg_nodes: &[Rc<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 NodeKind::Constant(val) = &node.kind {
                arg_values.push(val.clone());
            } else {
                return None;
            }
        }
        let func_val = match &callee.kind {
            NodeKind::Identifier {
                binding: IdentifierBinding::Reference(Address::Global(idx)),
                ..
            } => self.globals.as_ref()?.get(idx.0 as usize)?,
            NodeKind::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))
    }
}