python_ast/ast/tree/

module.rs

use std::{collections::HashMap, default::Default};

use tracing::info;
use proc_macro2::TokenStream;
use pyo3::{Borrowed, FromPyObject, PyAny, PyResult, prelude::PyAnyMethods};
use quote::{format_ident, quote};
use serde::{Deserialize, Serialize};

use crate::{CodeGen, CodeGenContext, Name, Object, PythonOptions, Statement, StatementType, ExprType, SymbolTableScopes};


#[derive(Clone, Debug, Serialize, Deserialize)]
pub enum Type {
    Unimplemented,
}

impl<'a, 'py> FromPyObject<'a, 'py> for Type {
    type Error = pyo3::PyErr;
    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
        info!("Type: {:?}", ob);
        Ok(Type::Unimplemented)
    }
}

/// Represents a module as imported from an ast. See the Module struct for the processed module.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct RawModule {
    pub body: Vec<Statement>,
    pub type_ignores: Vec<Type>,
}

// Extracted manually (not via derive) so a failing statement's precise error
// — which names the construct and its line — propagates instead of being
// replaced by a generic "failed to extract field RawModule.body" message.
impl<'a, 'py> FromPyObject<'a, 'py> for RawModule {
    type Error = pyo3::PyErr;
    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
        let body_attr = ob
            .getattr("body")
            .map_err(|e| crate::extraction_failure("module body", &ob, e))?;
        let body_list: Vec<pyo3::Bound<PyAny>> = body_attr
            .extract()
            .map_err(|e| crate::extraction_failure("module body", &ob, e))?;

        let mut body = Vec::with_capacity(body_list.len());
        for stmt in &body_list {
            body.push(Statement::extract(stmt.as_borrowed())?);
        }

        let type_ignores = ob
            .getattr("type_ignores")
            .and_then(|t| t.extract())
            .unwrap_or_default();

        Ok(Self { body, type_ignores })
    }
}

/// Represents a module as imported from an ast.
#[derive(Clone, Debug, Default, Serialize, Deserialize)]
pub struct Module {
    pub raw: RawModule,
    pub name: Option<Name>,
    pub doc: Option<String>,
    pub filename: Option<String>,
    pub attributes: HashMap<Name, String>,
}

impl<'a, 'py> FromPyObject<'a, 'py> for Module {
    type Error = pyo3::PyErr;
    fn extract(ob: Borrowed<'a, 'py, PyAny>) -> PyResult<Self> {
        // RawModule's extraction already produces precise per-statement
        // errors; don't re-wrap them here.
        let raw_module = ob.extract()?;

        Ok(Self {
            raw: raw_module,
            ..Default::default()
        })
    }
}

impl CodeGen for Module {
    type Context = CodeGenContext;
    type Options = PythonOptions;
    type SymbolTable = SymbolTableScopes;

    fn find_symbols(self, symbols: Self::SymbolTable) -> Self::SymbolTable {
        let mut symbols = symbols;
        symbols.new_scope();
        for s in self.raw.body {
            symbols = s.clone().find_symbols(symbols);
        }
        symbols
    }

    fn to_rust(
        self,
        ctx: Self::Context,
        options: Self::Options,
        symbols: Self::SymbolTable,
    ) -> Result<TokenStream, Box<dyn std::error::Error>> {
        let mut stream = TokenStream::new();

        // Capture the module's source filename before fields of `self` are
        // moved, so statement errors can point at the user's Python file.
        let module_filename = self
            .filename
            .clone()
            .or_else(|| self.name.as_ref().map(|n| format!("{}.py", n.id)))
            .unwrap_or_else(|| "<module>".to_string());

        // Add module-level documentation if available and not just an expression
        if let Some(docstring) = self.get_module_docstring() {
            // Only add module docs if there are multiple statements or if this seems to be a real module docstring
            if self.raw.body.len() > 1 || self.looks_like_module_docstring() {
                let doc_lines: Vec<_> = docstring
                    .lines()
                    .map(|line| {
                        if line.trim().is_empty() {
                            quote! { #![doc = ""] }
                        } else {
                            let doc_line = format!("{}", line);
                            quote! { #![doc = #doc_line] }
                        }
                    })
                    .collect();
                stream.extend(quote! { #(#doc_lines)* });
                
                // Add generated by comment only when we have actual module docs
                let generated_comment =
                    format!("Generated from Python file: {}", module_filename);
                stream.extend(quote! { #![doc = #generated_comment] });
            }
        }
        
        if options.with_std_python {
            // For imports, always use "stdpython" since that's the actual crate name
            // The runtime specification is just for dependency management
            stream.extend(quote!(use stdpython::*;));
        }

        // Under no_std the prelude has no String/Vec/format!: bring the
        // alloc surface generated code leans on into scope. Emitted per
        // module — each module file in the generated crate lowers through
        // here — while `extern crate alloc` itself binds locally too, so
        // the imports resolve regardless of what the crate root declares.
        // allow(unused_imports): this is harness plumbing, and the lint
        // posture should keep surfacing only source-Python weaknesses.
        if options.no_std {
            stream.extend(quote! {
                extern crate alloc;
                #[allow(unused_imports)]
                use alloc::{
                    format, vec, borrow::ToOwned, string::String, string::ToString,
                    vec::Vec,
                };
            });
        }
        
        // Add async runtime dependency if async functions are detected
        // We'll check this early so we can add the import at the top
        let needs_async_runtime = self.raw.body.iter().any(|s| {
            matches!(&s.statement, crate::StatementType::AsyncFunctionDef(_))
        });
        
        if needs_async_runtime {
            let runtime_import = format_ident!("{}", options.async_runtime.import());
            stream.extend(quote!(use #runtime_import;));
        }
        
        let mut main_body_stmts = Vec::new();
        let mut has_main_code = false;
        let mut has_async_functions = false;
        let mut module_init_stmts = Vec::new();
        let mut has_module_init_code = false;
        let mut is_simple_main_call_pattern = false;
        // The raw statements behind main_body_stmts/module_init_stmts, kept
        // so assigned names can be hoisted to declarations (assignments
        // lower to plain stores; see collect_assigned_names).
        let mut main_body_raw: Vec<crate::Statement> = Vec::new();
        let mut module_init_raw: Vec<crate::Statement> = Vec::new();

        // Module-level names assigned exactly once from a literal, with no
        // other store ANYWHERE in module scope, are CONSTANTS: they lower
        // to static items so functions in the module can see them (Python
        // module globals are visible everywhere; a store hidden inside
        // __module_init__ is not). The tally recurses through module-level
        // control flow — `DEBUG = False` conditionally overwritten inside
        // an `if` is NOT a constant, and treating it as one would silently
        // freeze the original value. (Function bodies need no scan: without
        // `global` — unsupported — their assignments create locals.)
        let mut module_assign_counts: std::collections::HashMap<String, usize> =
            std::collections::HashMap::new();
        count_module_stores(&self.raw.body, &mut module_assign_counts);

        // A user `main` that returns a value cannot serve as the Rust entry
        // point directly (Result<i64, _> does not implement Termination);
        // route it through the renaming wrapper, which discards the value —
        // exactly what Python's `if __name__: main()` does.
        let user_main_returns_value = self.raw.body.iter().any(|s| {
            matches!(
                &s.statement,
                crate::StatementType::FunctionDef(f)
                    if f.name == "main" && f.resolved_return_type().is_some()
            )
        });

        for s in self.raw.body {
            // Check if this statement is an async function
            if let crate::StatementType::AsyncFunctionDef(_) = &s.statement {
                has_async_functions = true;
            }

            // Check for if __name__ == "__main__" blocks at the AST level before generating code
            if let crate::StatementType::If(if_stmt) = &s.statement {
                let test_str = format!("{:?}", if_stmt.test);
                if test_str.contains("__name__") && test_str.contains("__main__") {
                    // Check if this is a simple main() call pattern
                    let is_simple_main_call = Self::is_simple_main_call_block(&if_stmt.body)
                        && !user_main_returns_value;
                    
                    if is_simple_main_call {
                        // For simple main() calls, we'll use the user's main function directly
                        // Set a flag to indicate we should not rename the main function
                        has_main_code = true;
                        is_simple_main_call_pattern = true;
                        // Don't collect the main body statements - we'll use user's main directly
                    } else {
                        // This is a complex __name__ == "__main__" block - collect its body for main function
                        for body_stmt in &if_stmt.body {
                            let stmt_token = body_stmt
                                .clone()
                                .to_rust(ctx.clone(), options.clone(), symbols.clone())
                                .map_err(|e| wrap_module_error(&module_filename, e))?;
                            if !stmt_token.to_string().trim().is_empty() {
                                main_body_stmts.push(stmt_token);
                                main_body_raw.push(body_stmt.clone());
                                has_main_code = true;
                            }
                        }
                    }
                    // Skip generating this if statement - we've processed its contents
                    continue;
                }
            }
            
            // Module-level constants become static items visible to every
            // function in the module.
            if let crate::StatementType::Assign(a) = &s.statement {
                if let [crate::ExprType::Name(n)] = a.targets.as_slice() {
                    if module_assign_counts.get(&n.id) == Some(&1) {
                        if let Some(ty) = const_static_type(&a.value) {
                            let ident = crate::safe_ident(&n.id);
                            let value = a.value.clone().to_rust(
                                ctx.clone(),
                                options.clone(),
                                symbols.clone(),
                            )?;
                            stream.extend(quote!(pub static #ident: #ty = #value;));
                            continue;
                        }
                    }
                }
            }

            // Categorize statements into declarations vs executable code
            let is_declaration = Self::is_declaration_statement(&s.statement);

            let statement = s
                .clone()
                .to_rust(ctx.clone(), options.clone(), symbols.clone())
                .map_err(|e| wrap_module_error(&module_filename, e))?;
            
            if statement.to_string() != "" {
                if is_declaration {
                    // Declarations go at module level (functions, classes, imports)
                    stream.extend(statement);
                } else {
                    // Executable statements go in module initialization function
                    module_init_stmts.push(statement);
                    module_init_raw.push(s.clone());
                    has_module_init_code = true;
                }
            }
        }

        // Hoist assigned names to declarations at the top of each generated
        // scope (assignments themselves lower to plain stores).
        let init_decls = hoisted_declarations(&module_init_raw, &ctx, &symbols);
        if !init_decls.is_empty() {
            module_init_stmts.insert(0, init_decls);
        }
        let main_decls = hoisted_declarations(&main_body_raw, &ctx, &symbols);
        if !main_decls.is_empty() {
            main_body_stmts.insert(0, main_decls);
        }

        // Generate module initialization function if needed. Like all
        // generated functions it returns Result so module-level raises and
        // calls propagate.
        if has_module_init_code {
            stream.extend(quote! {
                fn __module_init__() -> Result<(), PyException> {
                    #(#module_init_stmts;)*
                    Ok(())
                }
            });
        }
        
        // A `__main__` block wants a process entry point, and a no_std
        // target has no OS to enter from: refuse loudly instead of emitting
        // a fn main() that cannot link.
        if has_main_code && options.no_std {
            return Err(
                "`if __name__ == \"__main__\":` needs a process entry point, which \
                 the no_std profile does not provide; remove the block (convert the \
                 module as a library) or convert without the no_std profile"
                    .to_string()
                    .into(),
            );
        }

        // If we collected any main code, generate a single consolidated main function
        if has_main_code {
            if is_simple_main_call_pattern {
                // Simple main() call pattern - use user's main function directly as Rust entry point
                // Don't rename the user's main function, just add module init call if needed
                let stream_str = stream.to_string();
                
                // Check if the user's main function is async
                let user_main_is_async = stream_str.contains("pub async fn main (");
                
                if user_main_is_async {
                    // User's async main becomes the Rust entry point
                    let runtime_attr = options.async_runtime.main_attribute();
                    let attr_tokens: proc_macro2::TokenStream = runtime_attr.parse()
                        .unwrap_or_else(|_| quote!(tokio::main)); // fallback to tokio::main
                    
                    // Replace the user's function signature and add attributes
                    let new_stream_str = stream_str
                        .replace("pub async fn main (", &format!("#[{}] async fn main(", runtime_attr));
                    stream = new_stream_str.parse::<proc_macro2::TokenStream>()
                        .unwrap_or_else(|_| stream);
                        
                    // If we have module init code, we need to modify the user's main to call it first
                    if has_module_init_code {
                        // This is more complex - we'd need to modify the user's main function body
                        // For now, let's fall back to the rename approach for async functions with module init
                        let renamed_stream_str = Self::rename_main_function_and_references(&stream_str);
                        stream = renamed_stream_str.parse::<proc_macro2::TokenStream>()
                            .unwrap_or_else(|_| stream);

                        stream.extend(quote! {
                            #[#attr_tokens]
                            async fn main() {
                                let __rython_result: Result<(), PyException> = async {
                                    __module_init__()?;
                                    python_main().await?;
                                    Ok(())
                                }.await;
                                if let Err(e) = __rython_result {
                                    eprintln!("{}", e);
                                    std::process::exit(1);
                                }
                            }
                        });
                    }
                } else {
                    // User's sync main becomes the Rust entry point
                    // Need to modify the function to match Rust main signature requirements
                    let new_stream_str = Self::convert_python_main_to_rust_entry_point(&stream_str);
                    stream = new_stream_str.parse::<proc_macro2::TokenStream>()
                        .unwrap_or_else(|_| stream);
                    
                    // If we have module init code, we need to modify the user's main to call it first
                    if has_module_init_code {
                        // For simplicity, we'll use the rename approach when module init is needed
                        let renamed_stream_str = Self::rename_main_function_and_references(&stream_str);
                        stream = renamed_stream_str.parse::<proc_macro2::TokenStream>()
                            .unwrap_or_else(|_| stream);

                        stream.extend(quote! {
                            fn main() {
                                let __rython_result = (|| -> Result<(), PyException> {
                                    __module_init__()?;
                                    python_main()?;
                                    Ok(())
                                })();
                                if let Err(e) = __rython_result {
                                    eprintln!("{}", e);
                                    std::process::exit(1);
                                }
                            }
                        });
                    }
                }
            } else {
                // Complex main block - use existing behavior (rename user's main)
                let stream_str = stream.to_string();
                let has_python_main = stream_str.contains("pub fn main (") || stream_str.contains("pub async fn main (");
                
                if has_python_main {
                    // Rename the Python function to avoid conflict with Rust entry point
                    let new_stream_str = Self::rename_main_function_and_references(&stream_str);
                    stream = new_stream_str.parse::<proc_macro2::TokenStream>()
                        .unwrap_or_else(|_| stream);
                    
                    // Update main_body_stmts to call python_main instead of main
                    for stmt in &mut main_body_stmts {
                        let stmt_str = stmt.to_string();
                        let updated_stmt_str = Self::update_main_references(&stmt_str);
                        if updated_stmt_str != stmt_str {
                            if let Ok(new_stmt) = updated_stmt_str.parse::<proc_macro2::TokenStream>() {
                                *stmt = new_stmt;
                            }
                        }
                    }
                }
                
                // Generate the Rust entry point as main() - async if needed
                if needs_async_runtime || has_async_functions {
                    // Parse the runtime attribute string into tokens
                    let runtime_attr = options.async_runtime.main_attribute();
                    let attr_tokens: proc_macro2::TokenStream = runtime_attr.parse()
                        .unwrap_or_else(|_| quote!(tokio::main)); // fallback to tokio::main
                    
                    let init_call = if has_module_init_code {
                        quote!(__module_init__()?;)
                    } else {
                        quote!()
                    };
                    stream.extend(quote! {
                        #[#attr_tokens]
                        async fn main() {
                            let __rython_result: Result<(), PyException> = async {
                                #init_call
                                #(#main_body_stmts;)*
                                Ok(())
                            }.await;
                            if let Err(e) = __rython_result {
                                eprintln!("{}", e);
                                std::process::exit(1);
                            }
                        }
                    });
                } else {
                    let init_call = if has_module_init_code {
                        quote!(__module_init__()?;)
                    } else {
                        quote!()
                    };
                    stream.extend(quote! {
                        fn main() {
                            let __rython_result = (|| -> Result<(), PyException> {
                                #init_call
                                #(#main_body_stmts;)*
                                Ok(())
                            })();
                            if let Err(e) = __rython_result {
                                eprintln!("{}", e);
                                std::process::exit(1);
                            }
                        }
                    });
                }
            }
        } else if has_module_init_code {
            // No main block, but we have module initialization code
            // Generate a main function that just runs module initialization
            stream.extend(quote! {
                fn main() {
                    if let Err(e) = __module_init__() {
                        eprintln!("{}", e);
                        std::process::exit(1);
                    }
                }
            });
        }
        Ok(stream)
    }
}

/// Count stores to each name across MODULE scope: top-level assignments
/// plus everything nested in module-level control flow — if/while/for
/// bodies (and the for target itself, which rebinds every iteration),
/// with bodies and their `as` targets, try bodies and handlers. Nested
/// stores count double so a name assigned once at top level and again in
/// a branch never tallies as once-assigned. Function and class bodies are
/// their own scopes and are not walked.
fn count_module_stores(
    body: &[crate::Statement],
    counts: &mut std::collections::HashMap<String, usize>,
) {
    fn bump_target(target: &crate::ExprType, by: usize, counts: &mut std::collections::HashMap<String, usize>) {
        match target {
            crate::ExprType::Name(n) => {
                *counts.entry(n.id.clone()).or_insert(0) += by;
            }
            crate::ExprType::Tuple(t) => {
                for elt in &t.elts {
                    bump_target(elt, by, counts);
                }
            }
            _ => {}
        }
    }
    for s in body {
        match &s.statement {
            crate::StatementType::Assign(a) => {
                for target in &a.targets {
                    bump_target(target, 1, counts);
                }
            }
            crate::StatementType::AugAssign(a) => bump_target(&a.target, 2, counts),
            crate::StatementType::If(i) => {
                let mut nested = std::collections::HashMap::new();
                count_module_stores(&i.body, &mut nested);
                count_module_stores(&i.orelse, &mut nested);
                for (name, n) in nested {
                    *counts.entry(name).or_insert(0) += n * 2;
                }
            }
            crate::StatementType::While(w) => {
                let mut nested = std::collections::HashMap::new();
                count_module_stores(&w.body, &mut nested);
                count_module_stores(&w.orelse, &mut nested);
                for (name, n) in nested {
                    *counts.entry(name).or_insert(0) += n * 2;
                }
            }
            crate::StatementType::For(f) => {
                bump_target(&f.target, 2, counts);
                let mut nested = std::collections::HashMap::new();
                count_module_stores(&f.body, &mut nested);
                count_module_stores(&f.orelse, &mut nested);
                for (name, n) in nested {
                    *counts.entry(name).or_insert(0) += n * 2;
                }
            }
            crate::StatementType::With(w) => {
                for item in &w.items {
                    if let Some(vars) = &item.optional_vars {
                        bump_target(vars, 2, counts);
                    }
                }
                let mut nested = std::collections::HashMap::new();
                count_module_stores(&w.body, &mut nested);
                for (name, n) in nested {
                    *counts.entry(name).or_insert(0) += n * 2;
                }
            }
            crate::StatementType::Try(t) => {
                let mut nested = std::collections::HashMap::new();
                count_module_stores(&t.body, &mut nested);
                for h in &t.handlers {
                    count_module_stores(&h.body, &mut nested);
                }
                count_module_stores(&t.orelse, &mut nested);
                count_module_stores(&t.finalbody, &mut nested);
                for (name, n) in nested {
                    *counts.entry(name).or_insert(0) += n * 2;
                }
            }
            // Function and class bodies are separate scopes.
            _ => {}
        }
    }
}

/// The static-item type for a module-level constant, when its value is a
/// literal a static can hold (numbers, bools, strings — including a
/// leading unary minus). Non-literal or reassigned module globals keep
/// the old __module_init__ lowering, where referencing them from a
/// function is a loud compile error rather than a silent divergence.
fn const_static_type(value: &crate::ExprType) -> Option<TokenStream> {
    match value {
        crate::ExprType::Constant(c) => match &c.0 {
            Some(litrs::Literal::Integer(_)) => Some(quote!(i64)),
            Some(litrs::Literal::Float(_)) => Some(quote!(f64)),
            Some(litrs::Literal::Bool(_)) => Some(quote!(bool)),
            Some(litrs::Literal::String(_)) => Some(quote!(&'static str)),
            _ => None,
        },
        crate::ExprType::UnaryOp(op) => {
            if !matches!(op.op, crate::ast::tree::unary_op::Ops::USub) {
                return None;
            }
            match const_static_type(&op.operand) {
                Some(ty) if ty.to_string() == "i64" || ty.to_string() == "f64" => Some(ty),
                _ => None,
            }
        }
        _ => None,
    }
}

/// Declarations for every name assigned in a statement list, so
/// nested-block assignments store into scope-level variables instead of
/// creating shadowing bindings. Scope analysis decides which need `mut`.
fn hoisted_declarations(
    body: &[crate::Statement],
    ctx: &crate::CodeGenContext,
    symbols: &crate::SymbolTableScopes,
) -> TokenStream {
    // Class-aware mutation facts need the block's own assignments in the
    // symbol table (`c = Counter(...)` then `c.bump()` needs `c` mutable).
    let mut symbols = symbols.clone();
    for s in body {
        symbols = s.clone().find_symbols(symbols);
    }
    let scope =
        crate::analyze_scope_with(body, &[], &crate::class_call_resolver(ctx, &symbols));
    let mut out = TokenStream::new();
    for name in &scope.assigned {
        let ident = crate::safe_ident(name);
        if scope.needs_mut.contains(name) {
            out.extend(quote!(let mut #ident;));
        } else {
            out.extend(quote!(let #ident;));
        }
    }
    out
}

/// Rebuild a statement-level codegen error so it points at the module's real
/// source file. Statement errors carry a `<module>` placeholder filename in
/// their location; this substitutes the actual filename and preserves the
/// structured fields (message, help) so downstream consumers — the proc
/// macro in particular — can render precise diagnostics.
fn wrap_module_error(
    filename: &str,
    e: Box<dyn std::error::Error>,
) -> Box<dyn std::error::Error> {
    if let Some(inner) = e.downcast_ref::<crate::Error>() {
        let message = inner.get_field("message").unwrap_or_default().to_string();
        let location = inner
            .get_field("location")
            .unwrap_or("<module>")
            .replace("<module>", filename);
        let help = inner.get_field("help").unwrap_or_default().to_string();
        return Box::from(crate::codegen_error(
            crate::SourceLocation::new(location),
            message,
            help,
        ));
    }
    Box::from(crate::codegen_error(
        crate::SourceLocation::new(filename),
        crate::format_error_chain(e.as_ref()),
        "",
    ))
}

impl Module {
    /// Check if the __name__ == "__main__" block contains only a simple call to main()
    /// This includes patterns like:
    /// - main()
    /// - result = main()
    /// - sys.exit(main())
    fn is_simple_main_call_block(body: &[crate::Statement]) -> bool {
        // Must have exactly one statement
        if body.len() != 1 {
            return false;
        }
        
        let stmt = &body[0];
        match &stmt.statement {
            // Pattern 1: main() - direct call as expression statement
            crate::StatementType::Expr(expr) => {
                Self::is_main_function_call(&expr.value)
            },
            // Pattern 2: result = main() - assignment from main call
            crate::StatementType::Assign(assign) => {
                // Should have exactly one target and the value should be a main() call
                assign.targets.len() == 1 && Self::is_main_function_call(&assign.value)
            },
            // Pattern 3: sys.exit(main()) - call with main() as argument
            crate::StatementType::Call(call) => {
                // Check if any of the arguments is a main() call
                call.args.iter().any(|arg| Self::is_main_function_call(arg))
            },
            _ => false,
        }
    }
    
    /// Check if an expression is a call to a function named "main"
    fn is_main_function_call(expr: &crate::ExprType) -> bool {
        match expr {
            crate::ExprType::Call(call) => {
                match call.func.as_ref() {
                    crate::ExprType::Name(name) => name.id == "main",
                    _ => false,
                }
            },
            _ => false,
        }
    }
    
    /// Determine if a statement is a declaration (can stay at module level) or executable code (needs to go in init function)
    fn is_declaration_statement(stmt_type: &crate::StatementType) -> bool {
        use crate::StatementType::*;
        match stmt_type {
            // These are declarations that can stay at module level
            FunctionDef(_) | AsyncFunctionDef(_) | ClassDef(_) | Import(_) | ImportFrom(_) => true,
            
            // Standalone expressions can stay at module level (e.g., constants, simple values)
            // These are typically used in tests or simple modules
            Expr(expr) => Self::is_simple_expression(&expr.value),
            
            // These are executable statements that must go in the init function
            Assign(_) | AugAssign(_) | Call(_) | Return(_) |
            If(_) | For(_) | While(_) | Try(_) | With(_) | AsyncWith(_) | AsyncFor(_) |
            Raise(_) | Assert { .. } | Pass | Break | Continue => false,
            
            // Handle unimplemented statements conservatively as executable
            Unimplemented(_) => false,
        }
    }
    
    /// Check if an expression is simple enough to remain at module level
    fn is_simple_expression(expr: &crate::ExprType) -> bool {
        use crate::ExprType::*;
        match expr {
            // Simple constants and literals can stay at module level
            Constant(_) | Name(_) | NoneType(_) => true,
            
            // Allow unary operations for single-expression modules (test compatibility)
            UnaryOp(_) => true,
            
            // Function calls and complex expressions should go in init
            Call(_) | BinOp(_) | Compare(_) | BoolOp(_) | 
            IfExp(_) | Dict(_) | Set(_) | List(_) | Tuple(_) | ListComp(_) |
            Lambda(_) | Attribute(_) | Subscript(_) | Starred(_) |
            DictComp(_) | SetComp(_) | GeneratorExp(_) | Await(_) | 
            Yield(_) | YieldFrom(_) | FormattedValue(_) | JoinedStr(_) |
            NamedExpr(_) => false,
            
            // Be conservative about other expression types
            Unimplemented(_) | Unknown => false,
        }
    }
    
    /// Rename the main function definition and update all references to it throughout the code
    fn rename_main_function_and_references(code: &str) -> String {
        // First, rename the function definitions
        let code = code
            .replace("pub async fn main (", "pub async fn python_main (")
            .replace("pub fn main (", "pub fn python_main (");
        
        // Then update all references using the comprehensive reference updater
        Self::update_main_references(&code)
    }
    
    /// Convert a Python main function to be suitable as a Rust entry point
    /// This handles return value conversion and signature requirements
    fn convert_python_main_to_rust_entry_point(code: &str) -> String {
        use regex::Regex;
        
        // Replace "pub fn main (" with "fn main("
        let code = code.replace("pub fn main (", "fn main(");
        
        // Handle return statements in the main function
        // We need to wrap the function body to ignore return values
        let main_fn_pattern = Regex::new(r"fn main\(\s*\)\s*\{([^}]*)\}").unwrap();
        
        if let Some(captures) = main_fn_pattern.captures(&code) {
            let body = captures.get(1).map_or("", |m| m.as_str());
            
            // Check if the body contains return statements
            if body.contains("return ") {
                // Wrap the original function as python_main and create new main that ignores return
                let new_code = code.replace("fn main(", "fn python_main(");
                format!("{}\n\nfn main() {{\n    let _ = python_main();\n}}", new_code)
            } else {
                // No return statements, use the function as-is
                code
            }
        } else {
            // Couldn't parse the function, fall back to original
            code
        }
    }
    
    /// Update all references to main() function calls with python_main() calls
    /// This uses regex to handle various call patterns with parameters
    fn update_main_references(code: &str) -> String {
        use regex::Regex;
        
        // Pattern 1: main(...) - function calls with any arguments (including empty)
        // This pattern matches "main(" and lets us replace the function name
        let call_pattern = Regex::new(r"\bmain\s*\(").unwrap();
        let mut result = call_pattern.replace_all(code, "python_main(").to_string();
        
        // Pattern 2: Handle method calls like obj.call_main() -> obj.call_python_main()
        let method_pattern = Regex::new(r"\.call_main\s*\(").unwrap();
        result = method_pattern.replace_all(&result, ".call_python_main(").to_string();
        
        // Pattern 3: Handle assignment patterns like "result = main" (without parentheses)
        // We need to be careful not to match function definitions or other contexts
        let assignment_pattern = Regex::new(r"=\s+main\b").unwrap();
        result = assignment_pattern.replace_all(&result, "= python_main").to_string();
        
        // Pattern 4: Handle return statements like "return main"
        let return_pattern = Regex::new(r"return\s+main\b").unwrap();
        result = return_pattern.replace_all(&result, "return python_main").to_string();
        
        result
    }
    
    fn get_module_docstring(&self) -> Option<String> {
        if self.raw.body.is_empty() {
            return None;
        }
        
        // Check if the first statement is a string constant (docstring)
        let first_stmt = &self.raw.body[0];
        match &first_stmt.statement {
            StatementType::Expr(expr) => match &expr.value {
                ExprType::Constant(c) => {
                    let raw_string = c.to_string();
                    Some(self.format_module_docstring(&raw_string))
                },
                _ => None,
            },
            _ => None,
        }
    }
    
    fn format_module_docstring(&self, raw: &str) -> String {
        // Remove surrounding quotes
        let content = raw.trim_matches('"');
        
        // Split into lines and clean up Python-style indentation
        let lines: Vec<&str> = content.lines().collect();
        if lines.is_empty() {
            return String::new();
        }
        
        // For module docstrings, preserve more of the original formatting
        let mut formatted = Vec::new();
        
        for line in lines {
            let cleaned = line.trim();
            if !cleaned.is_empty() {
                formatted.push(cleaned.to_string());
            } else {
                formatted.push(String::new());
            }
        }
        
        formatted.join("\n")
    }
    
    fn looks_like_module_docstring(&self) -> bool {
        if self.raw.body.is_empty() {
            return false;
        }
        
        // Check if the first statement looks like a module docstring
        let first_stmt = &self.raw.body[0];
        if let StatementType::Expr(expr) = &first_stmt.statement {
            if let ExprType::Constant(c) = &expr.value {
                let raw_string = c.to_string();
                let content = raw_string.trim_matches('"');
                
                // Heuristics to detect if this is a module docstring vs just a string expression:
                // 1. Contains multiple lines
                // 2. Contains common docstring keywords
                // 3. Looks like documentation rather than a simple string
                return content.lines().count() > 1 
                    || content.to_lowercase().contains("module")
                    || content.to_lowercase().contains("this ")
                    || content.len() > 50; // Longer strings are more likely to be docstrings
            }
        }
        false
    }
}

impl Object for Module {
    /// __dir__ is called to list the attributes of the object.
    fn __dir__(&self) -> Vec<impl AsRef<str>> {
        // XXX - Make this meaningful.
        vec![
            "__class__",
            "__class_getitem__",
            "__contains__",
            "__delattr__",
            "__delitem__",
            "__dir__",
            "__doc__",
            "__eq__",
            "__format__",
            "__ge__",
            "__getattribute__",
            "__getitem__",
            "__getstate__",
            "__gt__",
            "__hash__",
            "__init__",
            "__init_subclass__",
            "__ior__",
            "__iter__",
            "__le__",
            "__len__",
            "__lt__",
            "__ne__",
            "__new__",
            "__or__",
            "__reduce__",
            "__reduce_ex__",
            "__repr__",
            "__reversed__",
            "__ror__",
            "__setattr__",
            "__setitem__",
            "__sizeof__",
            "__str__",
            "__subclasshook__",
            "clear",
            "copy",
            "fromkeys",
            "get",
            "items",
            "keys",
            "pop",
            "popitem",
            "setdefault",
            "update",
            "values",
        ]
    }
}

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

    #[test]
    fn can_we_print() {
        let options = PythonOptions::default();
        let result = crate::parse(
            "#test comment
def foo():
    print(\"Test print.\")
",
            "test_case.py",
        )
        .unwrap();
        info!("Python tree: {:?}", result);
        //info!("{}", result);

        let code = result.to_rust(
            CodeGenContext::Module("test_case".to_string()),
            options,
            SymbolTableScopes::new(),
        );
        info!("module: {:?}", code);
    }

    #[test]
    fn can_we_import() {
        let result = crate::parse("import ast", "ast.py").unwrap();
        let options = PythonOptions::default();
        info!("{:?}", result);

        let code = result.to_rust(
            CodeGenContext::Module("test_case".to_string()),
            options,
            SymbolTableScopes::new(),
        );
        info!("module: {:?}", code);
    }

    #[test]
    fn can_we_import2() {
        let result = crate::parse("import ast as test", "ast.py").unwrap();
        let options = PythonOptions::default();
        info!("{:?}", result);

        let code = result.to_rust(
            CodeGenContext::Module("test_case".to_string()),
            options,
            SymbolTableScopes::new(),
        );
        info!("module: {:?}", code);
    }
}