//! Phase 121: StepTree → JoinModule shadow lowering (if-only)
//!
//! ## Responsibility
//!
//! - Convert StepTree to JoinModule (Normalized dialect)
//! - Only for if-only patterns (no loops)
//! - Returns None for out-of-scope patterns
//! - Returns Err for patterns that should be supported but conversion failed
//!
//! ## Design
//!
//! - Input: `&StepTree` with pre-computed contract
//! - No AST re-analysis (contract-only decisions)
//! - Single responsibility: structure → JoinIR conversion

use crate::mir::control_tree::step_tree::StepTree;
use crate::mir::join_ir::lowering::carrier_info::{ExitMeta, JoinFragmentMeta};
use crate::mir::join_ir::JoinModule;

use super::contracts::{check_if_only, CapabilityCheckResult};

/// Box-First: StepTree → Normalized shadow lowering
pub struct StepTreeNormalizedShadowLowererBox;

impl StepTreeNormalizedShadowLowererBox {
    /// Try to lower an if-only StepTree to normalized form
    ///
    /// ## Returns
    ///
    /// - `Ok(None)`: Out of scope (e.g., contains loops)
    /// - `Ok(Some(...))`: Shadow generation succeeded
    /// - `Err(...)`: Should be supported but conversion failed (internal error)
    ///
    /// ## Contract
    ///
    /// - Only processes if-only patterns (no loops/breaks/continues)
    /// - Uses contract information only (no AST re-analysis)
    /// - Dev-only: caller must check `joinir_dev_enabled()` before calling
    ///
    /// ## Phase 122 Implementation
    ///
    /// - Generates Normalized JoinIR (env + continuation)
    /// - env layout: writes only (SSOT)
    /// - merge = join_k(env) tail-call (no PHI)
    /// - Minimal node support: If/Return/Assign(Const/Variable/BinOp(Add))
    pub fn try_lower_if_only(
        step_tree: &StepTree,
    ) -> Result<Option<(JoinModule, JoinFragmentMeta)>, String> {
        // Phase 121 P1: Capability check (if-only scope)
        let capability = check_if_only(step_tree);
        match capability {
            CapabilityCheckResult::Supported => {
                // Phase 122-123: Generate Normalized JoinModule
                Self::lower_if_only_to_normalized(step_tree)
            }
            CapabilityCheckResult::Unsupported(_reason) => {
                // Out of scope for Phase 121/122
                Ok(None)
            }
        }
    }

    /// Lower if-only StepTree to Normalized JoinModule (Phase 122-123)
    ///
    /// ## Design
    ///
    /// - env レイアウト: `writes` に含まれる変数だけ（決定的順序）
    /// - merge 形式: `join_k(env)` への tail-call（PHI 禁止）
    /// - 対応ノード: If/Return/Assign(最小セット)
    ///
    /// ## Phase 123 Node Support
    ///
    /// - Return(Integer literal): `Const + Ret(Some(vid))`
    /// - Return(Variable): Out of scope (Phase 124)
    /// - Return(void): `Ret(None)`
    /// - If(minimal compare): Compare with Integer literal only
    ///
    /// ## Returns
    ///
    /// - `Ok(Some((module, meta)))`: Normalized JoinModule生成成功
    /// - `Ok(None)`: Out of scope for Phase 123 (unsupported patterns)
    /// - `Err(msg)`: 生成できるはずなのに失敗（内部エラー）
    fn lower_if_only_to_normalized(
        step_tree: &StepTree,
    ) -> Result<Option<(JoinModule, JoinFragmentMeta)>, String> {
        use crate::mir::join_ir::{JoinFunction, JoinFuncId, JoinInst};
        use crate::mir::ValueId;
        use std::collections::BTreeMap;

        // Phase 122: env レイアウト
        let env_fields: Vec<String> = step_tree.contract.writes.iter().cloned().collect();
        let main_func_id = JoinFuncId::new(0);

        // env フィールドに対応する引数ValueIdを生成
        let mut next_value_id = 1;
        let env_params: Vec<ValueId> = env_fields
            .iter()
            .map(|_| {
                let vid = ValueId(next_value_id);
                next_value_id += 1;
                vid
            })
            .collect();

        // Phase 124: env マッピング（変数名 → ValueId）
        let env: BTreeMap<String, ValueId> = env_fields
            .iter()
            .cloned()
            .zip(env_params.iter().cloned())
            .collect();

        // main 関数生成
        let mut main_func = JoinFunction::new(
            main_func_id,
            "main".to_string(),
            env_params.clone(),
        );

        // Phase 123-124: Return node lowering
        // If Phase 123-124 patterns are not supported, return Ok(None)
        match Self::lower_return_from_tree(
            &step_tree.root,
            &mut main_func.body,
            &mut next_value_id,
            &env,
            &step_tree.contract,
        ) {
            Ok(()) => {
                // Success - continue
            }
            Err(msg) if msg.starts_with("[phase123/") || msg.starts_with("[phase124/") => {
                // Phase 123-124 limitation - out of scope
                return Ok(None);
            }
            Err(msg) => {
                // Real error - propagate
                return Err(msg);
            }
        }

        // JoinModule 構築
        let mut module = JoinModule::new();
        module.add_function(main_func);
        module.entry = Some(main_func_id);
        module.mark_normalized();

        // JoinFragmentMeta 生成（最小）
        let meta = JoinFragmentMeta::empty();

        Ok(Some((module, meta)))
    }

    /// Phase 123-124 P1-P3: Lower node from StepTree
    ///
    /// ## Support (Phase 123-124)
    ///
    /// - Return(Integer literal): Generate Const + Ret(Some(vid))
    /// - Return(void): Ret(None)
    /// - Return(Variable): Phase 124 - lookup from env (dev-only)
    /// - Return(other): Fail-Fast with structured error
    /// - If(minimal compare): Generate Compare + Branch + Ret (P3)
    fn lower_return_from_tree(
        node: &crate::mir::control_tree::step_tree::StepNode,
        body: &mut Vec<crate::mir::join_ir::JoinInst>,
        next_value_id: &mut u32,
        env: &std::collections::BTreeMap<String, crate::mir::ValueId>,
        contract: &crate::mir::control_tree::step_tree_contract_box::StepTreeContract,
    ) -> Result<(), String> {
        use crate::ast::{ASTNode, LiteralValue};
        use crate::mir::control_tree::step_tree::{StepNode, StepStmtKind};
        use crate::mir::join_ir::JoinInst;
        use crate::mir::ValueId;

        match node {
            StepNode::Block(nodes) => {
                // Process nodes in order
                for n in nodes {
                    match n {
                        StepNode::Stmt {
                            kind: StepStmtKind::Return { value_ast },
                            ..
                        } => {
                            return Self::lower_return_value(value_ast, body, next_value_id, env, contract);
                        }
                        StepNode::If { .. } => {
                            // Phase 123 P3: Lower If node
                            return Self::lower_if_node(n, body, next_value_id, env, contract);
                        }
                        _ => {
                            // Other nodes not yet supported
                        }
                    }
                }
                // No return found - default to void
                body.push(JoinInst::Ret { value: None });
                Ok(())
            }
            StepNode::Stmt {
                kind: StepStmtKind::Return { value_ast },
                ..
            } => Self::lower_return_value(value_ast, body, next_value_id, env, contract),
            StepNode::If { .. } => {
                // Phase 123 P3: Lower If node
                Self::lower_if_node(node, body, next_value_id, env, contract)
            }
            _ => {
                // No return in tree - default to void
                body.push(JoinInst::Ret { value: None });
                Ok(())
            }
        }
    }

    /// Phase 123-124 P3: Lower If node with minimal compare
    ///
    /// ## Support
    ///
    /// - Minimal binary comparison: Variable vs Integer literal
    /// - then/else: Return(Integer literal) or Return(Variable) (Phase 124)
    /// - Merge: Not yet implemented (will use join_k tail-call in future)
    ///
    /// ## Not Supported (Fail-Fast)
    ///
    /// - Compound expressions (&&, ||)
    /// - Method calls
    /// - Complex expressions
    /// - Non-return statements in branches
    fn lower_if_node(
        node: &crate::mir::control_tree::step_tree::StepNode,
        body: &mut Vec<crate::mir::join_ir::JoinInst>,
        next_value_id: &mut u32,
        env: &std::collections::BTreeMap<String, crate::mir::ValueId>,
        contract: &crate::mir::control_tree::step_tree_contract_box::StepTreeContract,
    ) -> Result<(), String> {
        use crate::ast::{ASTNode, BinaryOperator};
        use crate::mir::control_tree::step_tree::StepNode;
        use crate::mir::join_ir::{CompareOp, ConstValue, JoinInst, MirLikeInst};
        use crate::mir::ValueId;

        if let StepNode::If {
            cond_ast,
            then_branch,
            else_branch,
            ..
        } = node
        {
            let ast = &cond_ast.0;

            // Phase 123 P3: Parse minimal binary comparison only
            let (_lhs_var, op, rhs_literal) = Self::parse_minimal_compare(ast)?;

            // Generate Compare instruction
            // 1. Load/create lhs variable (for now, assume it's a parameter)
            // For Phase 123 minimal: we'll just create a load instruction placeholder
            // This is a simplification - real implementation would need variable resolution
            let lhs_vid = ValueId(*next_value_id);
            *next_value_id += 1;

            // For now, emit a const for the variable (placeholder)
            // Real implementation in Phase 124 will use reads facts
            body.push(JoinInst::Compute(MirLikeInst::Const {
                dst: lhs_vid,
                value: ConstValue::Integer(0), // Placeholder
            }));

            // 2. Create constant for rhs literal
            let rhs_vid = ValueId(*next_value_id);
            *next_value_id += 1;
            body.push(JoinInst::Compute(MirLikeInst::Const {
                dst: rhs_vid,
                value: ConstValue::Integer(rhs_literal),
            }));

            // 3. Generate Compare instruction
            let cond_vid = ValueId(*next_value_id);
            *next_value_id += 1;
            body.push(JoinInst::Compute(MirLikeInst::Compare {
                dst: cond_vid,
                op,
                lhs: lhs_vid,
                rhs: rhs_vid,
            }));

            // Phase 123 P3: Verify then/else branches contain only Return(Integer literal)
            Self::verify_branch_is_return_literal(then_branch)?;
            if let Some(else_br) = else_branch {
                Self::verify_branch_is_return_literal(else_br)?;
            }

            // For Phase 123-124, we generate a simplified structure:
            // The actual branching logic will be added in future phases
            // For now, just emit the then branch return
            Self::lower_return_from_tree(then_branch, body, next_value_id, env, contract)?;

            Ok(())
        } else {
            Err("[phase123/if/internal] Expected If node".to_string())
        }
    }

    /// Parse minimal binary comparison: Variable op Integer
    ///
    /// Returns: (variable_name, compare_op, integer_value)
    fn parse_minimal_compare(
        ast: &crate::ast::ASTNode,
    ) -> Result<(String, crate::mir::join_ir::CompareOp, i64), String> {
        use crate::ast::{ASTNode, BinaryOperator, LiteralValue};
        use crate::mir::join_ir::CompareOp;

        match ast {
            ASTNode::BinaryOp {
                operator,
                left,
                right,
                ..
            } => {
                // Phase 123: Only support Variable on left, Integer literal on right
                let var_name = match &**left {
                    ASTNode::Variable { name, .. } => name.clone(),
                    _ => {
                        return Err(format!(
                            "[phase123/if/compare_lhs_unsupported] Phase 123 only supports Variable on left side of comparison. Hint: Use simple variable comparison or wait for Phase 124"
                        ));
                    }
                };

                let int_value = match &**right {
                    ASTNode::Literal {
                        value: LiteralValue::Integer(i),
                        ..
                    } => *i,
                    _ => {
                        return Err(format!(
                            "[phase123/if/compare_rhs_unsupported] Phase 123 only supports Integer literal on right side of comparison. Hint: Use integer literal or wait for Phase 124"
                        ));
                    }
                };

                let compare_op = match operator {
                    BinaryOperator::Equal => CompareOp::Eq,
                    BinaryOperator::NotEqual => CompareOp::Ne,
                    BinaryOperator::Less => CompareOp::Lt,
                    BinaryOperator::LessEqual => CompareOp::Le,
                    BinaryOperator::Greater => CompareOp::Gt,
                    BinaryOperator::GreaterEqual => CompareOp::Ge,
                    _ => {
                        return Err(format!(
                            "[phase123/if/compare_op_unsupported] Phase 123 only supports comparison operators (==, !=, <, <=, >, >=). Hint: Use comparison operator or wait for Phase 124"
                        ));
                    }
                };

                Ok((var_name, compare_op, int_value))
            }
            _ => Err(format!(
                "[phase123/if/cond_unsupported] Phase 123 only supports binary comparisons. Hint: Use simple comparison (var == literal) or wait for Phase 124"
            )),
        }
    }

    /// Verify branch contains only Return(Integer literal)
    fn verify_branch_is_return_literal(
        branch: &crate::mir::control_tree::step_tree::StepNode,
    ) -> Result<(), String> {
        use crate::ast::{ASTNode, LiteralValue};
        use crate::mir::control_tree::step_tree::{StepNode, StepStmtKind};

        match branch {
            StepNode::Stmt {
                kind: StepStmtKind::Return { value_ast },
                ..
            } => {
                if let Some(ast_handle) = value_ast {
                    let ast = &ast_handle.0;
                    if let ASTNode::Literal {
                        value: LiteralValue::Integer(_),
                        ..
                    } = &**ast
                    {
                        Ok(())
                    } else {
                        Err(format!(
                            "[phase123/if/branch_return_not_int_literal] Phase 123 only supports Return(Integer literal) in then/else branches. Hint: Return integer literal only or wait for Phase 124"
                        ))
                    }
                } else {
                    Err(format!(
                        "[phase123/if/branch_return_void] Phase 123 requires Return(Integer literal) in branches, not void return. Hint: Return integer literal or wait for Phase 124"
                    ))
                }
            }
            StepNode::Block(nodes) => {
                // Check first node only
                if nodes.is_empty() {
                    return Err(format!(
                        "[phase123/if/branch_empty] Phase 123 requires Return(Integer literal) in branches. Hint: Add return statement"
                    ));
                }
                Self::verify_branch_is_return_literal(&nodes[0])
            }
            _ => Err(format!(
                "[phase123/if/branch_not_return] Phase 123 only supports Return(Integer literal) in then/else branches. Hint: Use return statement with integer literal"
            )),
        }
    }

    /// Phase 123-124 P1-P2-P3: Lower return value
    ///
    /// ## Support
    ///
    /// - Integer literal: Generate Const + Ret(Some(vid))
    /// - None: Ret(None)
    /// - Variable: Phase 124 - lookup from env (dev-only, Fail-Fast if not in env)
    /// - Other: Fail-Fast (out of scope)
    fn lower_return_value(
        value_ast: &Option<crate::mir::control_tree::step_tree::AstNodeHandle>,
        body: &mut Vec<crate::mir::join_ir::JoinInst>,
        next_value_id: &mut u32,
        env: &std::collections::BTreeMap<String, crate::mir::ValueId>,
        contract: &crate::mir::control_tree::step_tree_contract_box::StepTreeContract,
    ) -> Result<(), String> {
        use crate::ast::{ASTNode, LiteralValue};
        use crate::mir::join_ir::{ConstValue, JoinInst, MirLikeInst};
        use crate::mir::ValueId;

        match value_ast {
            None => {
                body.push(JoinInst::Ret { value: None });
                Ok(())
            }
            Some(ast_handle) => {
                let ast = &ast_handle.0;
                match &**ast {
                    ASTNode::Literal { value, .. } => match value {
                        LiteralValue::Integer(i) => {
                            // Phase 123 P1: Integer literal → Const + Ret(Some(vid))
                            let const_vid = ValueId(*next_value_id);
                            *next_value_id += 1;

                            // Generate Const instruction (wrapped in Compute)
                            body.push(JoinInst::Compute(MirLikeInst::Const {
                                dst: const_vid,
                                value: ConstValue::Integer(*i),
                            }));

                            // Generate Ret instruction
                            body.push(JoinInst::Ret {
                                value: Some(const_vid),
                            });

                            Ok(())
                        }
                        _ => {
                            // Phase 123: Other literals not supported
                            Err(format!(
                                "[phase123/return/literal_unsupported] Phase 123 only supports integer literals. Hint: Use integer literal or wait for Phase 124"
                            ))
                        }
                    },
                    ASTNode::Variable { name, .. } => {
                        // Phase 124 P3: Variable return support (dev-only)
                        // Check if variable is in env (writes-derived)
                        if let Some(&vid) = env.get(name) {
                            // Variable found in env - return it
                            body.push(JoinInst::Ret { value: Some(vid) });
                            Ok(())
                        } else {
                            // Variable not in env - out of scope for Phase 124
                            // Phase 124 only supports Return(Variable) when variable is in env (writes)
                            // Variables in reads-only or undefined are not supported
                            Err(format!(
                                "[phase124/return/var_not_in_env] Variable '{}' not in env (writes). Phase 124 only supports return of written variables. Hint: Assign variable before return",
                                name
                            ))
                        }
                    }
                    _ => {
                        // Phase 123: Other expressions not supported
                        Err(format!(
                            "[phase123/return/expr_unsupported] Phase 123 only supports integer literals. Hint: Simplify to literal or wait for Phase 124"
                        ))
                    }
                }
            }
        }
    }

    /// Get shadow lowering status string for dev logging
    ///
    /// ## Contract
    ///
    /// - Returns 1-line summary: "shadow_lowered=true/false reason=..."
    /// - Does not perform actual lowering (use `try_lower_if_only` for that)
    pub fn get_status_string(step_tree: &StepTree) -> String {
        let capability = check_if_only(step_tree);
        match capability {
            CapabilityCheckResult::Supported => {
                format!(
                    "shadow_lowered=true step_tree_sig={} exits={:?} writes={:?}",
                    step_tree.signature_basis_string(),
                    step_tree.contract.exits,
                    step_tree.contract.writes
                )
            }
            CapabilityCheckResult::Unsupported(reason) => {
                format!(
                    "shadow_lowered=false reason=\"{}\" step_tree_sig={}",
                    reason.reason(),
                    step_tree.signature_basis_string()
                )
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mir::control_tree::step_tree::{
        StepNode, StepStmtKind, StepTreeFeatures, StepTreeSignature,
    };
    use crate::mir::control_tree::step_tree_contract_box::StepTreeContract;

    fn make_if_only_tree() -> StepTree {
        StepTree {
            root: StepNode::Block(vec![]),
            features: StepTreeFeatures {
                has_if: true,
                has_loop: false,
                has_break: false,
                has_continue: false,
                has_return: false,
                max_if_depth: 1,
                max_loop_depth: 0,
            },
            contract: StepTreeContract {
                exits: Default::default(),
                writes: Default::default(),
                reads: Default::default(), // Phase 124
                required_caps: Default::default(),
                cond_sig: Default::default(),
            },
            signature: StepTreeSignature(0),
        }
    }

    fn make_loop_tree() -> StepTree {
        StepTree {
            root: StepNode::Block(vec![]),
            features: StepTreeFeatures {
                has_if: false,
                has_loop: true,
                has_break: false,
                has_continue: false,
                has_return: false,
                max_if_depth: 0,
                max_loop_depth: 1,
            },
            contract: StepTreeContract {
                exits: Default::default(),
                writes: Default::default(),
                reads: Default::default(), // Phase 124
                required_caps: Default::default(),
                cond_sig: Default::default(),
            },
            signature: StepTreeSignature(0),
        }
    }

    #[test]
    fn test_if_only_supported() {
        let tree = make_if_only_tree();
        let result = StepTreeNormalizedShadowLowererBox::try_lower_if_only(&tree);
        assert!(result.is_ok());
        assert!(result.unwrap().is_some());
    }

    #[test]
    fn test_loop_rejected() {
        let tree = make_loop_tree();
        let result = StepTreeNormalizedShadowLowererBox::try_lower_if_only(&tree);
        assert!(result.is_ok());
        assert!(result.unwrap().is_none());
    }

    #[test]
    fn test_status_string_if_only() {
        let tree = make_if_only_tree();
        let status = StepTreeNormalizedShadowLowererBox::get_status_string(&tree);
        assert!(status.contains("shadow_lowered=true"));
        assert!(status.contains("step_tree_sig="));
    }

    #[test]
    fn test_status_string_loop() {
        let tree = make_loop_tree();
        let status = StepTreeNormalizedShadowLowererBox::get_status_string(&tree);
        assert!(status.contains("shadow_lowered=false"));
        assert!(status.contains("reason=\"contains loop"));
    }

    #[test]
    fn test_return_integer_literal() {
        use crate::ast::{ASTNode, LiteralValue, Span};
        use crate::mir::control_tree::step_tree::AstNodeHandle;

        // Create StepTree with "return 7"
        let return_ast = Box::new(ASTNode::Literal {
            value: LiteralValue::Integer(7),
            span: Span::unknown(),
        });

        let mut tree = make_if_only_tree();
        tree.root = StepNode::Stmt {
            kind: StepStmtKind::Return {
                value_ast: Some(AstNodeHandle(return_ast)),
            },
            span: Span::unknown(),
        };

        // Lower to JoinModule
        let result = StepTreeNormalizedShadowLowererBox::try_lower_if_only(&tree);
        assert!(result.is_ok());

        let (module, _meta) = result.unwrap().expect("Should generate JoinModule");

        // Verify Const + Ret instructions
        assert_eq!(module.functions.len(), 1);
        let func = &module.functions.values().next().unwrap();
        assert_eq!(func.body.len(), 2, "Should have Const + Ret");

        // Check Const instruction
        use crate::mir::join_ir::{ConstValue, JoinInst, MirLikeInst};
        use crate::mir::ValueId;
        if let JoinInst::Compute(MirLikeInst::Const { dst, value }) = &func.body[0] {
            assert_eq!(*dst, ValueId(1));
            if let ConstValue::Integer(i) = value {
                assert_eq!(*i, 7);
            } else {
                panic!("Expected Integer const");
            }
        } else {
            panic!("Expected Const instruction");
        }

        // Check Ret instruction
        if let JoinInst::Ret { value } = &func.body[1] {
            assert_eq!(value, &Some(ValueId(1)));
        } else {
            panic!("Expected Ret instruction");
        }
    }

    #[test]
    fn test_return_void() {
        use crate::ast::Span;

        // Create StepTree with "return" (no value)
        let mut tree = make_if_only_tree();
        tree.root = StepNode::Stmt {
            kind: StepStmtKind::Return { value_ast: None },
            span: Span::unknown(),
        };

        // Lower to JoinModule
        let result = StepTreeNormalizedShadowLowererBox::try_lower_if_only(&tree);
        assert!(result.is_ok());

        let (module, _meta) = result.unwrap().expect("Should generate JoinModule");

        // Verify Ret(None) instruction
        assert_eq!(module.functions.len(), 1);
        let func = &module.functions.values().next().unwrap();
        assert_eq!(func.body.len(), 1, "Should have only Ret");

        // Check Ret instruction
        use crate::mir::join_ir::JoinInst;
        if let JoinInst::Ret { value } = &func.body[0] {
            assert_eq!(value, &None);
        } else {
            panic!("Expected Ret instruction");
        }
    }

    #[test]
    fn test_return_variable_out_of_scope() {
        use crate::ast::{ASTNode, Span};
        use crate::mir::control_tree::step_tree::AstNodeHandle;

        // Create StepTree with "return x" (variable)
        let return_ast = Box::new(ASTNode::Variable {
            name: "x".to_string(),
            span: Span::unknown(),
        });

        let mut tree = make_if_only_tree();
        tree.root = StepNode::Stmt {
            kind: StepStmtKind::Return {
                value_ast: Some(AstNodeHandle(return_ast)),
            },
            span: Span::unknown(),
        };

        // Lower to JoinModule - should return Ok(None) (out of scope for Phase 123)
        let result = StepTreeNormalizedShadowLowererBox::try_lower_if_only(&tree);
        assert!(result.is_ok());
        assert!(result.unwrap().is_none(), "Should return None for Phase 123 unsupported patterns");
    }

    #[test]
    fn test_if_minimal_compare() {
        use crate::ast::{ASTNode, BinaryOperator, LiteralValue, Span};
        use crate::mir::control_tree::step_tree::AstNodeHandle;

        // Create condition: flag == 1
        let cond_ast = Box::new(ASTNode::BinaryOp {
            operator: BinaryOperator::Equal,
            left: Box::new(ASTNode::Variable {
                name: "flag".to_string(),
                span: Span::unknown(),
            }),
            right: Box::new(ASTNode::Literal {
                value: LiteralValue::Integer(1),
                span: Span::unknown(),
            }),
            span: Span::unknown(),
        });

        // Create then branch: return 2
        let then_branch = Box::new(StepNode::Stmt {
            kind: StepStmtKind::Return {
                value_ast: Some(AstNodeHandle(Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(2),
                    span: Span::unknown(),
                }))),
            },
            span: Span::unknown(),
        });

        // Create else branch: return 3
        let else_branch = Some(Box::new(StepNode::Stmt {
            kind: StepStmtKind::Return {
                value_ast: Some(AstNodeHandle(Box::new(ASTNode::Literal {
                    value: LiteralValue::Integer(3),
                    span: Span::unknown(),
                }))),
            },
            span: Span::unknown(),
        }));

        // Create If node
        let mut tree = make_if_only_tree();
        tree.root = StepNode::If {
            cond: crate::mir::control_tree::step_tree::AstSummary::Other("test"),
            cond_ast: AstNodeHandle(cond_ast),
            then_branch,
            else_branch,
            span: Span::unknown(),
        };

        // Lower to JoinModule
        let result = StepTreeNormalizedShadowLowererBox::try_lower_if_only(&tree);
        assert!(result.is_ok());

        let (module, _meta) = result.unwrap().expect("Should generate JoinModule");

        // Verify structure
        assert_eq!(module.functions.len(), 1);
        let func = &module.functions.values().next().unwrap();

        // Should have: Const (lhs placeholder), Const (rhs), Compare, Const (return), Ret
        assert!(func.body.len() >= 4, "Should have at least 4 instructions");

        // Verify Compare instruction exists
        use crate::mir::join_ir::{JoinInst, MirLikeInst};
        let has_compare = func.body.iter().any(|inst| {
            matches!(
                inst,
                JoinInst::Compute(MirLikeInst::Compare { .. })
            )
        });
        assert!(has_compare, "Should have Compare instruction");
    }

    #[test]
    fn test_return_variable_from_env() {
        // Phase 124 P3: Test Return(Variable) when variable is in env (writes)
        use crate::ast::{ASTNode, Span};
        use crate::mir::control_tree::step_tree::AstNodeHandle;
        use std::collections::BTreeSet;

        // Create StepTree with "local x; return x"
        let mut tree = make_if_only_tree();

        // Add x to writes (simulating "local x" or assignment)
        tree.contract.writes.insert("x".to_string());

        // Create return x node
        tree.root = StepNode::Stmt {
            kind: StepStmtKind::Return {
                value_ast: Some(AstNodeHandle(Box::new(ASTNode::Variable {
                    name: "x".to_string(),
                    span: Span::unknown(),
                }))),
            },
            span: Span::unknown(),
        };

        // Lower to JoinModule - should succeed
        let result = StepTreeNormalizedShadowLowererBox::try_lower_if_only(&tree);
        assert!(result.is_ok());

        let option = result.unwrap();
        assert!(option.is_some(), "Should generate JoinModule for return x (x in writes)");

        let (module, _meta) = option.unwrap();

        // Verify structure
        assert_eq!(module.functions.len(), 1);
        let func = &module.functions.values().next().unwrap();

        // Should have exactly 1 parameter (x from writes/env)
        assert_eq!(func.params.len(), 1, "Should have 1 parameter (x)");

        // Should have exactly 1 instruction: Ret(Some(x))
        assert_eq!(func.body.len(), 1, "Should have 1 instruction (Ret)");

        // Verify Ret instruction returns the parameter (ValueId(1))
        use crate::mir::join_ir::JoinInst;
        use crate::mir::ValueId;
        if let JoinInst::Ret { value } = &func.body[0] {
            assert_eq!(value, &Some(ValueId(1)), "Should return ValueId(1) (parameter x)");
        } else {
            panic!("Expected Ret instruction");
        }
    }
}