//! L2 optimization passes: inline lets, constant fold, dead code elimination.
//!
//! Adapted from the platform evaluator. Platform-specific exemptions removed.

use std::sync::Arc;
use crate::closure_node::ClosureNode;
use pfcl_types::Value;

pub fn optimize(node: ClosureNode) -> ClosureNode {
    let node = inline_lets(node);
    let node = constant_fold(node);
    dead_let_elimination(node)
}

// ── Pass 1: inline single-use let bindings ────────────────────────────────────

fn inline_lets(node: ClosureNode) -> ClosureNode {
    match node {
        ClosureNode::Let { value, body } => {
            let value = inline_lets(*value);
            let body  = inline_lets(*body);
            let uses  = count_var_uses(&body, 0);
            if uses == 0 {
                shift_vars(body, 0, -1)
            } else if uses == 1 && is_cheap(&value) {
                let inlined = substitute_var(body, 0, &value);
                shift_vars(inlined, 0, -1)
            } else {
                ClosureNode::Let { value: Box::new(value), body: Box::new(body) }
            }
        }
        ClosureNode::Lambda { arity, body } =>
            ClosureNode::Lambda { arity, body: Box::new(inline_lets(*body)) },
        ClosureNode::Apply { func, args } =>
            ClosureNode::Apply {
                func: Box::new(inline_lets(*func)),
                args: args.into_iter().map(inline_lets).collect(),
            },
        ClosureNode::Match { scrutinee, arms } =>
            ClosureNode::Match {
                scrutinee: Box::new(inline_lets(*scrutinee)),
                arms: arms.into_iter().map(|a| crate::closure_node::ClosureArm {
                    pattern: a.pattern,
                    body: inline_lets(a.body),
                }).collect(),
            },
        ClosureNode::ListLit(elems) => ClosureNode::ListLit(elems.into_iter().map(inline_lets).collect()),
        ClosureNode::RecordLit(fields) =>
            ClosureNode::RecordLit(fields.into_iter().map(|(k,v)| (k, inline_lets(v))).collect()),
        ClosureNode::TupleLit(elems) => ClosureNode::TupleLit(elems.into_iter().map(inline_lets).collect()),
        ClosureNode::Constructor { name, args } =>
            ClosureNode::Constructor { name, args: args.into_iter().map(inline_lets).collect() },
        ClosureNode::FieldAccess { expr, fields } =>
            ClosureNode::FieldAccess { expr: Box::new(inline_lets(*expr)), fields },
        other => other,
    }
}

fn is_cheap(node: &ClosureNode) -> bool {
    matches!(node, ClosureNode::Literal(_) | ClosureNode::Var(_) | ClosureNode::Primitive {..} | ClosureNode::ComposedRef {..})
}

// ── Pass 2: constant folding ──────────────────────────────────────────────────

fn constant_fold(node: ClosureNode) -> ClosureNode {
    match node {
        ClosureNode::Apply { func, args } => {
            let func = Box::new(constant_fold(*func));
            let args: Vec<_> = args.into_iter().map(constant_fold).collect();
            // Fold: Primitive(f)(Literal(a), Literal(b), ...) → Literal(f([a,b,...]))
            if let ClosureNode::Primitive { func: f, ref name } = *func {
                if !is_higher_order_primitive(name) {
                    let lit_vals: Option<Vec<Value>> = args.iter()
                        .map(|a| if let ClosureNode::Literal(v) = a { Some(v.clone()) } else { None })
                        .collect();
                    if let Some(vals) = lit_vals {
                        return ClosureNode::Literal(f(&vals));
                    }
                }
            }
            ClosureNode::Apply { func, args }
        }
        ClosureNode::Lambda { arity, body } =>
            ClosureNode::Lambda { arity, body: Box::new(constant_fold(*body)) },
        ClosureNode::Let { value, body } =>
            ClosureNode::Let { value: Box::new(constant_fold(*value)), body: Box::new(constant_fold(*body)) },
        ClosureNode::Match { scrutinee, arms } =>
            ClosureNode::Match {
                scrutinee: Box::new(constant_fold(*scrutinee)),
                arms: arms.into_iter().map(|a| crate::closure_node::ClosureArm {
                    pattern: a.pattern,
                    body: constant_fold(a.body),
                }).collect(),
            },
        ClosureNode::ListLit(elems) => ClosureNode::ListLit(elems.into_iter().map(constant_fold).collect()),
        ClosureNode::RecordLit(fields) =>
            ClosureNode::RecordLit(fields.into_iter().map(|(k,v)| (k, constant_fold(v))).collect()),
        ClosureNode::TupleLit(elems) => ClosureNode::TupleLit(elems.into_iter().map(constant_fold).collect()),
        ClosureNode::Constructor { name, args } =>
            ClosureNode::Constructor { name, args: args.into_iter().map(constant_fold).collect() },
        ClosureNode::FieldAccess { expr, fields } =>
            ClosureNode::FieldAccess { expr: Box::new(constant_fold(*expr)), fields },
        other => other,
    }
}

fn is_higher_order_primitive(name: &str) -> bool {
    matches!(name,
        "list.map" | "list.filter" | "list.fold_left" | "list.fold_right" |
        "conditional" | "list.sort_by" | "list.take_while" | "list.drop_while" |
        "list.unfold" | "list.match"
    )
}

// ── Pass 3: dead let elimination ──────────────────────────────────────────────

fn dead_let_elimination(node: ClosureNode) -> ClosureNode {
    match node {
        ClosureNode::Let { value, body } => {
            let value = dead_let_elimination(*value);
            let body  = dead_let_elimination(*body);
            if count_var_uses(&body, 0) == 0 {
                shift_vars(body, 0, -1)
            } else {
                ClosureNode::Let { value: Box::new(value), body: Box::new(body) }
            }
        }
        ClosureNode::Lambda { arity, body } =>
            ClosureNode::Lambda { arity, body: Box::new(dead_let_elimination(*body)) },
        ClosureNode::Apply { func, args } =>
            ClosureNode::Apply {
                func: Box::new(dead_let_elimination(*func)),
                args: args.into_iter().map(dead_let_elimination).collect(),
            },
        other => other,
    }
}

// ── Helpers: variable counting, substitution, shifting ───────────────────────

pub fn count_var_uses(node: &ClosureNode, target: usize) -> usize {
    match node {
        ClosureNode::Var(i) => if *i == target { 1 } else { 0 },
        ClosureNode::Apply { func, args } =>
            count_var_uses(func, target) + args.iter().map(|a| count_var_uses(a, target)).sum::<usize>(),
        ClosureNode::Lambda { arity, body } => count_var_uses(body, target + arity),
        ClosureNode::Let { value, body } =>
            count_var_uses(value, target) + count_var_uses(body, target + 1),
        ClosureNode::Match { scrutinee, arms } => {
            count_var_uses(scrutinee, target)
                + arms.iter().map(|a| {
                    let bindings = pattern_binding_count(&a.pattern);
                    count_var_uses(&a.body, target + bindings)
                }).sum::<usize>()
        }
        ClosureNode::ListLit(elems) => elems.iter().map(|e| count_var_uses(e, target)).sum(),
        ClosureNode::RecordLit(fields) => fields.iter().map(|(_,v)| count_var_uses(v, target)).sum(),
        ClosureNode::TupleLit(elems) => elems.iter().map(|e| count_var_uses(e, target)).sum(),
        ClosureNode::Constructor { args, .. } => args.iter().map(|a| count_var_uses(a, target)).sum(),
        ClosureNode::FieldAccess { expr, .. } => count_var_uses(expr, target),
        _ => 0,
    }
}

pub fn pattern_binding_count(pat: &crate::closure_node::ClosurePattern) -> usize {
    use crate::closure_node::ClosurePattern;
    match pat {
        ClosurePattern::Wildcard | ClosurePattern::IntPat(_) |
        ClosurePattern::StringPat(_) | ClosurePattern::BoolPat(_) |
        ClosurePattern::NothingPat => 0,
        ClosurePattern::Var => 1,
        ClosurePattern::Constructor { args, .. } => args.iter().map(pattern_binding_count).sum(),
        ClosurePattern::TuplePat(pats) => pats.iter().map(pattern_binding_count).sum(),
    }
}

fn substitute_var(node: ClosureNode, target: usize, replacement: &ClosureNode) -> ClosureNode {
    match node {
        ClosureNode::Var(i) if i == target => replacement.clone(),
        ClosureNode::Var(i) => ClosureNode::Var(i),
        ClosureNode::Apply { func, args } =>
            ClosureNode::Apply {
                func: Box::new(substitute_var(*func, target, replacement)),
                args: args.into_iter().map(|a| substitute_var(a, target, replacement)).collect(),
            },
        ClosureNode::Lambda { arity, body } =>
            ClosureNode::Lambda { arity, body: Box::new(substitute_var(*body, target + arity, replacement)) },
        ClosureNode::Let { value, body } =>
            ClosureNode::Let {
                value: Box::new(substitute_var(*value, target, replacement)),
                body:  Box::new(substitute_var(*body, target + 1, replacement)),
            },
        other => other,
    }
}

fn shift_vars(node: ClosureNode, cutoff: usize, delta: i64) -> ClosureNode {
    match node {
        ClosureNode::Var(i) => {
            if i >= cutoff {
                ClosureNode::Var((i as i64 + delta) as usize)
            } else {
                ClosureNode::Var(i)
            }
        }
        ClosureNode::Apply { func, args } =>
            ClosureNode::Apply {
                func: Box::new(shift_vars(*func, cutoff, delta)),
                args: args.into_iter().map(|a| shift_vars(a, cutoff, delta)).collect(),
            },
        ClosureNode::Lambda { arity, body } =>
            ClosureNode::Lambda { arity, body: Box::new(shift_vars(*body, cutoff + arity, delta)) },
        ClosureNode::Let { value, body } =>
            ClosureNode::Let {
                value: Box::new(shift_vars(*value, cutoff, delta)),
                body:  Box::new(shift_vars(*body, cutoff + 1, delta)),
            },
        other => other,
    }
}
