use crate::ast::nodes::{
    Address, AnalyzedNode, GlobalIdx, IdentifierBinding, Node,
    NodeKind, NodeMetrics, TypedNode, TypedPhase,
};
use crate::ast::types::{Identity, Purity, Value};
use std::collections::{HashMap, HashSet};
use std::rc::Rc;

pub struct Analyzer<'a> {
    root_purity: &'a [Purity],
    /// Stack of currently visiting lambdas to detect direct recursion.
    lambda_stack: Vec<Identity>,
    /// Map of global index to its Lambda identity if known.
    globals_to_lambdas: HashMap<GlobalIdx, Identity>,
    /// Set of identities that were found to be recursive.
    recursive_identities: HashSet<Identity>,
}

impl<'a> Analyzer<'a> {
    pub fn analyze(
        node: &TypedNode,
        root_purity: &'a [Purity],
    ) -> AnalyzedNode {
        let mut analyzer = Self {
            root_purity,
            lambda_stack: Vec::new(),
            globals_to_lambdas: HashMap::new(),
            recursive_identities: HashSet::new(),
        };

        // First pass: map globals to their lambda identities
        analyzer.collect_globals(node);

        // Second pass: full analysis (decorating TypedNode into AnalyzedNode)
        analyzer.visit(Rc::new(node.clone()))
    }

    fn collect_globals(&mut self, node: &TypedNode) {
        match &node.kind {
            NodeKind::Def {
                pattern,
                value,
                ..
            } => {
                if let NodeKind::Identifier {
                    binding: IdentifierBinding::Declaration { addr: Address::Global(global_index), .. },
                    ..
                } = &pattern.kind
                    && let NodeKind::Lambda { .. } = &value.kind
                {
                    self.globals_to_lambdas
                        .insert(*global_index, value.identity.clone());
                }
                self.collect_globals(value);
            }
            NodeKind::Block { exprs } | NodeKind::Program { exprs } => {
                for e in exprs {
                    self.collect_globals(e);
                }
            }
            _ => {
                node.kind
                    .for_each_child(|child| self.collect_globals(child));
            }
        }
    }

    fn visit(&mut self, node_rc: Rc<TypedNode>) -> AnalyzedNode {
        let node = &*node_rc;
        let mut is_recursive = false;

        let (new_kind, purity) = match &node.kind {
            // Propagate the declared purity of function constants so the constant
            // folder does not evaluate impure factory closures at compile time.
            NodeKind::Constant(v) => {
                let purity = if let Value::Function(f) = v { f.purity } else { Purity::Pure };
                (NodeKind::Constant(v.clone()), purity)
            }
            NodeKind::Nop => (NodeKind::Nop, Purity::Pure),

            NodeKind::Identifier { symbol, binding } => {
                let p = if let IdentifierBinding::Reference(Address::Global(idx)) = binding {
                    self.root_purity.get(idx.0 as usize).cloned().unwrap_or(Purity::Pure)
                } else {
                    Purity::Pure
                };
                (
                    NodeKind::Identifier {
                        symbol: symbol.clone(),
                        binding: binding.clone(),
                    },
                    p,
                )
            }

            NodeKind::FieldAccessor(k) => (NodeKind::FieldAccessor(*k), Purity::Pure),

            NodeKind::GetField { rec, field } => {
                let rec_m = self.visit(rec.clone());
                let p = rec_m.ty.purity;
                (
                    NodeKind::GetField {
                        rec: Rc::new(rec_m),
                        field: *field,
                    },
                    p,
                )
            }

            NodeKind::Assign { target, value, info } => {
                let target_m = self.visit(target.clone());
                let val_m = self.visit(value.clone());
                (
                    NodeKind::Assign {
                        target: Rc::new(target_m),
                        value: Rc::new(val_m),
                        info: info.clone(),
                    },
                    Purity::Impure,
                )
            }

            NodeKind::Def {
                pattern,
                value,
                info,
            } => {
                let pat_m = self.visit(pattern.clone());
                let val_m = self.visit(value.clone());
                (
                    NodeKind::Def {
                        pattern: Rc::new(pat_m),
                        value: Rc::new(val_m),
                        info: info.clone(),
                    },
                    Purity::Impure,
                )
            }

            NodeKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond_m = self.visit(cond.clone());
                let then_m = self.visit(then_br.clone());
                let else_m = else_br.as_ref().map(|e| self.visit(e.clone()));

                let mut p = cond_m.ty.purity.min(then_m.ty.purity);
                if let Some(ref em) = else_m {
                    p = p.min(em.ty.purity);
                }
                (
                    NodeKind::If {
                        cond: Rc::new(cond_m),
                        then_br: Rc::new(then_m),
                        else_br: else_m.map(Rc::new),
                    },
                    p,
                )
            }

            NodeKind::Lambda {
                params,
                body,
                info,
            } => {
                self.lambda_stack.push(node.identity.clone());
                let params_m = self.visit(params.clone());
                let body_m = self.visit(body.clone());
                self.lambda_stack.pop();

                is_recursive = self.recursive_identities.contains(&node.identity);
                (
                    NodeKind::Lambda {
                        params: Rc::new(params_m),
                        body: Rc::new(body_m),
                        info: info.clone(),
                    },
                    Purity::Pure,
                )
            }

            NodeKind::Call { callee, args } => {
                let callee_m = self.visit(callee.clone());
                let args_m = self.visit(args.clone());

                if let NodeKind::Identifier {
                    binding: IdentifierBinding::Reference(Address::Global(idx)),
                    ..
                } = &callee.kind
                    && let Some(lambda_id) = self.globals_to_lambdas.get(idx)
                    && self.lambda_stack.contains(lambda_id)
                {
                    self.recursive_identities.insert(lambda_id.clone());
                    is_recursive = true;
                }

                let p_func = if let NodeKind::Identifier {
                    binding: IdentifierBinding::Reference(Address::Global(idx)),
                    ..
                } = &callee.kind
                {
                    self.root_purity
                        .get(idx.0 as usize)
                        .cloned()
                        .unwrap_or(Purity::Impure)
                } else {
                    Purity::Impure
                };
                let p = callee_m.ty.purity.min(args_m.ty.purity).min(p_func);
                (
                    NodeKind::Call {
                        callee: Rc::new(callee_m),
                        args: Rc::new(args_m),
                    },
                    p,
                )
            }

            NodeKind::Again { args } => {
                let args_m = self.visit(args.clone());
                if let Some(lambda_id) = self.lambda_stack.last() {
                    self.recursive_identities.insert(lambda_id.clone());
                    is_recursive = true;
                }
                (
                    NodeKind::Again {
                        args: Rc::new(args_m),
                    },
                    Purity::Impure,
                )
            }
            NodeKind::Block { exprs } => {
                let mut new_exprs = Vec::with_capacity(exprs.len());
                let mut p = Purity::Pure;
                for e in exprs {
                    let em = self.visit(e.clone());
                    p = p.min(em.ty.purity);
                    new_exprs.push(Rc::new(em));
                }
                (NodeKind::Block { exprs: new_exprs }, p)
            }
            NodeKind::Program { exprs } => {
                let mut new_exprs = Vec::with_capacity(exprs.len());
                let mut p = Purity::Pure;
                for e in exprs {
                    let em = self.visit(e.clone());
                    p = p.min(em.ty.purity);
                    new_exprs.push(Rc::new(em));
                }
                (NodeKind::Program { exprs: new_exprs }, p)
            }

            NodeKind::Tuple { elements } => {
                let mut new_elements = Vec::with_capacity(elements.len());
                let mut p = Purity::Pure;
                for e in elements {
                    let em = self.visit(e.clone());
                    p = p.min(em.ty.purity);
                    new_elements.push(Rc::new(em));
                }
                (
                    NodeKind::Tuple {
                        elements: new_elements,
                    },
                    p,
                )
            }

            NodeKind::Record { fields, layout } => {
                let mut new_fields = Vec::with_capacity(fields.len());
                let mut p = Purity::Pure;
                for (key_node, val_node) in fields {
                    let km = self.visit(key_node.clone());
                    let vm = self.visit(val_node.clone());
                    p = p.min(vm.ty.purity);
                    new_fields.push((Rc::new(km), Rc::new(vm)));
                }
                (
                    NodeKind::Record {
                        fields: new_fields,
                        layout: layout.clone(),
                    },
                    p,
                )
            }

            NodeKind::Expansion {
                original_call,
                expanded,
            } => {
                let expanded_m = self.visit(expanded.clone());
                (
                    NodeKind::Expansion {
                        original_call: original_call.clone(),
                        expanded: Rc::new(expanded_m.clone()),
                    },
                    expanded_m.ty.purity,
                )
            }

            NodeKind::Extension(_) => (NodeKind::Nop, Purity::Impure),
            NodeKind::Error => (NodeKind::Error, Purity::Impure),

            // Syntax-only variants should not appear in typed phases
            NodeKind::MacroDecl { .. }
            | NodeKind::Template(_)
            | NodeKind::Placeholder(_)
            | NodeKind::Splice(_) => (NodeKind::Error, Purity::Impure),
        };

        Node {
            identity: node.identity.clone(),
            kind: new_kind,
            comments: node.comments.clone(),
            ty: NodeMetrics {
                original: node_rc,
                purity,
                is_recursive,
            },
        }
    }
}

trait NodeExt {
    fn for_each_child<F: FnMut(&TypedNode)>(&self, f: F);
}

impl NodeExt for NodeKind<TypedPhase> {
    fn for_each_child<F: FnMut(&TypedNode)>(&self, mut f: F) {
        match self {
            NodeKind::If {
                cond,
                then_br,
                else_br,
            } => {
                f(cond);
                f(then_br);
                if let Some(e) = else_br {
                    f(e);
                }
            }
            NodeKind::Def { pattern, value, .. } => {
                f(pattern);
                f(value);
            }
            NodeKind::Assign { target, value, .. } => {
                f(target);
                f(value);
            }
            NodeKind::GetField { rec, .. } => {
                f(rec);
            }
            NodeKind::Lambda { params, body, .. } => {
                f(params);
                f(body);
            }
            NodeKind::Call { callee, args } => {
                f(callee);
                f(args);
            }
            NodeKind::Block { exprs } | NodeKind::Program { exprs } => {
                for e in exprs {
                    f(e);
                }
            }
            NodeKind::Tuple { elements } => {
                for e in elements {
                    f(e);
                }
            }
            NodeKind::Record { fields, .. } => {
                for (key, val) in fields {
                    f(key);
                    f(val);
                }
            }
            NodeKind::Expansion { expanded, .. } => {
                f(expanded);
            }
            NodeKind::Again { args } => {
                f(args);
            }
            NodeKind::MacroDecl { params, body, .. } => {
                f(params);
                f(body);
            }
            NodeKind::Template(inner)
            | NodeKind::Placeholder(inner)
            | NodeKind::Splice(inner) => {
                f(inner);
            }
            // Leaf nodes — no children to visit.
            NodeKind::Nop
            | NodeKind::Constant(_)
            | NodeKind::Identifier { .. }
            | NodeKind::FieldAccessor(_)
            | NodeKind::Error
            | NodeKind::Extension(_) => {}
        }
    }
}