dissolve_python/

stub_collector.rs

// Implementation of RuffDeprecatedFunctionCollector using rustpython-parser
use crate::core::types::*;
use anyhow::Result;
use rustpython_ast::{self as ast};
use rustpython_parser::{parse, Mode};
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};

pub struct RuffDeprecatedFunctionCollector {
    module_name: String,
    _file_path: Option<PathBuf>,
    replacements: HashMap<String, ReplaceInfo>,
    unreplaceable: HashMap<String, UnreplaceableNode>,
    imports: Vec<ImportInfo>,
    inheritance_map: HashMap<String, Vec<String>>,
    class_methods: HashMap<String, HashSet<String>>,
    class_stack: Vec<String>,
    source: String,
    builtins: HashSet<String>,
    local_classes: HashSet<String>,
    local_functions: HashSet<String>,
}

impl RuffDeprecatedFunctionCollector {
    pub fn new(module_name: String, file_path: Option<&Path>) -> Self {
        Self {
            module_name,
            _file_path: file_path.map(Path::to_path_buf),
            replacements: HashMap::new(),
            unreplaceable: HashMap::new(),
            imports: Vec::new(),
            inheritance_map: HashMap::new(),
            class_methods: HashMap::new(),
            class_stack: Vec::new(),
            source: String::new(),
            builtins: Self::get_all_builtins(),
            local_classes: HashSet::new(),
            local_functions: HashSet::new(),
        }
    }

    /// Collect from source string
    pub fn collect_from_source(mut self, source: String) -> Result<CollectorResult> {
        self.source = source;
        let parsed = parse(&self.source, Mode::Module, "<module>")?;

        match parsed {
            ast::Mod::Module(module) => {
                for stmt in &module.body {
                    self.visit_stmt(stmt);
                }
            }
            ast::Mod::Expression(_) => {
                // Not handling expression mode
            }
            _ => {
                // Not handling other modes
            }
        }

        Ok(CollectorResult {
            replacements: self.replacements,
            unreplaceable: self.unreplaceable,
            imports: self.imports,
            inheritance_map: self.inheritance_map,
            class_methods: self.class_methods,
        })
    }

    /// Build the full object path including module and class names
    fn build_full_path(&self, name: &str) -> String {
        let mut parts = Vec::with_capacity(2 + self.class_stack.len());
        parts.push(self.module_name.as_str());
        parts.extend(self.class_stack.iter().map(|s| s.as_str()));
        parts.push(name);
        parts.join(".")
    }

    /// Build a qualified name from an expression (e.g., module.Class)
    fn build_qualified_name_from_expr(&self, expr: &ast::Expr) -> String {
        match expr {
            ast::Expr::Name(name) => {
                // Simple name - assume it's in the current module
                format!("{}.{}", self.module_name, name.id)
            }
            ast::Expr::Attribute(attr) => {
                // Handle nested attributes like a.b.c
                let mut result = attr.attr.to_string();
                let mut current = &*attr.value;

                loop {
                    match current {
                        ast::Expr::Name(name) => {
                            result = format!("{}.{}", name.id, result);
                            break;
                        }
                        ast::Expr::Attribute(inner_attr) => {
                            result = format!("{}.{}", inner_attr.attr, result);
                            current = &*inner_attr.value;
                        }
                        _ => {
                            return attr.attr.to_string();
                        }
                    }
                }

                result
            }
            _ => "Unknown".to_string(),
        }
    }

    /// Check if a decorator list contains @replace_me
    fn has_replace_me_decorator(decorators: &[ast::Expr]) -> bool {
        decorators.iter().any(|dec| match dec {
            ast::Expr::Name(name) => name.id.as_str() == "replace_me",
            ast::Expr::Call(call) => {
                matches!(&*call.func, ast::Expr::Name(name) if name.id.as_str() == "replace_me")
            }
            _ => false,
        })
    }

    /// Extract the 'since' version from @replace_me decorator
    fn extract_since_version(&self, decorators: &[ast::Expr]) -> Option<String> {
        self.extract_decorator_version_arg(decorators, "since")
    }

    fn extract_remove_in_version(&self, decorators: &[ast::Expr]) -> Option<String> {
        self.extract_decorator_version_arg(decorators, "remove_in")
    }

    fn extract_message(decorators: &[ast::Expr]) -> Option<String> {
        Self::extract_decorator_string_arg(decorators, "message")
    }

    fn extract_decorator_version_arg(
        &self,
        decorators: &[ast::Expr],
        arg_name: &str,
    ) -> Option<String> {
        for dec in decorators {
            if let ast::Expr::Call(call) = dec {
                if matches!(&*call.func, ast::Expr::Name(name) if name.id.as_str() == "replace_me")
                {
                    for keyword in &call.keywords {
                        if let Some(arg) = &keyword.arg {
                            if arg.as_str() == arg_name {
                                match &keyword.value {
                                    // String literal: "1.2.3"
                                    ast::Expr::Constant(c) => {
                                        if let ast::Constant::Str(s) = &c.value {
                                            return Some(s.to_string());
                                        }
                                    }
                                    // Tuple literal: (1, 2, 3)
                                    ast::Expr::Tuple(tuple) => {
                                        let parts: Vec<String> = tuple
                                            .elts
                                            .iter()
                                            .filter_map(|elt| match elt {
                                                ast::Expr::Constant(c) => match &c.value {
                                                    ast::Constant::Int(i) => Some(i.to_string()),
                                                    ast::Constant::Str(s) => Some(s.to_string()),
                                                    _ => None,
                                                },
                                                _ => None,
                                            })
                                            .collect();
                                        if !parts.is_empty() {
                                            return Some(parts.join("."));
                                        }
                                    }
                                    _ => {}
                                }
                            }
                        }
                    }
                }
            }
        }
        None
    }

    fn extract_decorator_string_arg(decorators: &[ast::Expr], arg_name: &str) -> Option<String> {
        for dec in decorators {
            if let ast::Expr::Call(call) = dec {
                if matches!(&*call.func, ast::Expr::Name(name) if name.id.as_str() == "replace_me")
                {
                    for keyword in &call.keywords {
                        if let Some(arg) = &keyword.arg {
                            if arg.as_str() == arg_name {
                                if let ast::Expr::Constant(c) = &keyword.value {
                                    if let ast::Constant::Str(s) = &c.value {
                                        return Some(s.to_string());
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        None
    }

    /// Extract replacement expression from @replace_me decorator arguments
    fn extract_replacement_from_decorator(
        &self,
        decorators: &[ast::Expr],
    ) -> Option<(String, ast::Expr)> {
        for dec in decorators {
            if let ast::Expr::Call(call) = dec {
                if matches!(&*call.func, ast::Expr::Name(name) if name.id.as_str() == "replace_me")
                {
                    for keyword in &call.keywords {
                        if let Some(arg) = &keyword.arg {
                            if arg.as_str() == "replacement" {
                                let replacement_expr = self.expr_to_string(&keyword.value);
                                return Some((replacement_expr, keyword.value.clone()));
                            }
                        }
                    }
                }
            }
        }
        None
    }

    /// Extract parameters from a function
    fn extract_parameters(&self, func: &ast::StmtFunctionDef) -> Vec<ParameterInfo> {
        let mut params = Vec::new();

        // Regular parameters
        for arg in &func.args.args {
            let has_default = arg.default.is_some();
            let default_value = arg.default.as_ref().map(|e| self.expr_to_string(e));

            let mut param_info = ParameterInfo::new(&arg.def.arg);
            param_info.has_default = has_default;
            param_info.default_value = default_value;
            params.push(param_info);
        }

        // *args
        if let Some(vararg) = &func.args.vararg {
            params.push(ParameterInfo::vararg(&vararg.arg));
        }

        // Keyword-only parameters
        for arg in &func.args.kwonlyargs {
            let has_default = arg.default.is_some();
            let default_value = arg.default.as_ref().map(|e| self.expr_to_string(e));

            let mut param_info = ParameterInfo::new(&arg.def.arg);
            param_info.has_default = has_default;
            param_info.default_value = default_value;
            param_info.is_kwonly = true;
            params.push(param_info);
        }

        // **kwargs
        if let Some(kwarg) = &func.args.kwarg {
            params.push(ParameterInfo::kwarg(&kwarg.arg));
        }

        params
    }

    /// Extract parameters from an async function
    fn extract_async_parameters(&self, func: &ast::StmtAsyncFunctionDef) -> Vec<ParameterInfo> {
        let mut params = Vec::new();

        // Regular parameters
        for arg in &func.args.args {
            let has_default = arg.default.is_some();
            let default_value = arg.default.as_ref().map(|e| self.expr_to_string(e));
            let mut param_info = ParameterInfo::new(&arg.def.arg);
            param_info.has_default = has_default;
            param_info.default_value = default_value;
            params.push(param_info);
        }

        // *args
        if let Some(vararg) = &func.args.vararg {
            params.push(ParameterInfo::vararg(&vararg.arg));
        }

        // Keyword-only parameters
        for arg in &func.args.kwonlyargs {
            let has_default = arg.default.is_some();
            let default_value = arg.default.as_ref().map(|e| self.expr_to_string(e));
            let mut param_info = ParameterInfo::new(&arg.def.arg);
            param_info.has_default = has_default;
            param_info.default_value = default_value;
            param_info.is_kwonly = true;
            params.push(param_info);
        }

        // **kwargs
        if let Some(kwarg) = &func.args.kwarg {
            params.push(ParameterInfo::kwarg(&kwarg.arg));
        }

        params
    }

    /// Extract replacement expression from function body
    fn extract_replacement_from_function(
        &self,
        func: &ast::StmtFunctionDef,
    ) -> Result<(String, ast::Expr), ReplacementExtractionError> {
        // Skip docstring and pass statements
        let body_stmts: Vec<&ast::Stmt> = func
            .body
            .iter()
            .skip_while(|stmt| {
                matches!(stmt, ast::Stmt::Expr(expr_stmt) if matches!(expr_stmt.value.as_ref(),
                    ast::Expr::Constant(c) if matches!(&c.value, ast::Constant::Str(_))))
            })
            .filter(|stmt| !matches!(stmt, ast::Stmt::Pass(_)))
            .collect();

        if body_stmts.is_empty() {
            // Empty body is valid - means remove the function completely
            return Ok((
                "".to_string(),
                ast::Expr::Constant(ast::ExprConstant {
                    value: ast::Constant::Str("".to_string()),
                    kind: None,
                    range: Default::default(),
                }),
            ));
        }

        if body_stmts.len() > 1 {
            return Err(ReplacementExtractionError::new(
                func.name.to_string(),
                ReplacementFailureReason::MultipleStatements,
                "Function body contains multiple statements".to_string(),
            ));
        }

        // Extract return expression
        match body_stmts[0] {
            ast::Stmt::Return(ret_stmt) => {
                if let Some(value) = &ret_stmt.value {
                    // Get function parameters for placeholder conversion
                    let param_names: HashSet<String> = self
                        .extract_parameters(func)
                        .into_iter()
                        .filter(|p| !p.is_vararg && !p.is_kwarg)
                        .map(|p| p.name)
                        .collect();

                    // Convert the AST expression to string with placeholders
                    let replacement_expr =
                        self.expr_to_string_with_placeholders(value, &param_names);

                    tracing::debug!(
                        "Extracted replacement expression: {} for function {}",
                        replacement_expr,
                        func.name
                    );

                    Ok((replacement_expr, (**value).clone()))
                } else {
                    Err(ReplacementExtractionError::new(
                        func.name.to_string(),
                        ReplacementFailureReason::NoReturnStatement,
                        "Return statement has no value".to_string(),
                    ))
                }
            }
            _ => Err(ReplacementExtractionError::new(
                func.name.to_string(),
                ReplacementFailureReason::NoReturnStatement,
                "Function body does not contain a return statement".to_string(),
            )),
        }
    }

    /// Get all builtin names from Python
    fn get_all_builtins() -> HashSet<String> {
        use pyo3::prelude::*;

        Python::with_gil(|py| {
            let mut builtin_names = HashSet::new();

            if let Ok(builtins) = py.import("builtins") {
                if let Ok(dir_result) = builtins.dir() {
                    for item in dir_result.iter() {
                        if let Ok(name_str) = item.extract::<String>() {
                            builtin_names.insert(name_str);
                        }
                    }
                }
            }

            builtin_names
        })
    }

    /// Check if a name is imported from another module
    fn is_imported(&self, name: &str) -> bool {
        for import_info in &self.imports {
            for (imported_name, alias) in &import_info.names {
                let effective_name = alias.as_ref().unwrap_or(imported_name);
                if effective_name == name {
                    return true;
                }
            }
        }
        false
    }

    fn is_locally_available(&self, name: &str) -> bool {
        self.is_imported(name)
            || self.local_classes.contains(name)
            || self.local_functions.contains(name)
    }

    /// Get the builtins set
    pub fn builtins(&self) -> &HashSet<String> {
        &self.builtins
    }

    /// Convert expression to string
    fn expr_to_string(&self, expr: &ast::Expr) -> String {
        // Use the full implementation to avoid "..." fallbacks
        self.expr_to_string_with_placeholders(expr, &HashSet::new())
    }

    fn expr_to_string_with_placeholders(
        &self,
        expr: &ast::Expr,
        param_names: &HashSet<String>,
    ) -> String {
        match expr {
            ast::Expr::Name(name) => {
                let name_str = name.id.to_string();
                if param_names.contains(&name_str) {
                    format!("{{{}}}", name_str)
                } else {
                    name_str
                }
            }
            ast::Expr::Call(call) => {
                let func_str = match &*call.func {
                    ast::Expr::Name(name) => {
                        let name_str = name.id.to_string();
                        if param_names.contains(&name_str) {
                            format!("{{{}}}", name_str)
                        } else {
                            // For class replacements, add module prefix if the called name
                            // is not imported and not locally defined and looks like a class (starts with uppercase)
                            if name_str.chars().next().is_some_and(|c| c.is_uppercase()) {
                                if !self.is_locally_available(&name_str) {
                                    format!("{}.{}", self.module_name, name_str)
                                } else {
                                    name_str
                                }
                            } else {
                                // Don't add module prefix for lowercase names - let the migration handle it
                                name_str
                            }
                        }
                    }
                    _ => self.expr_to_string_with_placeholders(&call.func, param_names),
                };
                let mut args = Vec::new();

                // Handle positional arguments
                for arg in &call.args {
                    args.push(self.expr_to_string_with_placeholders(arg, param_names));
                }

                // Handle keyword arguments
                for keyword in &call.keywords {
                    if let Some(arg_name) = &keyword.arg {
                        let value_str =
                            self.expr_to_string_with_placeholders(&keyword.value, param_names);
                        args.push(format!("{}={}", arg_name, value_str));
                    } else {
                        // **kwargs expansion
                        let inner =
                            self.expr_to_string_with_placeholders(&keyword.value, param_names);
                        // Check if this is **kwargs
                        if let ast::Expr::Name(name) = &keyword.value {
                            if name.id.to_string() == "kwargs" {
                                args.push("{**kwargs}".to_string());
                                continue;
                            }
                        }
                        args.push(format!("**{}", inner));
                    }
                }

                format!("{}({})", func_str, args.join(", "))
            }
            ast::Expr::Attribute(attr) => {
                let value_str = self.expr_to_string_with_placeholders(&attr.value, param_names);
                format!("{}.{}", value_str, attr.attr)
            }
            ast::Expr::Starred(starred) => {
                let inner = self.expr_to_string_with_placeholders(&starred.value, param_names);
                // Check if this is a varargs parameter (*args)
                if let ast::Expr::Name(name) = starred.value.as_ref() {
                    if name.id.to_string() == "args" {
                        return "{*args}".to_string();
                    }
                }
                format!("*{}", inner)
            }
            ast::Expr::BinOp(binop) => {
                let left = self.expr_to_string_with_placeholders(&binop.left, param_names);
                let right = self.expr_to_string_with_placeholders(&binop.right, param_names);
                let op_str = match binop.op {
                    ast::Operator::Add => " + ",
                    ast::Operator::Sub => " - ",
                    ast::Operator::Mult => " * ",
                    ast::Operator::Div => " / ",
                    ast::Operator::Mod => " % ",
                    ast::Operator::Pow => " ** ",
                    ast::Operator::LShift => " << ",
                    ast::Operator::RShift => " >> ",
                    ast::Operator::BitOr => " | ",
                    ast::Operator::BitXor => " ^ ",
                    ast::Operator::BitAnd => " & ",
                    ast::Operator::FloorDiv => " // ",
                    ast::Operator::MatMult => " @ ",
                };
                format!("{}{}{}", left, op_str, right)
            }
            ast::Expr::Compare(compare) => {
                let mut result = self.expr_to_string_with_placeholders(&compare.left, param_names);
                for (op, comparator) in compare.ops.iter().zip(&compare.comparators) {
                    let op_str = match op {
                        ast::CmpOp::Eq => " == ",
                        ast::CmpOp::NotEq => " != ",
                        ast::CmpOp::Lt => " < ",
                        ast::CmpOp::LtE => " <= ",
                        ast::CmpOp::Gt => " > ",
                        ast::CmpOp::GtE => " >= ",
                        ast::CmpOp::Is => " is ",
                        ast::CmpOp::IsNot => " is not ",
                        ast::CmpOp::In => " in ",
                        ast::CmpOp::NotIn => " not in ",
                    };
                    result.push_str(op_str);
                    result
                        .push_str(&self.expr_to_string_with_placeholders(comparator, param_names));
                }
                result
            }
            ast::Expr::BoolOp(boolop) => {
                let op_str = match boolop.op {
                    ast::BoolOp::And => " and ",
                    ast::BoolOp::Or => " or ",
                };
                boolop
                    .values
                    .iter()
                    .map(|v| self.expr_to_string_with_placeholders(v, param_names))
                    .collect::<Vec<_>>()
                    .join(op_str)
            }
            ast::Expr::UnaryOp(unaryop) => {
                let operand = self.expr_to_string_with_placeholders(&unaryop.operand, param_names);
                let op_str = match unaryop.op {
                    ast::UnaryOp::Invert => "~",
                    ast::UnaryOp::Not => "not ",
                    ast::UnaryOp::UAdd => "+",
                    ast::UnaryOp::USub => "-",
                };
                format!("{}{}", op_str, operand)
            }
            ast::Expr::Await(await_expr) => {
                self.expr_to_string_with_placeholders(&await_expr.value, param_names)
            }
            ast::Expr::Subscript(subscript) => {
                let value_str =
                    self.expr_to_string_with_placeholders(&subscript.value, param_names);
                let slice_str =
                    self.expr_to_string_with_placeholders(&subscript.slice, param_names);
                format!("{}[{}]", value_str, slice_str)
            }
            ast::Expr::Slice(slice) => {
                let lower = slice
                    .lower
                    .as_ref()
                    .map(|e| self.expr_to_string_with_placeholders(e, param_names))
                    .unwrap_or_default();
                let upper = slice
                    .upper
                    .as_ref()
                    .map(|e| self.expr_to_string_with_placeholders(e, param_names))
                    .unwrap_or_default();
                let step = slice
                    .step
                    .as_ref()
                    .map(|e| format!(":{}", self.expr_to_string_with_placeholders(e, param_names)))
                    .unwrap_or_default();
                format!("{}:{}{}", lower, upper, step)
            }
            ast::Expr::Constant(c) => match &c.value {
                ast::Constant::Str(s) => {
                    // Use proper string literal escaping
                    format!(
                        "\"{}\"",
                        s.chars()
                            .map(|c| match c {
                                '"' => "\\\"".to_string(),
                                '\\' => "\\\\".to_string(),
                                '\n' => "\\n".to_string(),
                                '\r' => "\\r".to_string(),
                                '\t' => "\\t".to_string(),
                                c if c.is_control() => format!("\\x{:02x}", c as u8),
                                c => c.to_string(),
                            })
                            .collect::<String>()
                    )
                }
                ast::Constant::Int(i) => i.to_string(),
                ast::Constant::Float(f) => f.to_string(),
                ast::Constant::Bool(b) => if *b { "True" } else { "False" }.to_string(),
                ast::Constant::None => "None".to_string(),
                ast::Constant::Bytes(b) => {
                    // Properly escape bytes literal
                    let escaped = b
                        .iter()
                        .map(|&byte| match byte {
                            b'"' => "\\\"".to_string(),
                            b'\\' => "\\\\".to_string(),
                            b'\n' => "\\n".to_string(),
                            b'\r' => "\\r".to_string(),
                            b'\t' => "\\t".to_string(),
                            b'\0' => "\\x00".to_string(),
                            b if b.is_ascii_graphic() || b == b' ' => (b as char).to_string(),
                            b => format!("\\x{:02x}", b),
                        })
                        .collect::<String>();
                    format!("b\"{}\"", escaped)
                }
                ast::Constant::Complex { real, imag } => {
                    if *real == 0.0 {
                        format!("{}j", imag)
                    } else if *imag >= 0.0 {
                        format!("{} + {}j", real, imag)
                    } else {
                        format!("{} - {}j", real, -imag)
                    }
                }
                ast::Constant::Ellipsis => "...".to_string(),
                _ => "...".to_string(),
            },
            ast::Expr::ListComp(comp) => {
                let elt = self.expr_to_string_with_placeholders(&comp.elt, param_names);

                // Handle all generators (for multiple 'for' clauses)
                let mut generators = Vec::new();
                for gen in &comp.generators {
                    let target = self.expr_to_string_with_placeholders(&gen.target, param_names);
                    let iter = self.expr_to_string_with_placeholders(&gen.iter, param_names);
                    let mut gen_str = format!("for {} in {}", target, iter);

                    // Add conditions for this generator
                    if !gen.ifs.is_empty() {
                        let conds: Vec<String> = gen
                            .ifs
                            .iter()
                            .map(|if_expr| {
                                self.expr_to_string_with_placeholders(if_expr, param_names)
                            })
                            .collect();
                        gen_str.push_str(&format!(" if {}", conds.join(" if ")));
                    }

                    generators.push(gen_str);
                }

                format!("[{} {}]", elt, generators.join(" "))
            }
            ast::Expr::SetComp(comp) => {
                let elt = self.expr_to_string_with_placeholders(&comp.elt, param_names);

                // Handle all generators (for multiple 'for' clauses)
                let mut generators = Vec::new();
                for gen in &comp.generators {
                    let target = self.expr_to_string_with_placeholders(&gen.target, param_names);
                    let iter = self.expr_to_string_with_placeholders(&gen.iter, param_names);
                    let mut gen_str = format!("for {} in {}", target, iter);

                    // Add conditions for this generator
                    if !gen.ifs.is_empty() {
                        let conds: Vec<String> = gen
                            .ifs
                            .iter()
                            .map(|if_expr| {
                                self.expr_to_string_with_placeholders(if_expr, param_names)
                            })
                            .collect();
                        gen_str.push_str(&format!(" if {}", conds.join(" if ")));
                    }

                    generators.push(gen_str);
                }

                format!("{{{} {}}}", elt, generators.join(" "))
            }
            ast::Expr::DictComp(comp) => {
                let key = self.expr_to_string_with_placeholders(&comp.key, param_names);
                let value = self.expr_to_string_with_placeholders(&comp.value, param_names);

                // Handle all generators (for multiple 'for' clauses)
                let mut generators = Vec::new();
                for gen in &comp.generators {
                    let target = self.expr_to_string_with_placeholders(&gen.target, param_names);
                    let iter = self.expr_to_string_with_placeholders(&gen.iter, param_names);
                    let mut gen_str = format!("for {} in {}", target, iter);

                    // Add conditions for this generator
                    if !gen.ifs.is_empty() {
                        let conds: Vec<String> = gen
                            .ifs
                            .iter()
                            .map(|if_expr| {
                                self.expr_to_string_with_placeholders(if_expr, param_names)
                            })
                            .collect();
                        gen_str.push_str(&format!(" if {}", conds.join(" if ")));
                    }

                    generators.push(gen_str);
                }

                format!("{{{}: {} {}}}", key, value, generators.join(" "))
            }
            ast::Expr::GeneratorExp(gen) => {
                let elt = self.expr_to_string_with_placeholders(&gen.elt, param_names);

                // Handle all generators (for multiple 'for' clauses)
                let mut generators = Vec::new();
                for gen_item in &gen.generators {
                    let target =
                        self.expr_to_string_with_placeholders(&gen_item.target, param_names);
                    let iter = self.expr_to_string_with_placeholders(&gen_item.iter, param_names);
                    let mut gen_str = format!("for {} in {}", target, iter);

                    // Add conditions for this generator
                    if !gen_item.ifs.is_empty() {
                        let conds: Vec<String> = gen_item
                            .ifs
                            .iter()
                            .map(|if_expr| {
                                self.expr_to_string_with_placeholders(if_expr, param_names)
                            })
                            .collect();
                        gen_str.push_str(&format!(" if {}", conds.join(" if ")));
                    }

                    generators.push(gen_str);
                }

                format!("({} {})", elt, generators.join(" "))
            }
            ast::Expr::List(list) => {
                let elements: Vec<String> = list
                    .elts
                    .iter()
                    .map(|e| self.expr_to_string_with_placeholders(e, param_names))
                    .collect();
                format!("[{}]", elements.join(", "))
            }
            ast::Expr::Tuple(tuple) => {
                let elements: Vec<String> = tuple
                    .elts
                    .iter()
                    .map(|e| self.expr_to_string_with_placeholders(e, param_names))
                    .collect();
                if elements.len() == 1 {
                    format!("({},)", elements[0])
                } else {
                    format!("({})", elements.join(", "))
                }
            }
            ast::Expr::Dict(dict) => {
                let mut items = Vec::new();
                for (key, value) in dict.keys.iter().zip(&dict.values) {
                    if let Some(k) = key {
                        items.push(format!(
                            "{}: {}",
                            self.expr_to_string_with_placeholders(k, param_names),
                            self.expr_to_string_with_placeholders(value, param_names)
                        ));
                    } else {
                        items.push(format!(
                            "**{}",
                            self.expr_to_string_with_placeholders(value, param_names)
                        ));
                    }
                }
                format!("{{{}}}", items.join(", "))
            }
            _ => {
                // For unsupported expression types, try to provide a reasonable fallback
                match expr {
                    ast::Expr::IfExp(ifexp) => {
                        let body = self.expr_to_string_with_placeholders(&ifexp.body, param_names);
                        let test = self.expr_to_string_with_placeholders(&ifexp.test, param_names);
                        let orelse =
                            self.expr_to_string_with_placeholders(&ifexp.orelse, param_names);
                        format!("{} if {} else {}", body, test, orelse)
                    }
                    ast::Expr::Lambda(lambda) => {
                        let args: Vec<String> = lambda
                            .args
                            .args
                            .iter()
                            .map(|arg| arg.def.arg.to_string())
                            .collect();
                        let body = self.expr_to_string_with_placeholders(&lambda.body, param_names);
                        format!("lambda {}: {}", args.join(", "), body)
                    }
                    ast::Expr::JoinedStr(joined) => {
                        let mut result = String::from("f\"");
                        for value in &joined.values {
                            match value {
                                ast::Expr::Constant(c) => {
                                    if let ast::Constant::Str(s) = &c.value {
                                        result.push_str(s);
                                    }
                                }
                                ast::Expr::FormattedValue(fmt) => {
                                    result.push('{');
                                    let var_str = self
                                        .expr_to_string_with_placeholders(&fmt.value, param_names);
                                    result.push_str(&var_str);
                                    match fmt.conversion {
                                        ast::ConversionFlag::Str => result.push_str("!s"),
                                        ast::ConversionFlag::Repr => result.push_str("!r"),
                                        ast::ConversionFlag::Ascii => result.push_str("!a"),
                                        ast::ConversionFlag::None => {}
                                    }
                                    if let Some(spec) = &fmt.format_spec {
                                        result.push(':');
                                        match &**spec {
                                            ast::Expr::Constant(c) => {
                                                if let ast::Constant::Str(s) = &c.value {
                                                    result.push_str(s);
                                                }
                                            }
                                            ast::Expr::JoinedStr(_) => {
                                                // For f-strings inside format specs, we need special handling
                                                let reconstructed = self
                                                    .expr_to_string_with_placeholders(
                                                        spec,
                                                        param_names,
                                                    );
                                                // Remove the f" prefix and " suffix if it's an f-string
                                                if reconstructed.starts_with("f\"")
                                                    && reconstructed.ends_with('"')
                                                {
                                                    result.push_str(
                                                        &reconstructed[2..reconstructed.len() - 1],
                                                    );
                                                } else {
                                                    result.push_str(&reconstructed);
                                                }
                                            }
                                            _ => result.push_str(
                                                &self.expr_to_string_with_placeholders(
                                                    spec,
                                                    param_names,
                                                ),
                                            ),
                                        }
                                    }
                                    result.push('}');
                                }
                                _ => {
                                    result.push('{');
                                    let other_str =
                                        self.expr_to_string_with_placeholders(value, param_names);
                                    result.push_str(&other_str);
                                    result.push('}');
                                }
                            }
                        }
                        result.push('"');
                        result
                    }
                    ast::Expr::FormattedValue(fmt) => {
                        // This should normally be handled within JoinedStr
                        let mut result = String::new();
                        result.push('{');
                        result.push_str(
                            &self.expr_to_string_with_placeholders(&fmt.value, param_names),
                        );
                        match fmt.conversion {
                            ast::ConversionFlag::Str => result.push_str("!s"),
                            ast::ConversionFlag::Repr => result.push_str("!r"),
                            ast::ConversionFlag::Ascii => result.push_str("!a"),
                            ast::ConversionFlag::None => {}
                        }
                        if let Some(spec) = &fmt.format_spec {
                            result.push(':');
                            match &**spec {
                                ast::Expr::Constant(c) => {
                                    if let ast::Constant::Str(s) = &c.value {
                                        result.push_str(s);
                                    }
                                }
                                _ => result.push_str(
                                    &self.expr_to_string_with_placeholders(spec, param_names),
                                ),
                            }
                        }
                        result.push('}');
                        result
                    }
                    ast::Expr::NamedExpr(named) => {
                        format!(
                            "{} := {}",
                            self.expr_to_string_with_placeholders(&named.target, param_names),
                            self.expr_to_string_with_placeholders(&named.value, param_names)
                        )
                    }
                    ast::Expr::Set(set) => {
                        let elements: Vec<String> = set
                            .elts
                            .iter()
                            .map(|e| self.expr_to_string_with_placeholders(e, param_names))
                            .collect();
                        if elements.is_empty() {
                            "set()".to_string()
                        } else {
                            format!("{{{}}}", elements.join(", "))
                        }
                    }
                    ast::Expr::Yield(yield_expr) => {
                        if let Some(value) = &yield_expr.value {
                            format!(
                                "yield {}",
                                self.expr_to_string_with_placeholders(value, param_names)
                            )
                        } else {
                            "yield".to_string()
                        }
                    }
                    ast::Expr::YieldFrom(yieldfrom) => {
                        format!(
                            "yield from {}",
                            self.expr_to_string_with_placeholders(&yieldfrom.value, param_names)
                        )
                    }
                    _ => "/* unsupported expr */".to_string(),
                }
            }
        }
    }

    /// Extract replacement expression from async function body
    fn extract_replacement_from_async_function(
        &self,
        func: &ast::StmtAsyncFunctionDef,
    ) -> Result<(String, ast::Expr), ReplacementExtractionError> {
        // Skip docstring and pass statements
        let body_stmts: Vec<&ast::Stmt> = func
            .body
            .iter()
            .skip_while(|stmt| {
                matches!(stmt, ast::Stmt::Expr(expr_stmt) if matches!(expr_stmt.value.as_ref(),
                    ast::Expr::Constant(c) if matches!(&c.value, ast::Constant::Str(_))))
            })
            .filter(|stmt| !matches!(stmt, ast::Stmt::Pass(_)))
            .collect();

        if body_stmts.is_empty() {
            // Empty body is valid - means remove the function completely
            return Ok((
                "".to_string(),
                ast::Expr::Constant(ast::ExprConstant {
                    value: ast::Constant::Str("".to_string()),
                    kind: None,
                    range: Default::default(),
                }),
            ));
        }

        if body_stmts.len() > 1 {
            return Err(ReplacementExtractionError::new(
                func.name.to_string(),
                ReplacementFailureReason::MultipleStatements,
                "Function body contains multiple statements".to_string(),
            ));
        }

        // Extract return expression
        match body_stmts[0] {
            ast::Stmt::Return(ret_stmt) => {
                if let Some(value) = &ret_stmt.value {
                    // Get function parameters for placeholder conversion
                    let param_names: HashSet<String> = func
                        .args
                        .args
                        .iter()
                        .map(|arg| arg.def.arg.to_string())
                        .collect();

                    let replacement_expr =
                        self.expr_to_string_with_placeholders(value, &param_names);
                    Ok((replacement_expr, value.as_ref().clone()))
                } else {
                    Err(ReplacementExtractionError::new(
                        func.name.to_string(),
                        ReplacementFailureReason::NoReturnStatement,
                        "Return statement has no value".to_string(),
                    ))
                }
            }
            _ => Err(ReplacementExtractionError::new(
                func.name.to_string(),
                ReplacementFailureReason::NoReturnStatement,
                "Function body must contain only a return statement".to_string(),
            )),
        }
    }

    /// Visit a function definition
    fn visit_function(&mut self, func: &ast::StmtFunctionDef) {
        // Track this as a locally defined function (even if not decorated)
        if self.class_stack.is_empty() {
            // Only track top-level functions, not methods
            self.local_functions.insert(func.name.to_string());
        }

        if !Self::has_replace_me_decorator(&func.decorator_list) {
            return;
        }

        let full_path = self.build_full_path(&func.name);
        let parameters = self.extract_parameters(func);
        let since = self.extract_since_version(&func.decorator_list);
        let remove_in = self.extract_remove_in_version(&func.decorator_list);
        let message = Self::extract_message(&func.decorator_list);

        // Determine construct type
        let construct_type = if self.class_stack.is_empty() {
            ConstructType::Function
        } else {
            let decorator_names: Vec<&str> = func
                .decorator_list
                .iter()
                .filter_map(|dec| {
                    if let ast::Expr::Name(name) = dec {
                        Some(name.id.as_str())
                    } else {
                        None
                    }
                })
                .collect();

            if decorator_names.contains(&"property") {
                ConstructType::Property
            } else if decorator_names.contains(&"classmethod") {
                ConstructType::ClassMethod
            } else if decorator_names.contains(&"staticmethod") {
                ConstructType::StaticMethod
            } else {
                ConstructType::Function
            }
        };

        // Try to extract replacement from decorator first, then from function body
        let replacement_result = if let Some((replacement_expr, ast)) =
            self.extract_replacement_from_decorator(&func.decorator_list)
        {
            Ok((replacement_expr, ast))
        } else {
            self.extract_replacement_from_function(func)
        };

        match replacement_result {
            Ok((replacement_expr, ast)) => {
                let mut replace_info =
                    ReplaceInfo::new(&full_path, &replacement_expr, construct_type);
                // Store the AST
                replace_info.replacement_ast = Some(Box::new(ast));
                replace_info.parameters = parameters;
                replace_info.since = since.and_then(|s| s.parse().ok());
                replace_info.remove_in = remove_in.and_then(|s| s.parse().ok());
                replace_info.message = message;
                self.replacements.insert(full_path, replace_info);
            }
            Err(e) => {
                let unreplaceable = UnreplaceableNode::new(
                    full_path.clone(),
                    e.reason(),
                    e.to_string(),
                    construct_type,
                );
                self.unreplaceable.insert(full_path, unreplaceable);
            }
        }
    }

    /// Visit an async function definition
    fn visit_async_function(&mut self, func: &ast::StmtAsyncFunctionDef) {
        if !Self::has_replace_me_decorator(&func.decorator_list) {
            return;
        }

        let full_path = self.build_full_path(&func.name);
        let parameters = self.extract_async_parameters(func);
        let since = self.extract_since_version(&func.decorator_list);
        let remove_in = self.extract_remove_in_version(&func.decorator_list);
        let message = Self::extract_message(&func.decorator_list);

        // Determine construct type
        let construct_type = if self.class_stack.is_empty() {
            ConstructType::Function
        } else {
            let decorator_names: Vec<&str> = func
                .decorator_list
                .iter()
                .filter_map(|dec| {
                    if let ast::Expr::Name(name) = dec {
                        Some(name.id.as_str())
                    } else {
                        None
                    }
                })
                .collect();

            if decorator_names.contains(&"property") {
                ConstructType::Property
            } else if decorator_names.contains(&"classmethod") {
                ConstructType::ClassMethod
            } else if decorator_names.contains(&"staticmethod") {
                ConstructType::StaticMethod
            } else {
                ConstructType::Function
            }
        };

        // Try to extract replacement from decorator first, then from function body
        let replacement_result = if let Some((replacement_expr, ast)) =
            self.extract_replacement_from_decorator(&func.decorator_list)
        {
            Ok((replacement_expr, ast))
        } else {
            self.extract_replacement_from_async_function(func)
        };

        match replacement_result {
            Ok((replacement_expr, ast)) => {
                let mut replace_info =
                    ReplaceInfo::new(&full_path, &replacement_expr, construct_type);
                // Store the AST
                replace_info.replacement_ast = Some(Box::new(ast));
                replace_info.parameters = parameters;
                replace_info.since = since.and_then(|s| s.parse().ok());
                replace_info.remove_in = remove_in.and_then(|s| s.parse().ok());
                replace_info.message = message;
                self.replacements.insert(full_path, replace_info);
            }
            Err(e) => {
                let unreplaceable = UnreplaceableNode::new(
                    full_path.clone(),
                    e.reason(),
                    e.to_string(),
                    construct_type,
                );
                self.unreplaceable.insert(full_path, unreplaceable);
            }
        }
    }

    /// Visit a class definition
    fn visit_class(&mut self, class_def: &ast::StmtClassDef) {
        let class_name = class_def.name.to_string();
        let full_class_name = self.build_full_path(&class_name);

        // Track this as a locally defined class
        self.local_classes.insert(class_name.clone());

        // Record base classes
        let mut bases = Vec::new();
        for base in &class_def.bases {
            match base {
                ast::Expr::Name(name) => {
                    let base_name = name.id.to_string();
                    let qualified_name = format!("{}.{}", self.module_name, base_name);
                    bases.push(qualified_name);
                }
                ast::Expr::Attribute(_) => {
                    let qualified_name = self.build_qualified_name_from_expr(base);
                    bases.push(qualified_name);
                }
                _ => {}
            }
        }

        if !bases.is_empty() {
            tracing::debug!("Class {} inherits from: {:?}", full_class_name, bases);
            self.inheritance_map.insert(full_class_name.clone(), bases);
        }

        // Check if class itself has @replace_me
        if Self::has_replace_me_decorator(&class_def.decorator_list) {
            if let Some(init_replacement) = self.extract_class_replacement(class_def) {
                let mut replace_info =
                    ReplaceInfo::new(&full_class_name, &init_replacement, ConstructType::Class);

                // Extract __init__ parameters
                for stmt in &class_def.body {
                    if let ast::Stmt::FunctionDef(func) = stmt {
                        if func.name.as_str() == "__init__" {
                            replace_info.parameters = self
                                .extract_parameters(func)
                                .into_iter()
                                .filter(|p| p.name != "self")
                                .collect();
                            break;
                        }
                    }
                }

                self.replacements
                    .insert(full_class_name.clone(), replace_info);
            } else {
                let unreplaceable = UnreplaceableNode::new(
                    full_class_name.clone(),
                    ReplacementFailureReason::NoInitMethod,
                    "Class has @replace_me decorator but no __init__ method with clear replacement pattern".to_string(),
                    ConstructType::Class,
                );
                self.unreplaceable
                    .insert(full_class_name.clone(), unreplaceable);
            }
        }

        // Visit class body
        self.class_stack.push(class_name);
        for stmt in &class_def.body {
            self.visit_stmt(stmt);
        }
        self.class_stack.pop();
    }

    /// Extract replacement from class __init__ method
    fn extract_class_replacement(&self, class_def: &ast::StmtClassDef) -> Option<String> {
        for stmt in &class_def.body {
            if let ast::Stmt::FunctionDef(func) = stmt {
                if func.name.as_str() == "__init__" {
                    // Get parameter names from __init__ (excluding self)
                    let param_names: HashSet<String> = self
                        .extract_parameters(func)
                        .into_iter()
                        .filter(|p| p.name != "self" && !p.is_vararg && !p.is_kwarg)
                        .map(|p| p.name)
                        .collect();

                    for init_stmt in &func.body {
                        if let ast::Stmt::Assign(assign) = init_stmt {
                            if assign.targets.len() == 1 {
                                if let ast::Expr::Attribute(attr) = &assign.targets[0] {
                                    if let ast::Expr::Name(name) = &*attr.value {
                                        if name.id.as_str() == "self" {
                                            // Found self.attr = expr, use placeholders for parameters
                                            let replacement_expr = self
                                                .expr_to_string_with_placeholders(
                                                    &assign.value,
                                                    &param_names,
                                                );
                                            return Some(replacement_expr);
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
        None
    }

    fn visit_assign(&mut self, assign: &ast::StmtAssign) {
        // Handle module-level assignments like OLD_CONSTANT = replace_me(42)
        if assign.targets.len() == 1 {
            if let ast::Expr::Name(name) = &assign.targets[0] {
                // Check if the assignment value is a replace_me call
                if let ast::Expr::Call(call) = assign.value.as_ref() {
                    if matches!(&*call.func, ast::Expr::Name(func_name) if func_name.id.as_str() == "replace_me")
                    {
                        // Extract the replacement value
                        if let Some(arg) = call.args.first() {
                            let replacement_expr = self.expr_to_string(arg);

                            let full_name = if self.class_stack.is_empty() {
                                format!("{}.{}", self.module_name, name.id)
                            } else {
                                format!(
                                    "{}.{}.{}",
                                    self.module_name,
                                    self.class_stack.join("."),
                                    name.id
                                )
                            };

                            let construct_type = if self.class_stack.is_empty() {
                                ConstructType::ModuleAttribute
                            } else {
                                ConstructType::ClassAttribute
                            };

                            let replace_info = ReplaceInfo {
                                old_name: full_name.clone(),
                                replacement_expr,
                                replacement_ast: None, // rustpython AST types differ
                                construct_type,
                                parameters: vec![],
                                return_type: None,
                                since: None,
                                remove_in: None,
                                message: None,
                            };

                            self.replacements.insert(full_name, replace_info);
                        }
                    }
                }
            }
        }
    }

    fn visit_stmt(&mut self, stmt: &ast::Stmt) {
        match stmt {
            ast::Stmt::FunctionDef(func) => self.visit_function(func),
            ast::Stmt::AsyncFunctionDef(func) => self.visit_async_function(func),
            ast::Stmt::ClassDef(class) => self.visit_class(class),
            ast::Stmt::Import(import) => {
                for alias in &import.names {
                    self.imports.push(ImportInfo::new(
                        alias.name.to_string(),
                        vec![(
                            alias.name.to_string(),
                            alias.asname.as_ref().map(|n| n.to_string()),
                        )],
                    ));
                }
            }
            ast::Stmt::ImportFrom(import) => {
                let names: Vec<(String, Option<String>)> = import
                    .names
                    .iter()
                    .map(|alias| {
                        (
                            alias.name.to_string(),
                            alias.asname.as_ref().map(|n| n.to_string()),
                        )
                    })
                    .collect();

                let module_name = if let Some(module) = &import.module {
                    let level = import.level.as_ref().map_or(0, |i| i.to_usize());
                    let dots = ".".repeat(level);
                    format!("{}{}", dots, module)
                } else {
                    let level = import.level.as_ref().map_or(0, |i| i.to_usize());
                    ".".repeat(level)
                };

                self.imports.push(ImportInfo::new(module_name, names));
            }
            ast::Stmt::Assign(assign) => {
                self.visit_assign(assign);
            }
            ast::Stmt::AnnAssign(ann_assign) => {
                // Handle annotated assignments like DEFAULT_TIMEOUT: int = replace_me(30)
                if let Some(value) = &ann_assign.value {
                    if let ast::Expr::Name(name) = ann_assign.target.as_ref() {
                        if let ast::Expr::Call(call) = value.as_ref() {
                            if matches!(&*call.func, ast::Expr::Name(func_name) if func_name.id.as_str() == "replace_me")
                            {
                                if let Some(arg) = call.args.first() {
                                    let replacement_expr = self.expr_to_string(arg);

                                    let full_name = if self.class_stack.is_empty() {
                                        format!("{}.{}", self.module_name, name.id)
                                    } else {
                                        format!(
                                            "{}.{}.{}",
                                            self.module_name,
                                            self.class_stack.join("."),
                                            name.id
                                        )
                                    };

                                    let construct_type = if self.class_stack.is_empty() {
                                        ConstructType::ModuleAttribute
                                    } else {
                                        ConstructType::ClassAttribute
                                    };

                                    let replace_info = ReplaceInfo {
                                        old_name: full_name.clone(),
                                        replacement_expr,
                                        replacement_ast: None, // rustpython AST types differ
                                        construct_type,
                                        parameters: vec![],
                                        return_type: None,
                                        since: None,
                                        remove_in: None,
                                        message: None,
                                    };

                                    self.replacements.insert(full_name, replace_info);
                                }
                            }
                        }
                    }
                }
            }
            ast::Stmt::If(if_stmt) => {
                // Visit the if body
                for stmt in &if_stmt.body {
                    self.visit_stmt(stmt);
                }
                // Visit the else body
                for stmt in &if_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::Try(try_stmt) => {
                // Visit try body
                for stmt in &try_stmt.body {
                    self.visit_stmt(stmt);
                }
                // Visit except handlers
                for handler in &try_stmt.handlers {
                    let ast::ExceptHandler::ExceptHandler(h) = handler;
                    for stmt in &h.body {
                        self.visit_stmt(stmt);
                    }
                }
                // Visit else body
                for stmt in &try_stmt.orelse {
                    self.visit_stmt(stmt);
                }
                // Visit finally body
                for stmt in &try_stmt.finalbody {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::While(while_stmt) => {
                // Visit while body
                for stmt in &while_stmt.body {
                    self.visit_stmt(stmt);
                }
                // Visit else body
                for stmt in &while_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::For(for_stmt) => {
                // Visit for body
                for stmt in &for_stmt.body {
                    self.visit_stmt(stmt);
                }
                // Visit else body
                for stmt in &for_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::With(with_stmt) => {
                // Visit with body
                for stmt in &with_stmt.body {
                    self.visit_stmt(stmt);
                }
            }
            _ => {}
        }
    }
}

impl ReplacementExtractionError {
    fn reason(&self) -> ReplacementFailureReason {
        match self {
            Self::ExtractionFailed { reason, .. } => reason.clone(),
        }
    }
}