//! Phase 200-A: Function scope capture infrastructure
//!
//! This module provides types for capturing function-scoped variables
//! that are effectively immutable within a loop context.
//!
//! # Example
//!
//! For a function like JsonParser._atoi():
//!
//! ```nyash
//! method _atoi(s, pos, len) {
//!     local digits = "0123456789"  // <-- Captured variable
//!     local value = 0
//!     loop(pos < len) {
//!         local ch = s.charAt(pos)
//!         local digit = digits.indexOf(ch)  // Uses captured 'digits'
//!         if (digit < 0) { break }
//!         value = value * 10 + digit
//!         pos = pos + 1
//!     }
//!     return value
//! }
//! ```
//!
//! Here, `digits` is:
//! - Declared in function scope (before the loop)
//! - Never reassigned (effectively immutable)
//! - Referenced in loop body (digits.indexOf(ch))
//!
//! Phase 200-A creates the infrastructure to capture such variables.
//! Phase 200-B will implement the actual detection logic.

use crate::mir::ValueId;
use crate::mir::join_ir::lowering::loop_scope_shape::LoopScopeShape;
use crate::ast::ASTNode;

/// A variable captured from function scope for use in loop conditions/body.
///
/// Example: `local digits = "0123456789"` in JsonParser._atoi()
///
/// # Invariants
///
/// - `name`: Variable name as it appears in the source code
/// - `host_id`: MIR ValueId of the original definition in the host function
/// - `is_immutable`: True if the variable is never reassigned in the function
#[derive(Debug, Clone)]
pub struct CapturedVar {
    /// Variable name (e.g., "digits", "table")
    pub name: String,

    /// MIR ValueId of the original definition in the host function
    pub host_id: ValueId,

    /// Whether this variable is never reassigned in the function
    ///
    /// Phase 200-B will implement assignment analysis to determine this.
    /// For now, this is always set to true as a conservative default.
    pub is_immutable: bool,
}

/// Environment containing function-scoped captured variables.
///
/// Phase 200-A: Type definition only, not yet integrated with ConditionEnv.
/// Phase 200-B: Will be populated by FunctionScopeCaptureAnalyzer and
///               integrated into ConditionEnv via ConditionEnvBuilder v2.
#[derive(Debug, Clone, Default)]
pub struct CapturedEnv {
    /// List of captured variables
    pub vars: Vec<CapturedVar>,
}

impl CapturedEnv {
    /// Create a new empty environment
    pub fn new() -> Self {
        Self { vars: Vec::new() }
    }

    /// Check if the environment is empty
    pub fn is_empty(&self) -> bool {
        self.vars.is_empty()
    }

    /// Add a captured variable to the environment
    pub fn add_var(&mut self, var: CapturedVar) {
        self.vars.push(var);
    }

    /// Look up a captured variable by name
    ///
    /// Returns `Some(&CapturedVar)` if found, `None` otherwise.
    pub fn get(&self, name: &str) -> Option<&CapturedVar> {
        self.vars.iter().find(|v| v.name == name)
    }
}

/// Analyzes function-scoped variables that can be safely captured for loop conditions/body.
///
/// # Phase 200-B Implementation
///
/// Detects function-scoped variables that are effectively immutable constants
/// within a loop context (e.g., `digits` in JsonParser._atoi()).
///
/// # Detection Criteria
///
/// A variable is captured if ALL of the following conditions are met:
///
/// 1. **Declared before the loop**: Variable must be declared in function scope before the loop
/// 2. **Safe constant init**: Initialized with string/integer literal only
/// 3. **Never reassigned**: Variable is never reassigned within the function (is_immutable = true)
/// 4. **Referenced in loop**: Variable is referenced in loop condition or body
/// 5. **Not a loop parameter**: Variable is not in scope.loop_params
/// 6. **Not a body-local**: Variable is not in scope.body_locals
///
/// # Example
///
/// ```nyash
/// method _atoi(s, pos, len) {
///     local digits = "0123456789"  // ✅ Captured (declared before loop, never reassigned)
///     local value = 0               // ❌ Not captured (reassigned in loop body)
///     loop(pos < len) {
///         local ch = s.charAt(pos)  // ❌ Not captured (body-local)
///         local digit = digits.indexOf(ch)
///         value = value * 10 + digit
///         pos = pos + 1
///     }
/// }
/// ```
///
/// # Arguments
///
/// * `fn_body` - AST nodes of the function body (for analysis)
/// * `loop_ast` - AST node of the loop statement
/// * `scope` - LoopScopeShape (for excluding loop params and body-locals)
///
/// # Returns
///
/// `CapturedEnv` containing all captured variables
pub fn analyze_captured_vars(
    fn_body: &[ASTNode],
    loop_ast: &ASTNode,
    scope: &LoopScopeShape,
) -> CapturedEnv {
    use std::env;

    let debug = env::var("NYASH_CAPTURE_DEBUG").is_ok();

    if debug {
        eprintln!("[capture/debug] Starting capture analysis");
    }

    // Step 1: Find loop position in fn_body
    let loop_index = match find_stmt_index(fn_body, loop_ast) {
        Some(idx) => idx,
        None => {
            if debug {
                eprintln!("[capture/debug] Loop not found in function body, returning empty CapturedEnv");
            }
            return CapturedEnv::new();
        }
    };

    if debug {
        eprintln!("[capture/debug] Loop found at index {}", loop_index);
    }

    // Step 2: Collect local declarations BEFORE the loop
    let pre_loop_locals = collect_local_declarations(&fn_body[..loop_index]);

    if debug {
        eprintln!("[capture/debug] Found {} pre-loop local declarations", pre_loop_locals.len());
    }

    let mut env = CapturedEnv::new();

    // Step 3: For each pre-loop local, check capture criteria
    for (name, init_expr) in pre_loop_locals {
        if debug {
            eprintln!("[capture/check] Checking variable '{}'", name);
        }

        // 3a: Is init expression a safe constant?
        if !is_safe_const_init(&init_expr) {
            if debug {
                eprintln!("[capture/reject] '{}': init is not a safe constant", name);
            }
            continue;
        }

        // 3b: Is this variable reassigned anywhere in fn_body?
        if is_reassigned_in_fn(fn_body, &name) {
            if debug {
                eprintln!("[capture/reject] '{}': reassigned in function", name);
            }
            continue;
        }

        // 3c: Is this variable used in loop (condition or body)?
        if !is_used_in_loop(loop_ast, &name) {
            if debug {
                eprintln!("[capture/reject] '{}': not used in loop", name);
            }
            continue;
        }

        // 3d: Skip if already in pinned, carriers, or body_locals
        if scope.pinned.contains(&name) {
            if debug {
                eprintln!("[capture/reject] '{}': is a pinned variable", name);
            }
            continue;
        }

        if scope.carriers.contains(&name) {
            if debug {
                eprintln!("[capture/reject] '{}': is a carrier variable", name);
            }
            continue;
        }

        if scope.body_locals.contains(&name) {
            if debug {
                eprintln!("[capture/reject] '{}': is a body-local variable", name);
            }
            continue;
        }

        // All checks passed: add to CapturedEnv
        // Note: We don't have access to variable_map here, so we use a placeholder ValueId
        // The actual host_id will be resolved in ConditionEnvBuilder
        if debug {
            eprintln!("[capture/accept] '{}': ALL CHECKS PASSED, adding to CapturedEnv", name);
        }

        env.add_var(CapturedVar {
            name: name.clone(),
            host_id: ValueId(0), // Placeholder, will be resolved in ConditionEnvBuilder
            is_immutable: true,
        });
    }

    if debug {
        eprintln!("[capture/result] Captured {} variables: {:?}",
            env.vars.len(),
            env.vars.iter().map(|v| &v.name).collect::<Vec<_>>()
        );
    }

    env
}

/// Phase 200-C: Analyze captured vars with condition/body instead of loop_ast
///
/// This variant solves the pointer comparison problem when the loop AST is constructed
/// dynamically (e.g., in Pattern 2). Instead of passing a loop_ast reference,
/// we pass the condition and body directly and perform structural matching.
///
/// # Arguments
///
/// * `fn_body` - AST nodes of the function body (for analysis)
/// * `loop_condition` - Condition expression of the loop
/// * `loop_body` - Body statements of the loop
/// * `scope` - LoopScopeShape (for excluding loop params and body-locals)
///
/// # Returns
///
/// `CapturedEnv` containing all captured variables
pub fn analyze_captured_vars_v2(
    fn_body: &[ASTNode],
    loop_condition: &ASTNode,
    loop_body: &[ASTNode],
    scope: &LoopScopeShape,
) -> CapturedEnv {
    use std::env;

    let debug = env::var("NYASH_CAPTURE_DEBUG").is_ok();

    if debug {
        eprintln!("[capture/debug] Starting capture analysis v2 (structural matching)");
    }

    // Step 1: Find loop position in fn_body by structural matching
    let loop_index = match find_loop_index_by_structure(fn_body, loop_condition, loop_body) {
        Some(idx) => idx,
        None => {
            if debug {
                eprintln!("[capture/debug] Loop not found in function body by structure, returning empty CapturedEnv");
            }
            return CapturedEnv::new();
        }
    };

    if debug {
        eprintln!("[capture/debug] Loop found at index {} by structure", loop_index);
    }

    // Step 2: Collect local declarations BEFORE the loop
    let pre_loop_locals = collect_local_declarations(&fn_body[..loop_index]);

    if debug {
        eprintln!("[capture/debug] Found {} pre-loop local declarations", pre_loop_locals.len());
    }

    let mut env = CapturedEnv::new();

    // Step 3: For each pre-loop local, check capture criteria
    for (name, init_expr) in pre_loop_locals {
        if debug {
            eprintln!("[capture/check] Checking variable '{}'", name);
        }

        // 3a: Is init expression a safe constant?
        if !is_safe_const_init(&init_expr) {
            if debug {
                eprintln!("[capture/reject] '{}': init is not a safe constant", name);
            }
            continue;
        }

        // 3b: Is this variable reassigned anywhere in fn_body?
        if is_reassigned_in_fn(fn_body, &name) {
            if debug {
                eprintln!("[capture/reject] '{}': reassigned in function", name);
            }
            continue;
        }

        // 3c: Is this variable used in loop (condition or body)?
        if !is_used_in_loop_parts(loop_condition, loop_body, &name) {
            if debug {
                eprintln!("[capture/reject] '{}': not used in loop", name);
            }
            continue;
        }

        // 3d: Skip if already in pinned, carriers, or body_locals
        if scope.pinned.contains(&name) {
            if debug {
                eprintln!("[capture/reject] '{}': is a pinned variable", name);
            }
            continue;
        }

        if scope.carriers.contains(&name) {
            if debug {
                eprintln!("[capture/reject] '{}': is a carrier variable", name);
            }
            continue;
        }

        if scope.body_locals.contains(&name) {
            if debug {
                eprintln!("[capture/reject] '{}': is a body-local variable", name);
            }
            continue;
        }

        // All checks passed: add to CapturedEnv
        if debug {
            eprintln!("[capture/accept] '{}': ALL CHECKS PASSED, adding to CapturedEnv", name);
        }

        env.add_var(CapturedVar {
            name: name.clone(),
            host_id: ValueId(0), // Placeholder, will be resolved in ConditionEnvBuilder
            is_immutable: true,
        });
    }

    if debug {
        eprintln!("[capture/result] Captured {} variables: {:?}",
            env.vars.len(),
            env.vars.iter().map(|v| &v.name).collect::<Vec<_>>()
        );
    }

    env
}

/// Find the index of a loop statement in the function body
///
/// Returns Some(index) if found, None otherwise.
fn find_stmt_index(fn_body: &[ASTNode], loop_ast: &ASTNode) -> Option<usize> {
    // Compare by pointer address (same AST node instance)
    fn_body.iter().position(|stmt| {
        std::ptr::eq(stmt as *const ASTNode, loop_ast as *const ASTNode)
    })
}

/// Phase 200-C: Find loop index by structure matching (condition + body comparison)
///
/// Instead of pointer comparison, compare the loop structure.
/// This is useful when the loop AST is constructed dynamically.
fn find_loop_index_by_structure(
    fn_body: &[ASTNode],
    target_condition: &ASTNode,
    target_body: &[ASTNode],
) -> Option<usize> {
    for (idx, stmt) in fn_body.iter().enumerate() {
        if let ASTNode::Loop { condition, body, .. } = stmt {
            // Compare condition and body by structure
            if ast_matches(condition, target_condition) && body_matches(body, target_body) {
                return Some(idx);
            }
        }
    }
    None
}

/// Simple structural AST comparison
///
/// Uses Debug string comparison as a heuristic. This is not perfect but
/// works well enough for finding loops by structure.
fn ast_matches(a: &ASTNode, b: &ASTNode) -> bool {
    format!("{:?}", a) == format!("{:?}", b)
}

/// Compare two body slices by structure
fn body_matches(a: &[ASTNode], b: &[ASTNode]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    a.iter().zip(b.iter()).all(|(x, y)| ast_matches(x, y))
}

/// Collect local variable declarations from statements
///
/// Returns Vec<(name, init_expr)> for each variable declared with `local`.
fn collect_local_declarations(stmts: &[ASTNode]) -> Vec<(String, Option<Box<ASTNode>>)> {
    let mut locals = Vec::new();

    for stmt in stmts {
        if let ASTNode::Local { variables, initial_values, .. } = stmt {
            // Local declaration can have multiple variables (e.g., local a, b, c)
            for (i, name) in variables.iter().enumerate() {
                let init_expr = initial_values.get(i).and_then(|opt| opt.clone());
                locals.push((name.clone(), init_expr));
            }
        }
    }

    locals
}

/// Check if expression is a safe constant (string/integer literal)
///
/// Phase 200-B: Only string and integer literals are allowed.
/// Future: May expand to include other safe constant patterns.
fn is_safe_const_init(expr: &Option<Box<ASTNode>>) -> bool {
    match expr {
        Some(boxed) => match boxed.as_ref() {
            ASTNode::Literal { value, .. } => matches!(
                value,
                crate::ast::LiteralValue::String(_) | crate::ast::LiteralValue::Integer(_)
            ),
            _ => false,
        },
        None => false,
    }
}

/// Check if variable is reassigned anywhere in function body
///
/// Walks the entire function body AST to detect any assignments to the variable.
/// Returns true if the variable is reassigned (excluding the initial local declaration).
fn is_reassigned_in_fn(fn_body: &[ASTNode], name: &str) -> bool {
    fn check_node(node: &ASTNode, name: &str) -> bool {
        match node {
            // Assignment to this variable
            ASTNode::Assignment { target, value, .. } => {
                // Check if target is the variable we're looking for
                let is_target_match = match target.as_ref() {
                    ASTNode::Variable { name: var_name, .. } => var_name == name,
                    ASTNode::FieldAccess { .. } | ASTNode::Index { .. } => {
                        // Field access or index assignment doesn't count as reassignment
                        false
                    }
                    _ => false,
                };

                is_target_match || check_node(value, name)
            }

            // Grouped assignment expression: (x = expr)
            ASTNode::GroupedAssignmentExpr { lhs, rhs, .. } => {
                lhs == name || check_node(rhs, name)
            }

            // Recursive cases
            ASTNode::If { condition, then_body, else_body, .. } => {
                check_node(condition, name)
                    || then_body.iter().any(|n| check_node(n, name))
                    || else_body.as_ref().map_or(false, |body| body.iter().any(|n| check_node(n, name)))
            }

            ASTNode::Loop { condition, body, .. } => {
                check_node(condition, name) || body.iter().any(|n| check_node(n, name))
            }

            ASTNode::While { condition, body, .. } => {
                check_node(condition, name) || body.iter().any(|n| check_node(n, name))
            }

            ASTNode::TryCatch { try_body, catch_clauses, finally_body, .. } => {
                try_body.iter().any(|n| check_node(n, name))
                    || catch_clauses.iter().any(|clause| clause.body.iter().any(|n| check_node(n, name)))
                    || finally_body.as_ref().map_or(false, |body| body.iter().any(|n| check_node(n, name)))
            }

            ASTNode::UnaryOp { operand, .. } => check_node(operand, name),

            ASTNode::BinaryOp { left, right, .. } => {
                check_node(left, name) || check_node(right, name)
            }

            ASTNode::MethodCall { object, arguments, .. } => {
                check_node(object, name) || arguments.iter().any(|arg| check_node(arg, name))
            }

            ASTNode::FunctionCall { arguments, .. } => {
                arguments.iter().any(|arg| check_node(arg, name))
            }

            ASTNode::FieldAccess { object, .. } => check_node(object, name),

            ASTNode::Index { target, index, .. } => {
                check_node(target, name) || check_node(index, name)
            }

            ASTNode::Return { value, .. } => {
                value.as_ref().map_or(false, |v| check_node(v, name))
            }

            ASTNode::Local { .. } => {
                // Local declarations are not reassignments
                false
            }

            _ => false,
        }
    }

    fn_body.iter().any(|stmt| check_node(stmt, name))
}

/// Check if variable is referenced in loop condition or body
///
/// Returns true if the variable name appears anywhere in the loop AST.
fn is_used_in_loop(loop_ast: &ASTNode, name: &str) -> bool {
    fn check_usage(node: &ASTNode, name: &str) -> bool {
        match node {
            ASTNode::Variable { name: var_name, .. } => var_name == name,

            ASTNode::Loop { condition, body, .. } => {
                check_usage(condition, name) || body.iter().any(|n| check_usage(n, name))
            }

            ASTNode::If { condition, then_body, else_body, .. } => {
                check_usage(condition, name)
                    || then_body.iter().any(|n| check_usage(n, name))
                    || else_body.as_ref().map_or(false, |body| body.iter().any(|n| check_usage(n, name)))
            }

            ASTNode::Assignment { target, value, .. } => {
                check_usage(target, name) || check_usage(value, name)
            }

            ASTNode::UnaryOp { operand, .. } => check_usage(operand, name),

            ASTNode::BinaryOp { left, right, .. } => {
                check_usage(left, name) || check_usage(right, name)
            }

            ASTNode::MethodCall { object, arguments, .. } => {
                check_usage(object, name) || arguments.iter().any(|arg| check_usage(arg, name))
            }

            ASTNode::FunctionCall { arguments, .. } => {
                arguments.iter().any(|arg| check_usage(arg, name))
            }

            ASTNode::FieldAccess { object, .. } => check_usage(object, name),

            ASTNode::Index { target, index, .. } => {
                check_usage(target, name) || check_usage(index, name)
            }

            ASTNode::Return { value, .. } => {
                value.as_ref().map_or(false, |v| check_usage(v, name))
            }

            ASTNode::Local { initial_values, .. } => {
                initial_values.iter().any(|opt| {
                    opt.as_ref().map_or(false, |init| check_usage(init, name))
                })
            }

            _ => false,
        }
    }

    check_usage(loop_ast, name)
}

/// Phase 200-C: Check if variable is used in loop condition or body (separate parts)
///
/// This is used by analyze_captured_vars_v2 when condition and body are passed separately.
fn is_used_in_loop_parts(condition: &ASTNode, body: &[ASTNode], name: &str) -> bool {
    fn check_usage(node: &ASTNode, name: &str) -> bool {
        match node {
            ASTNode::Variable { name: var_name, .. } => var_name == name,

            ASTNode::Loop { condition, body, .. } => {
                check_usage(condition, name) || body.iter().any(|n| check_usage(n, name))
            }

            ASTNode::If { condition, then_body, else_body, .. } => {
                check_usage(condition, name)
                    || then_body.iter().any(|n| check_usage(n, name))
                    || else_body.as_ref().map_or(false, |body| body.iter().any(|n| check_usage(n, name)))
            }

            ASTNode::Assignment { target, value, .. } => {
                check_usage(target, name) || check_usage(value, name)
            }

            ASTNode::UnaryOp { operand, .. } => check_usage(operand, name),

            ASTNode::BinaryOp { left, right, .. } => {
                check_usage(left, name) || check_usage(right, name)
            }

            ASTNode::MethodCall { object, arguments, .. } => {
                check_usage(object, name) || arguments.iter().any(|arg| check_usage(arg, name))
            }

            ASTNode::FunctionCall { arguments, .. } => {
                arguments.iter().any(|arg| check_usage(arg, name))
            }

            ASTNode::FieldAccess { object, .. } => check_usage(object, name),

            ASTNode::Index { target, index, .. } => {
                check_usage(target, name) || check_usage(index, name)
            }

            ASTNode::Return { value, .. } => {
                value.as_ref().map_or(false, |v| check_usage(v, name))
            }

            ASTNode::Local { initial_values, .. } => {
                initial_values.iter().any(|opt| {
                    opt.as_ref().map_or(false, |init| check_usage(init, name))
                })
            }

            _ => false,
        }
    }

    check_usage(condition, name) || body.iter().any(|n| check_usage(n, name))
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_captured_env_empty() {
        let env = CapturedEnv::new();
        assert!(env.is_empty());
        assert!(env.get("digits").is_none());
    }

    #[test]
    fn test_captured_env_add_and_get() {
        let mut env = CapturedEnv::new();
        env.add_var(CapturedVar {
            name: "digits".to_string(),
            host_id: ValueId(42),
            is_immutable: true,
        });

        assert!(!env.is_empty());
        let var = env.get("digits").unwrap();
        assert_eq!(var.name, "digits");
        assert_eq!(var.host_id, ValueId(42));
        assert!(var.is_immutable);
    }

    #[test]
    fn test_captured_env_multiple_vars() {
        let mut env = CapturedEnv::new();
        env.add_var(CapturedVar {
            name: "digits".to_string(),
            host_id: ValueId(42),
            is_immutable: true,
        });
        env.add_var(CapturedVar {
            name: "table".to_string(),
            host_id: ValueId(100),
            is_immutable: true,
        });

        assert_eq!(env.vars.len(), 2);
        assert!(env.get("digits").is_some());
        assert!(env.get("table").is_some());
        assert!(env.get("nonexistent").is_none());
    }

    // Phase 200-B: Capture analysis tests

    #[test]
    fn test_capture_simple_digits() {
        use crate::ast::{ASTNode, LiteralValue, Span};

        // Build AST for:
        // local digits = "0123456789"
        // loop(i < 10) {
        //     local pos = digits.indexOf(ch)
        // }

        let digits_decl = ASTNode::Local {
            variables: vec!["digits".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::Literal {
                value: LiteralValue::String("0123456789".to_string()),
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        };

        let loop_body = vec![ASTNode::Local {
            variables: vec!["pos".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::MethodCall {
                object: Box::new(ASTNode::Variable {
                    name: "digits".to_string(),
                    span: Span::unknown(),
                }),
                method: "indexOf".to_string(),
                arguments: vec![ASTNode::Variable {
                    name: "ch".to_string(),
                    span: Span::unknown(),
                }],
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        }];

        let loop_node = ASTNode::Loop {
            condition: Box::new(ASTNode::BinaryOp {
                operator: crate::ast::BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "i".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(10),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            body: loop_body,
            span: Span::unknown(),
        };

        let fn_body = vec![digits_decl, loop_node.clone()];

        use std::collections::{BTreeSet, BTreeMap};
        use crate::mir::BasicBlockId;

        let scope = crate::mir::join_ir::lowering::loop_scope_shape::LoopScopeShape {
            header: BasicBlockId(0),
            body: BasicBlockId(1),
            latch: BasicBlockId(2),
            exit: BasicBlockId(3),
            pinned: BTreeSet::from(["i".to_string()]),
            carriers: BTreeSet::new(),
            body_locals: BTreeSet::from(["pos".to_string()]),
            exit_live: BTreeSet::new(),
            progress_carrier: None,
            variable_definitions: BTreeMap::new(),
        };

        // IMPORTANT: Pass a reference to the same loop_node instance that's in fn_body
        // find_stmt_index uses pointer comparison, so we must use &fn_body[1] instead of &loop_node
        let env = analyze_captured_vars(&fn_body, &fn_body[1], &scope);

        assert_eq!(env.vars.len(), 1);
        assert!(env.get("digits").is_some());
        let var = env.get("digits").unwrap();
        assert_eq!(var.name, "digits");
        assert!(var.is_immutable);
    }

    #[test]
    fn test_capture_reassigned_rejected() {
        use crate::ast::{ASTNode, LiteralValue, Span};

        // Build AST for:
        // local digits = "0123456789"
        // digits = "abc"  // reassignment
        // loop(i < 10) {
        //     local pos = digits.indexOf(ch)
        // }

        let digits_decl = ASTNode::Local {
            variables: vec!["digits".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::Literal {
                value: LiteralValue::String("0123456789".to_string()),
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        };

        let reassignment = ASTNode::Assignment {
            target: Box::new(ASTNode::Variable {
                name: "digits".to_string(),
                span: Span::unknown(),
            }),
            value: Box::new(ASTNode::Literal {
                value: LiteralValue::String("abc".to_string()),
                span: Span::unknown(),
            }),
            span: Span::unknown(),
        };

        let loop_body = vec![ASTNode::Local {
            variables: vec!["pos".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::MethodCall {
                object: Box::new(ASTNode::Variable {
                    name: "digits".to_string(),
                    span: Span::unknown(),
                }),
                method: "indexOf".to_string(),
                arguments: vec![],
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        }];

        let loop_node = ASTNode::Loop {
            condition: Box::new(ASTNode::BinaryOp {
                operator: crate::ast::BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "i".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(10),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            body: loop_body,
            span: Span::unknown(),
        };

        let fn_body = vec![digits_decl, reassignment, loop_node.clone()];

        use std::collections::{BTreeSet, BTreeMap};
        use crate::mir::BasicBlockId;

        let scope = crate::mir::join_ir::lowering::loop_scope_shape::LoopScopeShape {
            header: BasicBlockId(0),
            body: BasicBlockId(1),
            latch: BasicBlockId(2),
            exit: BasicBlockId(3),
            pinned: BTreeSet::from(["i".to_string()]),
            carriers: BTreeSet::new(),
            body_locals: BTreeSet::from(["pos".to_string()]),
            exit_live: BTreeSet::new(),
            progress_carrier: None,
            variable_definitions: BTreeMap::new(),
        };

        let env = analyze_captured_vars(&fn_body, &loop_node, &scope);

        // Should reject because digits is reassigned
        assert_eq!(env.vars.len(), 0);
    }

    #[test]
    fn test_capture_after_loop_rejected() {
        use crate::ast::{ASTNode, LiteralValue, Span};

        // Build AST for:
        // loop(i < 10) { }
        // local digits = "0123456789"  // defined AFTER loop

        let loop_node = ASTNode::Loop {
            condition: Box::new(ASTNode::BinaryOp {
                operator: crate::ast::BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "i".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(10),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            body: vec![],
            span: Span::unknown(),
        };

        let digits_decl = ASTNode::Local {
            variables: vec!["digits".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::Literal {
                value: LiteralValue::String("0123456789".to_string()),
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        };

        let fn_body = vec![loop_node.clone(), digits_decl];

        use std::collections::{BTreeSet, BTreeMap};
        use crate::mir::BasicBlockId;

        let scope = crate::mir::join_ir::lowering::loop_scope_shape::LoopScopeShape {
            header: BasicBlockId(0),
            body: BasicBlockId(1),
            latch: BasicBlockId(2),
            exit: BasicBlockId(3),
            pinned: BTreeSet::from(["i".to_string()]),
            carriers: BTreeSet::new(),
            body_locals: BTreeSet::new(),
            exit_live: BTreeSet::new(),
            progress_carrier: None,
            variable_definitions: BTreeMap::new(),
        };

        let env = analyze_captured_vars(&fn_body, &loop_node, &scope);

        // Should reject because digits is defined after the loop
        assert_eq!(env.vars.len(), 0);
    }

    #[test]
    fn test_capture_method_call_init_rejected() {
        use crate::ast::{ASTNode, LiteralValue, Span};

        // Build AST for:
        // local result = someBox.getValue()  // MethodCall init
        // loop(i < 10) {
        //     local x = result.something()
        // }

        let result_decl = ASTNode::Local {
            variables: vec!["result".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::MethodCall {
                object: Box::new(ASTNode::Variable {
                    name: "someBox".to_string(),
                    span: Span::unknown(),
                }),
                method: "getValue".to_string(),
                arguments: vec![],
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        };

        let loop_body = vec![ASTNode::Local {
            variables: vec!["x".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::MethodCall {
                object: Box::new(ASTNode::Variable {
                    name: "result".to_string(),
                    span: Span::unknown(),
                }),
                method: "something".to_string(),
                arguments: vec![],
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        }];

        let loop_node = ASTNode::Loop {
            condition: Box::new(ASTNode::BinaryOp {
                operator: crate::ast::BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "i".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(10),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            body: loop_body,
            span: Span::unknown(),
        };

        let fn_body = vec![result_decl, loop_node.clone()];

        use std::collections::{BTreeSet, BTreeMap};
        use crate::mir::BasicBlockId;

        let scope = crate::mir::join_ir::lowering::loop_scope_shape::LoopScopeShape {
            header: BasicBlockId(0),
            body: BasicBlockId(1),
            latch: BasicBlockId(2),
            exit: BasicBlockId(3),
            pinned: BTreeSet::from(["i".to_string()]),
            carriers: BTreeSet::new(),
            body_locals: BTreeSet::from(["x".to_string()]),
            exit_live: BTreeSet::new(),
            progress_carrier: None,
            variable_definitions: BTreeMap::new(),
        };

        let env = analyze_captured_vars(&fn_body, &loop_node, &scope);

        // Should reject because result is initialized with MethodCall (not safe constant)
        assert_eq!(env.vars.len(), 0);
    }

    #[test]
    fn test_capture_unused_in_loop_rejected() {
        use crate::ast::{ASTNode, LiteralValue, Span};

        // Build AST for:
        // local digits = "0123456789"  // not used in loop
        // loop(i < 10) {
        //     print(i)  // doesn't use digits
        // }

        let digits_decl = ASTNode::Local {
            variables: vec!["digits".to_string()],
            initial_values: vec![Some(Box::new(ASTNode::Literal {
                value: LiteralValue::String("0123456789".to_string()),
                span: Span::unknown(),
            }))],
            span: Span::unknown(),
        };

        let loop_node = ASTNode::Loop {
            condition: Box::new(ASTNode::BinaryOp {
                operator: crate::ast::BinaryOperator::Less,
                left: Box::new(ASTNode::Variable {
                    name: "i".to_string(),
                    span: Span::unknown(),
                }),
                right: Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(10),
                    span: Span::unknown(),
                }),
                span: Span::unknown(),
            }),
            body: vec![],  // empty body, no usage of digits
            span: Span::unknown(),
        };

        let fn_body = vec![digits_decl, loop_node.clone()];

        use std::collections::{BTreeSet, BTreeMap};
        use crate::mir::BasicBlockId;

        let scope = crate::mir::join_ir::lowering::loop_scope_shape::LoopScopeShape {
            header: BasicBlockId(0),
            body: BasicBlockId(1),
            latch: BasicBlockId(2),
            exit: BasicBlockId(3),
            pinned: BTreeSet::from(["i".to_string()]),
            carriers: BTreeSet::new(),
            body_locals: BTreeSet::new(),
            exit_live: BTreeSet::new(),
            progress_carrier: None,
            variable_definitions: BTreeMap::new(),
        };

        let env = analyze_captured_vars(&fn_body, &loop_node, &scope);

        // Should reject because digits is not used in loop
        assert_eq!(env.vars.len(), 0);
    }
}