//! Lowering of `Term::Lam` and the surrounding closure machinery.
//!
//! `lower_lambda` walks the body to find free vars (captures), lifts
//! the body into a top-level thunk fn `@ail_<m>_<def>_lam<id>`,
//! allocates a heap env for the captures, and returns a closure
//! pair `{ thunk_ptr, env_ptr }`. Under `--alloc=rc` it also defers
//! per-pair drop-fn IR via `build_env_drop_fn` /
//! `build_pair_drop_fn` (in `drop.rs`) and registers the pair-drop
//! symbol in `closure_drops` so the surrounding `Term::Let` close
//! can call it.
//!
//! The two helpers `collect_captures` and `pattern_bound_names` are
//! free-variable analyses needed by the env-construction step. They
//! stay private to this module — `collect_captures` is recursive
//! and only ever entered through `lower_lambda`'s prelude;
//! `pattern_bound_names` only feeds the match-arm branch of
//! `collect_captures`.
//!
//! Cross-references back into the parent module:
//!   - `start_block`, `lower_term`, `fresh_ssa`,
//!     `build_env_drop_fn`, `build_pair_drop_fn`: lifecycle and
//!     drop-emission primitives, all `pub(crate)`.
//!   - `synth::llvm_type`, `synth::fn_sig_from_type`: free-function
//!     helpers for IR shaping.

use ailang_core::ast::*;
use std::collections::BTreeSet;

use super::escape;
use super::synth::{fn_sig_from_type, llvm_type};
use super::{AllocStrategy, CodegenError, Emitter, FnSig, Result};

impl<'a> Emitter<'a> {
    /// lower a `Term::Lam`. Walks the body to find captures
    /// (free vars relative to the enclosing scope, minus builtins and
    /// top-level fns), allocates a heap env for the captures, lifts the
    /// body to a top-level thunk fn `@ail_<m>_<def>_lam<id>`, and
    /// returns a closure-pair pointer that pairs the thunk with the env.
    pub(crate) fn lower_lambda(
        &mut self,
        lam_params: &[String],
        lam_param_tys: &[Type],
        lam_ret_ty: &Type,
        lam_body: &Term,
        term_ptr: usize,
    ) -> Result<(String, String)> {
        let llvm_param_tys: Vec<String> =
            lam_param_tys.iter().map(llvm_type).collect::<Result<_>>()?;
        let llvm_ret = llvm_type(lam_ret_ty)?;

        // 1. Capture analysis. The "bound" set seeds with everything
        //    that's a local in the enclosing scope OR a lambda param —
        //    but params are bound only INSIDE the body, not before. So
        //    pass them through as part of the recursion.
        let top_level: BTreeSet<String> = self
            .module_user_fns
            .get(self.module_name)
            .map(|m| m.keys().cloned().collect())
            .unwrap_or_default();
        let builtins_owned: Vec<String> = ailang_check::builtins::value_names()
            .into_iter()
            .map(|s| s.to_string())
            .collect();
        let builtins: BTreeSet<&str> =
            builtins_owned.iter().map(|s| s.as_str()).collect();

        let mut bound: BTreeSet<String> = BTreeSet::new();
        for p in lam_params {
            bound.insert(p.clone());
        }
        let mut captures: Vec<String> = Vec::new();
        let mut captures_set: BTreeSet<String> = BTreeSet::new();
        Self::collect_captures(
            lam_body,
            &mut bound,
            &mut captures,
            &mut captures_set,
            &builtins,
            &top_level,
        );

        // Outer scope provides every capture as a local — assert and
        // pull SSA + LLVM type from `self.locals`. Unknown captures
        // would mean the typechecker let through an unbound var, so a
        // hard internal error is right.
        // Tuple: (name, outer_ssa, llvm_type, ail_type, optional_fn_sig).
        //
        // The typechecker rejects lambda-captures-of-loop-binder via
        // `CheckError::LoopBinderCapturedByLambda`. Reaching the
        // `None` arm below means typecheck was skipped or a bug let
        // the AST through — both are bugs in upstream layers, so
        // panic rather than return an Internal error.
        let mut cap_meta: Vec<(String, String, String, Type, Option<FnSig>)> = Vec::new();
        for c in &captures {
            let (_, outer_ssa, lty, ail_ty) = self
                .locals
                .iter()
                .rev()
                .find(|(n, _, _, _)| n == c)
                .unwrap_or_else(|| {
                    unreachable!(
                        "lambda capture `{c}` not in locals — typecheck should have rejected via LoopBinderCapturedByLambda",
                    )
                })
                .clone();
            let sig = self.ssa_fn_sigs.get(&outer_ssa).cloned();
            cap_meta.push((c.clone(), outer_ssa, lty, ail_ty, sig));
        }

        // 2. Pick a thunk name and switch the emitter into "thunk
        //    emission" mode by saving and resetting per-fn state.
        let lam_id = self.lam_counter;
        self.lam_counter += 1;
        let parent = self.current_def.clone();
        let thunk_name = format!("{parent}_lam{lam_id}");
        let thunk_symbol = format!("@ail_{m}_{thunk_name}", m = self.module_name);

        let saved_body = std::mem::take(&mut self.body);
        let saved_locals = std::mem::take(&mut self.locals);
        let saved_counter = self.counter;
        let saved_block = std::mem::take(&mut self.current_block);
        let saved_terminated = self.block_terminated;
        self.block_terminated = false;
        let saved_sigs = std::mem::take(&mut self.ssa_fn_sigs);
        // a lambda thunk is its own fn frame for escape
        // analysis. Save the outer fn's non-escape set and recompute
        // for the lambda body. Restored at the end alongside the rest
        // of the saved emitter state.
        let saved_non_escape = std::mem::take(&mut self.non_escape);
        self.non_escape = escape::analyze_fn_body(lam_body);
        // a lambda thunk is its own fn frame for move
        // tracking — the outer fn's binders are not in scope inside
        // the thunk body (only its captures + params). Save and reset.
        let saved_moved = std::mem::take(&mut self.moved_slots);
        // a lambda thunk is its own fn frame for the
        // mut-var bookkeeping introduced in iter mut.3 — the outer
        // fn's loop binders are not in scope inside the thunk (a
        // lambda body cannot reference enclosing loop binders; the
        // typechecker rejects such captures, and the
        // `binder_allocas` map is thunk-local). Save and reset the
        // three fields so the thunk emits its own hoisted-alloca
        // block at the thunk's entry, not the outer fn's.
        let saved_binder_allocas = std::mem::take(&mut self.binder_allocas);
        // loop-recur iter 3: a lambda thunk is its own fn frame —
        // a `recur` cannot target a loop enclosing the lambda
        // (iter-2 typecheck already rejects it; codegen resets so
        // the thunk's own loops work and the outer frames cannot
        // leak in). Save + reset (mem::take empties the Vec),
        // restored below alongside the binder-alloca map.
        let saved_loop_frames = std::mem::take(&mut self.loop_frames);
        let saved_pending_allocas = std::mem::take(&mut self.pending_entry_allocas);
        let saved_entry_marker = self.entry_block_end_marker.take();
        // A lambda thunk is its own fn frame for
        // param-mode lookup. The outer fn's params are not in scope
        // inside the thunk; the thunk's own params are pushed below
        // and (currently) carry no mode annotation, so they default
        // to `Implicit` — `lower_match`'s Iter A gate will skip arm-
        // close pattern-binder dec for matches on lambda params,
        // mirroring the fn-level Implicit-param treatment.
        let saved_param_modes = std::mem::take(&mut self.current_param_modes);
        for pname in lam_params.iter() {
            self.current_param_modes
                .insert(pname.clone(), ParamMode::Implicit);
        }
        // Lambdas inside lambdas are fine: they get their own counter
        // namespace within the enclosing thunk. They share the
        // `deferred_thunks` queue (top-level body owns it).
        self.counter = 0;

        // Thunk header: `(ptr %env, params...)`.
        let mut thunk_sig = format!("define {ret} {thunk_symbol}(ptr %env", ret = llvm_ret);
        for (pname, pty) in lam_params.iter().zip(llvm_param_tys.iter()) {
            thunk_sig.push_str(&format!(", {pty} %arg_{pname}"));
        }
        thunk_sig.push_str(") {\n");
        self.body.push_str(&thunk_sig);
        self.start_block("entry");

        // Unpack captures back into named locals at the start of the
        // body. Layout: 8 bytes per slot, regardless of LLVM type
        // (typed load reads only the needed bytes; padding is wasted
        // but uniform).
        for (i, (cname, _outer, cty, c_ail, sig_opt)) in cap_meta.iter().enumerate() {
            let offset = (i * 8) as i64;
            let slot = self.fresh_ssa();
            let val = self.fresh_ssa();
            self.body.push_str(&format!(
                "  {slot} = getelementptr inbounds i8, ptr %env, i64 {offset}\n"
            ));
            self.body
                .push_str(&format!("  {val} = load {cty}, ptr {slot}\n"));
            self.locals.push((
                cname.clone(),
                val.clone(),
                cty.clone(),
                c_ail.clone(),
            ));
            if let Some(sig) = sig_opt {
                self.ssa_fn_sigs.insert(val, sig.clone());
            }
        }

        // Push lambda params as locals (after captures so shadowing
        // honors lexical order — params win).
        for ((pname, pty), pty_ail) in lam_params
            .iter()
            .zip(llvm_param_tys.iter())
            .zip(lam_param_tys.iter())
        {
            let pssa = format!("%arg_{pname}");
            self.locals.push((
                pname.clone(),
                pssa.clone(),
                pty.clone(),
                pty_ail.clone(),
            ));
            if let Some(fs) = fn_sig_from_type(pty_ail) {
                self.ssa_fn_sigs.insert(pssa, fs);
            }
        }

        let (body_v, body_ty) = self.lower_term(lam_body)?;
        if !self.block_terminated {
            if body_ty != llvm_ret {
                return Err(CodegenError::Internal(format!(
                    "lambda `{thunk_name}`: body type {body_ty} != return type {llvm_ret}"
                )));
            }
            self.body
                .push_str(&format!("  ret {body_ty} {body_v}\n}}\n\n"));
        } else {
            self.body.push_str("}\n\n");
        }

        // Splice any hoisted loop-binder allocas into the thunk's
        // entry block at the marker captured during the thunk's
        // `start_block("entry")` above. Empty `pending` is a no-op
        // (the thunk had no `Term::Loop` in its body).
        if !self.pending_entry_allocas.is_empty() {
            let marker = self.entry_block_end_marker.ok_or_else(|| {
                CodegenError::Internal(
                    "lambda thunk: pending_entry_allocas accumulated but no entry-block \
                     marker was captured".into(),
                )
            })?;
            let allocas = std::mem::take(&mut self.pending_entry_allocas);
            self.body.insert_str(marker, &allocas);
        }

        // Park the thunk text in the deferred queue and restore outer
        // emitter state.
        let thunk_text = std::mem::take(&mut self.body);
        self.deferred_thunks.push(thunk_text);
        self.body = saved_body;
        self.locals = saved_locals;
        self.counter = saved_counter;
        self.current_block = saved_block;
        self.block_terminated = saved_terminated;
        self.ssa_fn_sigs = saved_sigs;
        self.non_escape = saved_non_escape;
        self.moved_slots = saved_moved;
        self.current_param_modes = saved_param_modes;
        // Restore outer fn's loop-binder bookkeeping.
        self.binder_allocas = saved_binder_allocas;
        self.loop_frames = saved_loop_frames;
        self.pending_entry_allocas = saved_pending_allocas;
        self.entry_block_end_marker = saved_entry_marker;

        // 3. Emit allocation + capture filling + closure-pair packing
        //    in the OUTER body. Captures use 8 bytes each; closure-pair
        //    is 16 bytes.
        // env and closure-pair share the Lam term's escape
        // status — they have parallel lifetimes. If the closure pair
        // does not escape (only used as a callee in the outer body),
        // both can be `alloca`'d. The escape-analysis lookup runs
        // against the outer fn's `non_escape` set (already restored
        // above).
        let lam_local = self.non_escape.contains(&term_ptr);
        let env_size = (cap_meta.len() * 8) as i64;
        let env_ssa = if env_size > 0 {
            let env = self.fresh_ssa();
            if lam_local {
                self.body.push_str(&format!(
                    "  {env} = alloca i8, i64 {env_size}, align 8\n"
                ));
            } else {
                self.body.push_str(&format!(
                    "  {env} = call ptr @{}(i64 {env_size})\n",
                    self.alloc.fn_name()
                ));
            }
            for (i, (_cname, outer_ssa, cty, _c_ail, _sig)) in cap_meta.iter().enumerate() {
                let offset = (i * 8) as i64;
                let slot = self.fresh_ssa();
                self.body.push_str(&format!(
                    "  {slot} = getelementptr inbounds i8, ptr {env}, i64 {offset}\n"
                ));
                self.body.push_str(&format!(
                    "  store {cty} {outer_ssa}, ptr {slot}\n"
                ));
            }
            env
        } else {
            "null".into()
        };

        let clos = self.fresh_ssa();
        if lam_local {
            self.body
                .push_str(&format!("  {clos} = alloca i8, i64 16, align 8\n"));
        } else {
            self.body.push_str(&format!(
                "  {clos} = call ptr @{}(i64 16)\n",
                self.alloc.fn_name()
            ));
        }
        let cs_t = self.fresh_ssa();
        self.body.push_str(&format!(
            "  {cs_t} = getelementptr inbounds {{ ptr, ptr }}, ptr {clos}, i64 0, i32 0\n"
        ));
        self.body
            .push_str(&format!("  store ptr {thunk_symbol}, ptr {cs_t}\n"));
        let cs_e = self.fresh_ssa();
        self.body.push_str(&format!(
            "  {cs_e} = getelementptr inbounds {{ ptr, ptr }}, ptr {clos}, i64 0, i32 1\n"
        ));
        self.body
            .push_str(&format!("  store ptr {env_ssa}, ptr {cs_e}\n"));

        // Register sig so subsequent indirect calls on `clos` work.
        let fs = FnSig {
            params: llvm_param_tys,
            ret: llvm_ret,
        };
        self.ssa_fn_sigs.insert(clos.clone(), fs);

        // emit per-pair drop fns for heap-allocated
        // closures under `--alloc=rc`. `lam_local` closures live on
        // the stack and need no drop fn — LLVM `alloca` reclaims the
        // pair on fn return. For heap closures we emit two symbols:
        //
        //  - `drop_<thunk>_env(env)` — dec each pointer-typed
        //    capture, then dec the env block itself. ADT captures
        //    cascade through the per-type drop fn; closure-typed
        //    captures fall back to plain `ailang_rc_dec` because we
        //    don't keep a per-pair drop pointer alongside the env
        //    cell.
        //
        //  - `drop_<thunk>_pair(p)` — load the env from offset 8,
        //    call `drop_<thunk>_env(env)`, then dec the pair box.
        //
        // The pair drop is the symbol `Term::Let` lowering calls
        // when the binder is a closure; we record it in
        // `closure_drops` keyed by the closure-pair SSA.
        if matches!(self.alloc, AllocStrategy::Rc) && !lam_local {
            let pair_drop = format!("drop_{m}_{thunk_name}_pair", m = self.module_name);
            let env_drop = format!("drop_{m}_{thunk_name}_env", m = self.module_name);
            let env_text = self.build_env_drop_fn(&env_drop, &cap_meta);
            self.deferred_thunks.push(env_text);
            let pair_text = self.build_pair_drop_fn(&pair_drop, &env_drop, !cap_meta.is_empty());
            self.deferred_thunks.push(pair_text);
            self.closure_drops.insert(clos.clone(), pair_drop);
        }

        Ok((clos, "ptr".into()))
    }

    /// walk a Term collecting free-variable names that should
    /// be captured by an enclosing lambda. Builtins and top-level fns
    /// are excluded — they're globally accessible. Qualified names
    /// (`prefix.def`) are excluded too.
    fn collect_captures(
        t: &Term,
        bound: &mut BTreeSet<String>,
        captures: &mut Vec<String>,
        captures_set: &mut BTreeSet<String>,
        builtins: &BTreeSet<&str>,
        top_level: &BTreeSet<String>,
    ) {
        match t {
            Term::Lit { .. } => {}
            Term::Var { name } => {
                if name.contains('.') {
                    return; // qualified ref is module-level
                }
                if bound.contains(name) {
                    return;
                }
                if builtins.contains(name.as_str()) {
                    return;
                }
                if top_level.contains(name) {
                    return;
                }
                if captures_set.insert(name.clone()) {
                    captures.push(name.clone());
                }
            }
            Term::App { callee, args, .. } => {
                Self::collect_captures(callee, bound, captures, captures_set, builtins, top_level);
                for a in args {
                    Self::collect_captures(a, bound, captures, captures_set, builtins, top_level);
                }
            }
            Term::Let { name, value, body } => {
                Self::collect_captures(value, bound, captures, captures_set, builtins, top_level);
                let inserted = bound.insert(name.clone());
                Self::collect_captures(body, bound, captures, captures_set, builtins, top_level);
                if inserted {
                    bound.remove(name);
                }
            }
            Term::If { cond, then, else_ } => {
                Self::collect_captures(cond, bound, captures, captures_set, builtins, top_level);
                Self::collect_captures(then, bound, captures, captures_set, builtins, top_level);
                Self::collect_captures(else_, bound, captures, captures_set, builtins, top_level);
            }
            Term::Do { args, .. } => {
                for a in args {
                    Self::collect_captures(a, bound, captures, captures_set, builtins, top_level);
                }
            }
            Term::Ctor { args, .. } => {
                for a in args {
                    Self::collect_captures(a, bound, captures, captures_set, builtins, top_level);
                }
            }
            Term::Match { scrutinee, arms } => {
                Self::collect_captures(scrutinee, bound, captures, captures_set, builtins, top_level);
                for arm in arms {
                    let mut pat_bindings = Vec::new();
                    Self::pattern_bound_names(&arm.pat, &mut pat_bindings);
                    let mut newly_bound = Vec::new();
                    for n in &pat_bindings {
                        if bound.insert(n.clone()) {
                            newly_bound.push(n.clone());
                        }
                    }
                    Self::collect_captures(&arm.body, bound, captures, captures_set, builtins, top_level);
                    for n in newly_bound {
                        bound.remove(&n);
                    }
                }
            }
            Term::Lam { params, body, .. } => {
                // Inner lambda's params don't escape; its free vars
                // (relative to the outer) DO contribute to outer captures.
                let mut newly_bound = Vec::new();
                for p in params {
                    if bound.insert(p.clone()) {
                        newly_bound.push(p.clone());
                    }
                }
                Self::collect_captures(body, bound, captures, captures_set, builtins, top_level);
                for n in newly_bound {
                    bound.remove(&n);
                }
            }
            Term::Seq { lhs, rhs } => {
                Self::collect_captures(lhs, bound, captures, captures_set, builtins, top_level);
                Self::collect_captures(rhs, bound, captures, captures_set, builtins, top_level);
            }
            Term::LetRec { .. } => {
                // eliminated by desugar before codegen.
                unreachable!("Term::LetRec eliminated by desugar")
            }
            Term::Clone { value } => {
                // clone is identity for capture analysis —
                // free vars of `(clone X)` are exactly the free vars of `X`.
                Self::collect_captures(value, bound, captures, captures_set, builtins, top_level);
            }
            Term::ReuseAs { source, body } => {
                // free vars are the union of source and body.
                Self::collect_captures(source, bound, captures, captures_set, builtins, top_level);
                Self::collect_captures(body, bound, captures, captures_set, builtins, top_level);
            }
            Term::Loop { binders, body } => {
                let mut newly_bound: Vec<String> = Vec::new();
                for b in binders {
                    Self::collect_captures(&b.init, bound, captures, captures_set, builtins, top_level);
                    if bound.insert(b.name.clone()) {
                        newly_bound.push(b.name.clone());
                    }
                }
                Self::collect_captures(body, bound, captures, captures_set, builtins, top_level);
                for n in newly_bound {
                    bound.remove(&n);
                }
            }
            Term::Recur { args } => {
                for a in args {
                    Self::collect_captures(a, bound, captures, captures_set, builtins, top_level);
                }
            }
        }
    }

    fn pattern_bound_names(p: &Pattern, out: &mut Vec<String>) {
        match p {
            Pattern::Wild | Pattern::Lit { .. } => {}
            Pattern::Var { name } => out.push(name.clone()),
            Pattern::Ctor { fields, .. } => {
                for f in fields {
                    Self::pattern_bound_names(f, out);
                }
            }
        }
    }
}