use crate::ast::compiler::bound_nodes::{Address, AnalyzedNode, BoundKind, NodeMetrics, TypedNode};
use crate::ast::types::Purity;
use std::collections::{HashMap, HashSet};
use std::rc::Rc;

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

impl<'a> Analyzer<'a> {
    pub fn analyze(node: &TypedNode, global_purity: &'a HashMap<u32, Purity>) -> AnalyzedNode {
        let mut analyzer = Self {
            global_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 {
            BoundKind::DefGlobal {
                global_index,
                value,
                ..
            } => {
                if let BoundKind::Lambda { .. } = &value.kind {
                    self.globals_to_lambdas
                        .insert(*global_index, value.identity.clone());
                }
                self.collect_globals(value);
            }
            BoundKind::Block { 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 {
            BoundKind::Constant(v) => (BoundKind::Constant(v.clone()), Purity::Pure),
            BoundKind::Nop => (BoundKind::Nop, Purity::Pure),
            BoundKind::Parameter { name, slot } => (
                BoundKind::Parameter {
                    name: name.clone(),
                    slot: *slot,
                },
                Purity::Pure,
            ),

            BoundKind::Get { addr, name } => {
                let p = match addr {
                    Address::Global(idx) => {
                        self.global_purity.get(idx).cloned().unwrap_or(Purity::Pure)
                    }
                    _ => Purity::Pure,
                };
                (
                    BoundKind::Get {
                        addr: *addr,
                        name: name.clone(),
                    },
                    p,
                )
            }

            BoundKind::Set { addr, value } => {
                let val_m = self.visit(Rc::new((**value).clone()));
                (
                    BoundKind::Set {
                        addr: *addr,
                        value: Box::new(val_m),
                    },
                    Purity::Impure,
                )
            }

            BoundKind::DefLocal {
                name,
                slot,
                value,
                captured_by,
            } => {
                let val_m = self.visit(Rc::new((**value).clone()));
                let p = val_m.ty.purity;
                (
                    BoundKind::DefLocal {
                        name: name.clone(),
                        slot: *slot,
                        value: Box::new(val_m),
                        captured_by: captured_by.clone(),
                    },
                    p,
                )
            }

            BoundKind::DefGlobal {
                name,
                global_index,
                value,
            } => {
                let val_m = self.visit(Rc::new((**value).clone()));
                let p = val_m.ty.purity;
                (
                    BoundKind::DefGlobal {
                        name: name.clone(),
                        global_index: *global_index,
                        value: Box::new(val_m),
                    },
                    p,
                )
            }

            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => {
                let cond_m = self.visit(Rc::new((**cond).clone()));
                let then_m = self.visit(Rc::new((**then_br).clone()));
                let else_m = else_br.as_ref().map(|e| self.visit(Rc::new((**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);
                }
                (
                    BoundKind::If {
                        cond: Box::new(cond_m),
                        then_br: Box::new(then_m),
                        else_br: else_m.map(Box::new),
                    },
                    p,
                )
            }

            BoundKind::Lambda {
                params,
                upvalues,
                body,
                positional_count,
            } => {
                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);
                (
                    BoundKind::Lambda {
                        params: Rc::new(params_m),
                        upvalues: upvalues.clone(),
                        body: Rc::new(body_m),
                        positional_count: *positional_count,
                    },
                    Purity::Pure,
                )
            }

            BoundKind::DefDestructure { pattern, value } => {
                let pat_m = self.visit(Rc::new((**pattern).clone()));
                let val_m = self.visit(Rc::new((**value).clone()));
                (
                    BoundKind::DefDestructure {
                        pattern: Box::new(pat_m),
                        value: Box::new(val_m),
                    },
                    Purity::Impure,
                )
            }

            BoundKind::Call { callee, args } => {
                let callee_m = self.visit(Rc::new((**callee).clone()));
                let args_m = self.visit(Rc::new((**args).clone()));

                if let BoundKind::Get {
                    addr: 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 BoundKind::Get {
                    addr: Address::Global(idx),
                    ..
                } = &callee.kind
                {
                    self.global_purity
                        .get(idx)
                        .cloned()
                        .unwrap_or(Purity::Impure)
                } else {
                    Purity::Impure
                };
                let p = callee_m.ty.purity.min(args_m.ty.purity).min(p_func);
                (
                    BoundKind::Call {
                        callee: Box::new(callee_m),
                        args: Box::new(args_m),
                    },
                    p,
                )
            }

            BoundKind::Again { args } => {
                let args_m = self.visit(Rc::new((**args).clone()));
                if let Some(lambda_id) = self.lambda_stack.last() {
                    self.recursive_identities.insert(lambda_id.clone());
                    is_recursive = true;
                }
                (
                    BoundKind::Again {
                        args: Box::new(args_m),
                    },
                    Purity::Impure,
                )
            }

            BoundKind::Block { exprs } => {
                let mut new_exprs = Vec::with_capacity(exprs.len());
                let mut p = Purity::Pure;
                for e in exprs {
                    let em = self.visit(Rc::new(e.clone()));
                    p = p.min(em.ty.purity);
                    new_exprs.push(em);
                }
                (BoundKind::Block { exprs: new_exprs }, p)
            }

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

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

            BoundKind::Expansion {
                original_call,
                bound_expanded,
            } => {
                let expanded_m = self.visit(Rc::new((**bound_expanded).clone()));
                (
                    BoundKind::Expansion {
                        original_call: original_call.clone(),
                        bound_expanded: Box::new(expanded_m.clone()),
                    },
                    expanded_m.ty.purity,
                )
            }

            BoundKind::Extension(_) => (BoundKind::Nop, Purity::Impure),
        };

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

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

impl NodeExt for BoundKind<crate::ast::types::StaticType> {
    fn for_each_child<F: FnMut(&TypedNode)>(&self, mut f: F) {
        match self {
            BoundKind::If {
                cond,
                then_br,
                else_br,
            } => {
                f(cond);
                f(then_br);
                if let Some(e) = else_br {
                    f(e);
                }
            }
            BoundKind::DefLocal { value, .. }
            | BoundKind::DefGlobal { value, .. }
            | BoundKind::Set { value, .. } => {
                f(value);
            }
            BoundKind::Lambda { params, body, .. } => {
                f(params);
                f(body);
            }
            BoundKind::Call { callee, args } => {
                f(callee);
                f(args);
            }
            BoundKind::Block { exprs } => {
                for e in exprs {
                    f(e);
                }
            }
            BoundKind::Tuple { elements } => {
                for e in elements {
                    f(e);
                }
            }
            BoundKind::Record { fields } => {
                for (k, v) in fields {
                    f(k);
                    f(v);
                }
            }
            BoundKind::Expansion { bound_expanded, .. } => {
                f(bound_expanded);
            }
            _ => {}
        }
    }
}