oxur_repl/

wrapper.rs

//! Rust AST wrapper for REPL scaffolding
//!
//! Wraps user code with necessary boilerplate for REPL execution:
//! - Entry point functions
//! - Variable store access
//! - Export declarations
//!
//! Based on ODD-0026 Section 2.1

use thiserror::Error;

/// Wrapper errors
#[derive(Debug, Error)]
pub enum WrapperError {
    #[error("Invalid code structure: {0}")]
    InvalidCode(String),

    #[error("Variable name conflict: {0}")]
    NameConflict(String),
}

pub type Result<T> = std::result::Result<T, WrapperError>;

/// Wraps Rust code for REPL execution
///
/// Adds necessary scaffolding for dynamic library compilation and execution.
///
/// # Examples
///
/// ```
/// use oxur_repl::wrapper::RustAstWrapper;
///
/// let wrapper = RustAstWrapper::new();
///
/// let user_code = "let x = 42;\nprintln!(\"x = {}\", x);";
/// let wrapped = wrapper.wrap("test_key", user_code, None).unwrap();
///
/// // wrapped contains:
/// // - #[no_mangle]
/// // - pub extern "C" fn oxur_eval_test_key()
/// // - user code inside the function
/// ```
#[derive(Debug, Clone)]
pub struct RustAstWrapper {
    /// Whether to include debug information
    debug: bool,
}

impl RustAstWrapper {
    /// Create a new wrapper
    pub fn new() -> Self {
        Self { debug: false }
    }

    /// Create a new wrapper with debug mode
    pub fn with_debug(mut self, debug: bool) -> Self {
        self.debug = debug;
        self
    }

    /// Wrap user code for REPL execution
    ///
    /// # Arguments
    ///
    /// * `cache_key` - Cache key used for function naming
    /// * `user_code` - User's Rust code to wrap
    ///
    /// # Returns
    ///
    /// Complete Rust source code ready for compilation
    ///
    /// # Errors
    ///
    /// Returns error if the user code contains invalid patterns
    pub fn wrap(
        &self,
        cache_key: impl AsRef<str>,
        user_code: impl AsRef<str>,
        source_map: Option<&oxur_smap::SourceMap>,
    ) -> Result<String> {
        let cache_key = cache_key.as_ref();
        let user_code = user_code.as_ref();

        // Validate cache key is valid identifier
        if !is_valid_identifier(cache_key) {
            return Err(WrapperError::InvalidCode(format!(
                "Cache key must be valid identifier: {}",
                cache_key
            )));
        }

        // Parse user code into AST
        let mut user_stmts = self.parse_user_code(user_code)?;

        // Check if last statement is an expression (needs result capture)
        let needs_result_capture = if let Some(last_stmt) = user_stmts.last() {
            matches!(last_stmt, syn::Stmt::Expr(_, None))
        } else {
            false
        };

        // If last statement is an expression, wrap it to capture the result
        if needs_result_capture {
            if let Some(syn::Stmt::Expr(expr, None)) = user_stmts.pop() {
                // Generate code to capture result in a global static
                // The subprocess will export a function to read this after execution
                let result_capture = quote::quote! {
                    {
                        let __result_value = #expr;
                        let __result_string = format!("{:?}", __result_value);
                        // Store in a global variable that subprocess can read
                        unsafe {
                            OXUR_RESULT_BUFFER = Some(__result_string);
                        }
                    }
                };

                // Parse the generated code into a statement
                let result_stmt: syn::Stmt = syn::parse2(result_capture).map_err(|e| {
                    WrapperError::InvalidCode(format!("Failed to parse result capture: {}", e))
                })?;

                user_stmts.push(result_stmt);
            }
        }

        // Generate function name identifier
        let fn_name =
            syn::Ident::new(&format!("oxur_eval_{}", cache_key), proc_macro2::Span::call_site());

        // Generate wrapper function with global result buffer
        let wrapped_fn = if needs_result_capture {
            quote::quote! {
                // Global buffer for storing expression results
                static mut OXUR_RESULT_BUFFER: Option<String> = None;

                #[no_mangle]
                pub extern "C" fn #fn_name() {
                    #(#user_stmts)*
                }

                // Export function to read the result buffer
                #[no_mangle]
                pub extern "C" fn oxur_get_result() -> *const u8 {
                    unsafe {
                        if let Some(ref s) = OXUR_RESULT_BUFFER {
                            s.as_ptr()
                        } else {
                            std::ptr::null()
                        }
                    }
                }

                #[no_mangle]
                pub extern "C" fn oxur_get_result_len() -> usize {
                    unsafe {
                        OXUR_RESULT_BUFFER.as_ref().map(|s| s.len()).unwrap_or(0)
                    }
                }
            }
        } else {
            quote::quote! {
                #[no_mangle]
                pub extern "C" fn #fn_name() {
                    #(#user_stmts)*
                }
            }
        };

        // Parse back into syn::File for proper formatting
        let file: syn::File = syn::parse2(wrapped_fn).map_err(|e| {
            WrapperError::InvalidCode(format!("Failed to parse generated code: {}", e))
        })?;

        // Format with prettyplease
        let formatted = prettyplease::unparse(&file);

        // Add header
        let mut output = String::new();
        output.push_str("// Generated by Oxur REPL wrapper\n");
        output.push_str("// Do not edit manually\n\n");

        if self.debug {
            output.push_str("// DEBUG MODE ENABLED\n\n");
        }

        output.push_str(&formatted);

        // Post-process: insert source map comments if source_map is provided
        let final_output = if let Some(smap) = source_map {
            self.insert_source_map_comments(&output, user_stmts.len(), smap)
        } else {
            output
        };

        Ok(final_output)
    }

    /// Parse user code into a list of statements
    ///
    /// Handles both complete files and bare statements.
    fn parse_user_code(&self, code: &str) -> Result<Vec<syn::Stmt>> {
        // Try parsing as a file first
        if let Ok(file) = syn::parse_str::<syn::File>(code) {
            // Convert all items to statements
            let mut stmts = Vec::new();
            for item in file.items {
                // Keep all items as statements, including function definitions
                stmts.push(syn::Stmt::Item(item));
            }
            return Ok(stmts);
        }

        // Try parsing as a block (statements)
        let wrapped_code = format!("{{ {} }}", code);
        if let Ok(block) = syn::parse_str::<syn::Block>(&wrapped_code) {
            return Ok(block.stmts);
        }

        // Try parsing as a single expression
        if let Ok(expr) = syn::parse_str::<syn::Expr>(code) {
            // Wrap expression in a statement
            return Ok(vec![syn::Stmt::Expr(expr, None)]);
        }

        // Failed to parse
        Err(WrapperError::InvalidCode(format!("Could not parse user code as valid Rust: {}", code)))
    }

    /// Wrap code with variable store access
    ///
    /// Generates code that:
    /// 1. Loads specified variables from VariableStore
    /// 2. Executes user code
    /// 3. Stores variables back to VariableStore
    /// 4. (Phase 4) Adds source map comments for error translation
    ///
    /// # Arguments
    ///
    /// * `cache_key` - Cache key for function naming
    /// * `user_code` - User's Rust code to wrap
    /// * `variables` - Variables to load/store (name, type) pairs
    /// * `_source_map` - Optional source map for Phase 4 integration (currently unused)
    ///
    /// # Returns
    ///
    /// Complete Rust source code with VariableStore integration
    ///
    /// # Examples
    ///
    /// ```no_run
    /// use oxur_repl::wrapper::RustAstWrapper;
    ///
    /// let wrapper = RustAstWrapper::new();
    /// let code = "let z = x + y;";
    /// let vars = vec![("x".to_string(), "i32".to_string()),
    ///                 ("y".to_string(), "i32".to_string())];
    /// let wrapped = wrapper.wrap_with_store("test", code, &vars, None).unwrap();
    /// // Generated code will load x and y, execute code, store z back
    /// ```
    pub fn wrap_with_store(
        &self,
        cache_key: impl AsRef<str>,
        user_code: impl AsRef<str>,
        variables: &[(String, String)],
        _source_map: Option<&oxur_smap::SourceMap>,
    ) -> Result<String> {
        let cache_key = cache_key.as_ref();
        let user_code = user_code.as_ref();

        // Validate cache key
        if !is_valid_identifier(cache_key) {
            return Err(WrapperError::InvalidCode(format!(
                "Cache key must be valid identifier: {}",
                cache_key
            )));
        }

        // Parse user code
        let user_stmts = self.parse_user_code(user_code)?;

        // Extract new variables defined in user code
        let new_vars = self.extract_variables(user_code);

        // Generate variable load statements
        let var_loads = self.generate_var_loads(variables)?;

        // Generate variable store statements (for existing + new variables)
        let mut all_vars = variables.to_vec();
        for (name, ty) in &new_vars {
            // Add new variables that aren't already in the list
            if !all_vars.iter().any(|(n, _)| n == name) {
                all_vars.push((name.clone(), ty.clone()));
            }
        }
        let var_stores = self.generate_var_stores(&all_vars)?;

        // Generate function name
        let fn_name =
            syn::Ident::new(&format!("oxur_eval_{}", cache_key), proc_macro2::Span::call_site());

        // Generate complete function (WITHOUT source map comments in AST)
        // We'll add comments via string post-processing below
        let wrapped_fn = quote::quote! {
            #[no_mangle]
            pub extern "C" fn #fn_name() {
                #(#var_loads)*

                #(#user_stmts)*

                #(#var_stores)*
            }
        };

        // Parse and format
        let file: syn::File = syn::parse2(wrapped_fn).map_err(|e| {
            WrapperError::InvalidCode(format!("Failed to parse generated code: {}", e))
        })?;
        let formatted = prettyplease::unparse(&file);

        // Add header
        let mut output = String::new();
        output.push_str("// Generated by Oxur REPL wrapper\n");
        output.push_str("// Do not edit manually\n\n");

        if self.debug {
            output.push_str("// DEBUG MODE ENABLED\n\n");
        }

        output.push_str(&formatted);

        // Post-process: insert source map comments if source_map is provided
        let final_output = if let Some(source_map) = _source_map {
            self.insert_source_map_comments(&output, user_stmts.len(), source_map)
        } else {
            output
        };

        Ok(final_output)
    }

    /// Generate variable load statements
    ///
    /// Creates code to load variables from the global VariableStore.
    fn generate_var_loads(&self, variables: &[(String, String)]) -> Result<Vec<syn::Stmt>> {
        let mut stmts = Vec::new();

        for (var_name, var_type) in variables {
            // Parse the type
            let ty: syn::Type = syn::parse_str(var_type).map_err(|e| {
                WrapperError::InvalidCode(format!("Invalid type '{}': {}", var_type, e))
            })?;

            let var_ident = syn::Ident::new(var_name, proc_macro2::Span::call_site());

            // Generate: let var_name: var_type = oxur_repl::subprocess::with_store(|store| { ... });
            let stmt = quote::quote! {
                let #var_ident: #ty = oxur_repl::subprocess::with_store(|store| {
                    store.get::<#ty>(#var_name)
                        .cloned()
                        .unwrap_or_default()
                });
            };

            let parsed_stmt: syn::Stmt = syn::parse2(stmt).map_err(|e| {
                WrapperError::InvalidCode(format!("Failed to parse load statement: {}", e))
            })?;
            stmts.push(parsed_stmt);
        }

        Ok(stmts)
    }

    /// Generate variable store statements
    ///
    /// Creates code to store variables back to the global VariableStore.
    fn generate_var_stores(&self, variables: &[(String, String)]) -> Result<Vec<syn::Stmt>> {
        let mut stmts = Vec::new();

        for (var_name, _var_type) in variables {
            let var_ident = syn::Ident::new(var_name, proc_macro2::Span::call_site());

            // Generate: oxur_repl::subprocess::with_store(|store| { store.set(...); });
            let stmt = quote::quote! {
                oxur_repl::subprocess::with_store(|store| {
                    store.set(#var_name.to_string(), #var_ident);
                });
            };

            let parsed_stmt: syn::Stmt = syn::parse2(stmt).map_err(|e| {
                WrapperError::InvalidCode(format!("Failed to parse store statement: {}", e))
            })?;
            stmts.push(parsed_stmt);
        }

        Ok(stmts)
    }

    /// Extract variable definitions from user code
    ///
    /// Uses TypeInference to find all variable bindings and their types.
    ///
    /// # Returns
    ///
    /// A list of (variable_name, type_name) pairs found in the code.
    ///
    /// # Examples
    ///
    /// ```
    /// use oxur_repl::wrapper::RustAstWrapper;
    ///
    /// let wrapper = RustAstWrapper::new();
    /// let code = "let x: i32 = 42; let y = 100;";
    /// let vars = wrapper.extract_variables(code);
    ///
    /// assert!(vars.iter().any(|(name, _)| name == "x"));
    /// ```
    pub fn extract_variables(&self, user_code: impl AsRef<str>) -> Vec<(String, String)> {
        use crate::type_inference::TypeInference;

        let inference = TypeInference::new();
        let type_infos = inference.infer_from_code(user_code);

        // Convert TypeInfo to (name, type_string) pairs
        type_infos
            .into_iter()
            .map(|(name, type_info)| (name, type_info.type_name.clone()))
            .collect()
    }

    /// Annotate statements with source map comments for error translation
    ///
    /// Adds /* oxur_node=N */ comments before each statement based on the
    /// source map. This enables ErrorTranslator to map Rust compiler errors
    /// back to original Oxur source positions.
    ///
    /// # Phase 4: Source Map Integration
    ///
    /// This method implements the core of Phase 4's error position translation.
    /// The generated comments allow ErrorTranslator to:
    /// 1. Parse rustc error with file position
    /// 2. Extract /* oxur_node=N */ comment near error column
    /// 3. Look up NodeId in SourceMap to find original Oxur position
    ///
    /// # Arguments
    ///
    /// * `stmts` - Original statements to annotate
    /// * `source_map` - Source map containing NodeId → SourcePos mappings
    ///
    /// # Returns
    ///
    /// New list of statements with source map comments inserted
    ///
    /// # Note
    ///
    /// Currently uses a simple heuristic: assign sequential NodeIds to statements.
    /// In a full implementation, the parser/lowerer would track which NodeId
    /// corresponds to which AST node, and we'd look those up here.
    /// Insert source map comments into generated code (string-based post-processing)
    ///
    /// Since syn/quote don't preserve comments, we post-process the generated string.
    /// This inserts /* oxur_node=N */ comments before user code statements.
    fn insert_source_map_comments(
        &self,
        generated_code: &str,
        num_user_stmts: usize,
        source_map: &oxur_smap::SourceMap,
    ) -> String {
        let stats = source_map.stats();
        let mut lines: Vec<String> = generated_code.lines().map(|s| s.to_string()).collect();

        // Find the function body (after the opening brace of the function)
        let mut in_function_body = false;
        let mut inserted_count = 0;

        // We'll insert comments before user statements
        // User statements start after variable loads and before variable stores
        // Look for patterns like "let x: i32 = oxur_repl::subprocess::with_store"
        // which are variable loads, and statements after those are user code

        let mut insert_indices = Vec::new();

        for (idx, line) in lines.iter().enumerate() {
            // Skip until we're inside the function body
            if line.contains("pub extern \"C\" fn") {
                in_function_body = true;
                continue;
            }

            if !in_function_body {
                continue;
            }

            // Look for user code statements (not variable load/store)
            // User statements are after all "let x: T = ... with_store" lines
            // and before "oxur_repl::subprocess::with_store(|store| { store.set"
            let is_var_load = line.contains("with_store(|store| {") && line.contains("store.get");
            let is_var_store = line.contains("with_store(|store| {") && line.contains("store.set");

            // If it's not a variable load/store and has actual code (not just braces/comments)
            if !is_var_load
                && !is_var_store
                && !line.trim().is_empty()
                && !line.trim().starts_with('}')
                && !line.trim().starts_with("//")
                && inserted_count < num_user_stmts
            {
                // Check if this looks like a user statement
                // (contains let, or is an expression statement)
                if line.trim().starts_with("let ")
                    || (!line.contains("with_store") && line.trim().ends_with(';'))
                {
                    insert_indices.push(idx);
                    inserted_count += 1;
                }
            }
        }

        // Insert comments at the identified positions (in reverse to maintain indices)
        for (comment_idx, &line_idx) in insert_indices.iter().enumerate().rev() {
            let offset = num_user_stmts.saturating_sub(comment_idx);
            let node_id = stats.surface_nodes.saturating_sub(offset) as u32;
            let comment = format!(
                "{}/* oxur_node={} */",
                " ".repeat(lines[line_idx].len() - lines[line_idx].trim_start().len()), // Match indentation
                node_id
            );
            lines.insert(line_idx, comment);
        }

        lines.join("\n")
    }
}

impl Default for RustAstWrapper {
    fn default() -> Self {
        Self::new()
    }
}

/// Check if a string is a valid Rust identifier
fn is_valid_identifier(s: &str) -> bool {
    if s.is_empty() {
        return false;
    }

    let mut chars = s.chars();

    // First character must be alphabetic or underscore
    if let Some(first) = chars.next() {
        if !first.is_alphabetic() && first != '_' {
            return false;
        }
    }

    // Rest must be alphanumeric or underscore
    chars.all(|c| c.is_alphanumeric() || c == '_')
}

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

    #[test]
    fn test_wrapper_creation() {
        let wrapper = RustAstWrapper::new();
        assert!(!wrapper.debug);
    }

    #[test]
    fn test_wrapper_with_debug() {
        let wrapper = RustAstWrapper::new().with_debug(true);
        assert!(wrapper.debug);
    }

    #[test]
    fn test_wrap_simple_code() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let x = 42;";
        let wrapped = wrapper.wrap("test", user_code, None).expect("Wrapping failed");

        assert!(wrapped.contains("#[no_mangle]"));
        assert!(wrapped.contains("pub extern \"C\" fn oxur_eval_test()"));
        assert!(wrapped.contains("let x = 42;"));
    }

    #[test]
    fn test_wrap_multiline_code() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let x = 42;\nlet y = 100;\nprintln!(\"sum = {}\", x + y);";
        let wrapped = wrapper.wrap("multiline", user_code, None).expect("Wrapping failed");

        assert!(wrapped.contains("let x = 42;"));
        assert!(wrapped.contains("let y = 100;"));
        assert!(wrapped.contains("println!"));
    }

    #[test]
    fn test_wrap_with_debug() {
        let wrapper = RustAstWrapper::new().with_debug(true);

        let user_code = "let x = 1;";
        let wrapped = wrapper.wrap("debug_test", user_code, None).expect("Wrapping failed");

        assert!(wrapped.contains("DEBUG MODE ENABLED"));
    }

    #[test]
    fn test_wrap_invalid_cache_key() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let x = 1;";

        // Invalid cache keys
        assert!(wrapper.wrap("123invalid", user_code, None).is_err());
        assert!(wrapper.wrap("invalid-key", user_code, None).is_err());
        assert!(wrapper.wrap("", user_code, None).is_err());
    }

    #[test]
    fn test_wrap_valid_cache_keys() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let x = 1;";

        // Valid cache keys
        assert!(wrapper.wrap("valid", user_code, None).is_ok());
        assert!(wrapper.wrap("valid_key", user_code, None).is_ok());
        assert!(wrapper.wrap("_valid", user_code, None).is_ok());
        assert!(wrapper.wrap("valid123", user_code, None).is_ok());
    }

    #[test]
    fn test_is_valid_identifier() {
        assert!(is_valid_identifier("valid"));
        assert!(is_valid_identifier("_valid"));
        assert!(is_valid_identifier("valid123"));
        assert!(is_valid_identifier("Valid_Identifier_123"));

        assert!(!is_valid_identifier(""));
        assert!(!is_valid_identifier("123invalid"));
        assert!(!is_valid_identifier("invalid-key"));
        assert!(!is_valid_identifier("invalid.key"));
        assert!(!is_valid_identifier("invalid key"));
    }

    #[test]
    fn test_extract_variables() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let x: i32 = 42;\nlet y: i32 = 100;";
        let vars = wrapper.extract_variables(user_code);

        // Should extract both variables with their types
        assert!(vars.len() >= 2);
        assert!(vars.iter().any(|(name, ty)| name == "x" && ty.contains("i32")));
        assert!(vars.iter().any(|(name, ty)| name == "y" && ty.contains("i32")));
    }

    #[test]
    fn test_extract_variables_type_inference() {
        let wrapper = RustAstWrapper::new();

        // Type inference should work without explicit type annotations
        let user_code = "let x = 42;";
        let vars = wrapper.extract_variables(user_code);

        // Should infer i32 type
        assert!(vars.len() >= 1);
        // Type inference should detect integer literal
        let has_x = vars.iter().any(|(name, _ty)| name == "x");
        assert!(has_x, "Should extract variable 'x'");
    }

    #[test]
    fn test_extract_variables_shadowing() {
        let wrapper = RustAstWrapper::new();

        // Test variable shadowing
        let user_code = r#"
            let x = 42;
            let x = "hello";
        "#;
        let vars = wrapper.extract_variables(user_code);

        // Should detect both bindings (shadowing)
        // TypeInference tracks all bindings, even shadowed ones
        assert!(vars.len() >= 2);
        let x_count = vars.iter().filter(|(name, _)| name == "x").count();
        assert!(x_count >= 1, "Should extract at least one 'x' binding");
    }

    #[test]
    fn test_wrap_with_store_simple() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let z = x + y;";
        let variables =
            vec![("x".to_string(), "i32".to_string()), ("y".to_string(), "i32".to_string())];

        let wrapped = wrapper
            .wrap_with_store("store_test", user_code, &variables, None)
            .expect("Wrapping failed");

        // Should have function declaration
        assert!(wrapped.contains("#[no_mangle]"));
        assert!(wrapped.contains("pub extern \"C\" fn oxur_eval_store_test()"));

        // Should load x and y
        assert!(wrapped.contains("with_store"));
        assert!(wrapped.contains("\"x\"") || wrapped.contains("let x"));
        assert!(wrapped.contains("\"y\"") || wrapped.contains("let y"));

        // Should have user code
        assert!(wrapped.contains("let z = x + y"));

        // Should store variables back
        assert!(wrapped.contains("store.set"));
    }

    #[test]
    fn test_wrap_with_store_no_existing_vars() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let x = 42;";
        let variables = vec![]; // No existing variables

        let wrapped = wrapper
            .wrap_with_store("no_vars", user_code, &variables, None)
            .expect("Wrapping failed");

        // Should still generate function
        assert!(wrapped.contains("#[no_mangle]"));
        assert!(wrapped.contains("oxur_eval_no_vars"));

        // Should have user code
        assert!(wrapped.contains("let x = 42"));

        // Should store newly created variable
        assert!(wrapped.contains("store.set") || wrapped.contains("with_store"));
    }

    #[test]
    fn test_wrap_with_store_multiple_vars() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let sum = a + b + c;";
        let variables = vec![
            ("a".to_string(), "i32".to_string()),
            ("b".to_string(), "i32".to_string()),
            ("c".to_string(), "i32".to_string()),
        ];

        let wrapped = wrapper
            .wrap_with_store("multi_vars", user_code, &variables, None)
            .expect("Wrapping failed");

        // Should load all three variables
        assert!(wrapped.contains("\"a\"") || wrapped.contains("let a"));
        assert!(wrapped.contains("\"b\"") || wrapped.contains("let b"));
        assert!(wrapped.contains("\"c\"") || wrapped.contains("let c"));

        // Should store sum back
        assert!(wrapped.contains("store.set"));
    }

    #[test]
    fn test_default() {
        let wrapper1 = RustAstWrapper::default();
        let wrapper2 = RustAstWrapper::new();

        assert_eq!(wrapper1.debug, wrapper2.debug);
    }

    // Task 2.5: Integration tests

    #[test]
    fn test_generated_code_parses_as_valid_rust() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let result = x * 2 + y;";
        let variables =
            vec![("x".to_string(), "i32".to_string()), ("y".to_string(), "i32".to_string())];

        let wrapped =
            wrapper.wrap_with_store("test", user_code, &variables, None).expect("Wrapping failed");

        // The generated code should parse as valid Rust
        let parsed = syn::parse_file(&wrapped);
        assert!(parsed.is_ok(), "Generated code should parse as valid Rust: {}", wrapped);

        if let Ok(file) = parsed {
            // Should have at least one item (the function)
            assert!(!file.items.is_empty());
        }
    }

    #[test]
    fn test_round_trip_wrap_and_parse() {
        let wrapper = RustAstWrapper::new();

        let user_code = r#"
            let a = 10;
            let b = 20;
            let sum = a + b;
        "#;

        // Wrap without variables
        let wrapped1 = wrapper.wrap("trip1", user_code, None).expect("Wrap failed");

        // Parse the generated code
        let parsed = syn::parse_file(&wrapped1);
        assert!(parsed.is_ok());

        // Wrap with variables
        let vars = vec![("x".to_string(), "i32".to_string())];
        let wrapped2 = wrapper
            .wrap_with_store("trip2", user_code, &vars, None)
            .expect("Wrap with store failed");

        let parsed2 = syn::parse_file(&wrapped2);
        assert!(parsed2.is_ok());
    }

    #[test]
    fn test_end_to_end_variable_flow() {
        let wrapper = RustAstWrapper::new();

        // Scenario: REPL session with multiple evaluations

        // Evaluation 1: Create variables
        let code1 = "let x: i32 = 42; let y: i32 = 100;";
        let wrapped1 = wrapper.wrap_with_store("eval1", code1, &[], None).expect("Eval 1 failed");

        // Should create and store x and y
        assert!(wrapped1.contains("store.set"));
        assert!(syn::parse_file(&wrapped1).is_ok());

        // Extract variables from first evaluation
        let vars_after_eval1 = wrapper.extract_variables(code1);
        assert!(vars_after_eval1.len() >= 2);

        // Evaluation 2: Use existing variables
        let code2 = "let sum = x + y;";
        let wrapped2 = wrapper
            .wrap_with_store("eval2", code2, &vars_after_eval1, None)
            .expect("Eval 2 failed");

        // Should load x and y, create sum
        assert!(wrapped2.contains("with_store"));
        assert!(wrapped2.contains("\"x\"") || wrapped2.contains("let x"));
        assert!(wrapped2.contains("\"y\"") || wrapped2.contains("let y"));
        assert!(syn::parse_file(&wrapped2).is_ok());

        // Extract all variables after eval 2
        let mut all_vars = vars_after_eval1.clone();
        let new_vars = wrapper.extract_variables(code2);
        for (name, ty) in new_vars {
            if !all_vars.iter().any(|(n, _)| n == &name) {
                all_vars.push((name, ty));
            }
        }

        // Should now have x, y, and sum
        assert!(all_vars.len() >= 3);

        // Evaluation 3: Use all variables
        let code3 = "let product = sum * x;";
        let wrapped3 =
            wrapper.wrap_with_store("eval3", code3, &all_vars, None).expect("Eval 3 failed");

        assert!(wrapped3.contains("with_store"));
        assert!(syn::parse_file(&wrapped3).is_ok());
    }

    #[test]
    fn test_complex_type_handling() {
        let wrapper = RustAstWrapper::new();

        // Test with various Rust types
        let code = r#"
            let s: String = String::from("hello");
            let v: Vec<i32> = vec![1, 2, 3];
            let opt: Option<i32> = Some(42);
        "#;

        let vars = vec![
            ("existing_str".to_string(), "String".to_string()),
            ("existing_vec".to_string(), "Vec<i32>".to_string()),
        ];

        let wrapped =
            wrapper.wrap_with_store("complex", code, &vars, None).expect("Complex types failed");

        // Should handle complex types in variable loads/stores
        assert!(wrapped.contains("with_store"));
        assert!(syn::parse_file(&wrapped).is_ok());
    }

    // New tests for AST-based wrapping (Task 2.1)

    #[test]
    fn test_wrap_expression() {
        let wrapper = RustAstWrapper::new();

        // Test wrapping a single expression
        let user_code = "42";
        let wrapped = wrapper.wrap("expr_test", user_code, None).expect("Wrapping failed");

        assert!(wrapped.contains("#[no_mangle]"));
        assert!(wrapped.contains("pub extern \"C\" fn oxur_eval_expr_test()"));
        assert!(wrapped.contains("42"));
    }

    #[test]
    fn test_wrap_multiple_statements() {
        let wrapper = RustAstWrapper::new();

        let user_code = r#"
            let x = 10;
            let y = 20;
            let z = x + y;
        "#;
        let wrapped = wrapper.wrap("multi", user_code, None).expect("Wrapping failed");

        assert!(wrapped.contains("let x = 10;"));
        assert!(wrapped.contains("let y = 20;"));
        assert!(wrapped.contains("let z = x + y;"));
    }

    #[test]
    fn test_wrap_function_definition_as_expression() {
        let wrapper = RustAstWrapper::new();

        // When parsing a function, we extract the body statements
        let user_code = r#"
            fn add(a: i32, b: i32) -> i32 {
                a + b
            }
        "#;
        let wrapped = wrapper.wrap("func_def", user_code, None).expect("Wrapping failed");

        // The body expression gets wrapped
        assert!(wrapped.contains("oxur_eval_func_def"));
        assert!(wrapped.contains("a + b"));
    }

    #[test]
    fn test_wrap_invalid_syntax() {
        let wrapper = RustAstWrapper::new();

        // Invalid Rust code should return an error
        let user_code = "let x = ;";
        let result = wrapper.wrap("invalid", user_code, None);

        assert!(result.is_err());
        match result {
            Err(WrapperError::InvalidCode(msg)) => {
                assert!(msg.contains("Could not parse"));
            }
            _ => panic!("Expected InvalidCode error"),
        }
    }

    #[test]
    fn test_parse_user_code_as_file() {
        let wrapper = RustAstWrapper::new();

        let code = "fn main() { println!(\"hello\"); }";
        let stmts = wrapper.parse_user_code(code).expect("Parse failed");

        // Should extract statements from the main function
        assert!(!stmts.is_empty());
    }

    #[test]
    fn test_parse_user_code_as_statements() {
        let wrapper = RustAstWrapper::new();

        let code = "let x = 42; let y = 100;";
        let stmts = wrapper.parse_user_code(code).expect("Parse failed");

        assert_eq!(stmts.len(), 2);
    }

    #[test]
    fn test_parse_user_code_as_expression() {
        let wrapper = RustAstWrapper::new();

        let code = "42 + 58";
        let stmts = wrapper.parse_user_code(code).expect("Parse failed");

        assert_eq!(stmts.len(), 1);
    }

    // ===== Phase 4 Tests: Source Map Annotation =====

    #[test]
    fn test_wrap_with_store_and_source_map() {
        use oxur_smap::{new_node_id, SourceMap, SourcePos};

        let wrapper = RustAstWrapper::new();
        let mut source_map = SourceMap::new();

        // Simulate some surface nodes being added (would normally be done by parser)
        // For testing, we just need a source map with some stats
        let node1 = new_node_id();
        source_map.record_surface_node(node1, SourcePos::repl(1, 1, 10));
        let node2 = new_node_id();
        source_map.record_surface_node(node2, SourcePos::repl(2, 1, 15));

        let user_code = "let result = x + y;";
        let variables =
            vec![("x".to_string(), "i32".to_string()), ("y".to_string(), "i32".to_string())];

        // Wrap with source map
        let wrapped = wrapper
            .wrap_with_store("smap_test", user_code, &variables, Some(&source_map))
            .expect("Wrapping with source map failed");

        // Should contain source map comments
        assert!(wrapped.contains("/* oxur_node="), "Should have source map comments: {}", wrapped);

        // Should still have all the normal wrapping
        assert!(wrapped.contains("oxur_eval_smap_test"));
        assert!(wrapped.contains("with_store"));

        // Generated code should still parse
        let parsed = syn::parse_file(&wrapped);
        assert!(parsed.is_ok(), "Code with source map comments should parse: {}", wrapped);
    }

    #[test]
    fn test_wrap_with_store_without_source_map() {
        let wrapper = RustAstWrapper::new();

        let user_code = "let z = 42;";
        let variables = vec![];

        // Wrap without source map (None)
        let wrapped = wrapper
            .wrap_with_store("no_smap", user_code, &variables, None)
            .expect("Wrapping without source map failed");

        // Should NOT contain source map comments
        assert!(
            !wrapped.contains("/* oxur_node="),
            "Should not have source map comments when source_map is None"
        );

        // But should still be valid
        assert!(wrapped.contains("oxur_eval_no_smap"));
        let parsed = syn::parse_file(&wrapped);
        assert!(parsed.is_ok());
    }

    #[test]
    fn test_source_map_comments_format() {
        use oxur_smap::{new_node_id, SourceMap, SourcePos};

        let wrapper = RustAstWrapper::new();
        let mut source_map = SourceMap::new();

        // Add multiple nodes
        for i in 0..5 {
            let node = new_node_id();
            source_map.record_surface_node(node, SourcePos::repl(i + 1, 1, 10));
        }

        let user_code = r#"
            let a = 1;
            let b = 2;
            let c = a + b;
        "#;

        let wrapped = wrapper
            .wrap_with_store("format_test", user_code, &[], Some(&source_map))
            .expect("Wrapping failed");

        // Should have properly formatted comments
        // Check that comments match the /* oxur_node=N */ pattern
        assert!(wrapped.contains("/* oxur_node="));

        // Comments should be followed by code
        if let Some(idx) = wrapped.find("/* oxur_node=") {
            let after_comment = &wrapped[idx..];
            // After comment, there should be valid Rust code
            assert!(
                after_comment.contains("let") || after_comment.contains("fn"),
                "Comments should be followed by code"
            );
        }
    }

    #[test]
    fn test_insert_source_map_comments_preserves_semantics() {
        use oxur_smap::{new_node_id, SourceMap, SourcePos};

        let wrapper = RustAstWrapper::new();
        let mut source_map = SourceMap::new();

        let node = new_node_id();
        source_map.record_surface_node(node, SourcePos::repl(1, 1, 20));

        let user_code = "let x: i32 = 42;";

        // Generate wrapped code with source map
        let wrapped = wrapper
            .wrap_with_store("semantics_test", user_code, &[], Some(&source_map))
            .expect("Wrapping with source map failed");

        // Should contain the original variable and value
        assert!(
            wrapped.contains("x") && wrapped.contains("42"),
            "Annotated code should preserve semantics"
        );

        // Should contain source map comments
        assert!(wrapped.contains("/* oxur_node="), "Should have source map comments");

        // Code should still parse as valid Rust
        let parsed = syn::parse_file(&wrapped);
        assert!(parsed.is_ok(), "Code with comments should be valid Rust: {}", wrapped);
    }
}