python-ast 1.0.2

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

use log::info;
use proc_macro2::TokenStream;
use pyo3::{Bound, 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> FromPyObject<'a> for Type {
    fn extract_bound(ob: &Bound<'a, 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, FromPyObject, Serialize, Deserialize)]
pub struct RawModule {
    pub body: Vec<Statement>,
    pub type_ignores: Vec<Type>,
}

/// 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> FromPyObject<'a> for Module {
    fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult<Self> {
        let raw_module = ob.extract().expect("Failed parsing module.");

        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();
        
        // 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: {}", 
                    self.filename.unwrap_or_else(|| "unknown.py".to_string()));
                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::*;));
        }
        
        // 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;
        
        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);
                    
                    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())
                                .expect("parsing if __name__ body statement");
                            if !stmt_token.to_string().trim().is_empty() {
                                main_body_stmts.push(stmt_token);
                                has_main_code = true;
                            }
                        }
                    }
                    // Skip generating this if statement - we've processed its contents
                    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())
                .expect(format!("parsing statement {:?} in module", s).as_str());
            
            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);
                    has_module_init_code = true;
                }
            }
        }
        
        // Generate module initialization function if needed
        if has_module_init_code {
            stream.extend(quote! {
                fn __module_init__() {
                    #(#module_init_stmts)*
                }
            });
        }
        
        // 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() {
                                __module_init__();
                                python_main();
                            }
                        });
                    }
                } 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() {
                                __module_init__();
                                python_main();
                            }
                        });
                    }
                }
            } 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
                    
                    if has_module_init_code {
                        stream.extend(quote! {
                            #[#attr_tokens]
                            async fn main() {
                                __module_init__();
                                #(#main_body_stmts)*
                            }
                        });
                    } else {
                        stream.extend(quote! {
                            #[#attr_tokens]
                            async fn main() {
                                #(#main_body_stmts)*
                            }
                        });
                    }
                } else {
                    if has_module_init_code {
                        stream.extend(quote! {
                            fn main() {
                                __module_init__();
                                #(#main_body_stmts)*
                            }
                        });
                    } else {
                        stream.extend(quote! {
                            fn main() {
                                #(#main_body_stmts)*
                            }
                        });
                    }
                }
            }
        } 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() {
                    __module_init__();
                }
            });
        }
        Ok(stream)
    }
}

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(_) | 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);
    }
}