dissolve_python/

unified_visitor.rs

// Copyright (C) 2024 Jelmer Vernooij <jelmer@samba.org>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! Unified AST visitor for deprecated function collection and cleanup operations
//!
//! This visitor consolidates the shared logic between collection ('dissolve info')
//! and cleanup ('dissolve cleanup') operations. Migration continues to use the
//! specialized rustpython_visitor approach due to its complexity.

use crate::core::types::Version;
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;

/// Callback type for determining if a replacement should be removed
pub type ShouldRemoveCallback = Box<dyn Fn(&ReplaceInfo) -> bool>;

/// Operation mode for the unified visitor
#[derive(Debug, Clone, PartialEq)]
pub enum OperationMode {
    /// Collect deprecated functions and metadata
    Collect,
    /// Remove deprecated functions using a callback
    RemoveWithCallback,
}

/// Result type for unified operations
#[derive(Debug)]
pub enum UnifiedResult {
    Collection(CollectorResult),
    Removal(String), // cleaned source
}

/// Unified visitor that handles collection and cleanup operations
pub struct UnifiedVisitor {
    // Common fields
    module_name: String,
    _file_path: Option<std::path::PathBuf>,
    source: String,
    class_stack: Vec<String>,
    builtins: HashSet<String>,
    local_classes: HashSet<String>,
    local_functions: HashSet<String>,

    // Collection fields
    replacements: HashMap<String, ReplaceInfo>,
    unreplaceable: HashMap<String, UnreplaceableNode>,
    imports: Vec<ImportInfo>,
    inheritance_map: HashMap<String, Vec<String>>,
    class_methods: HashMap<String, HashSet<String>>,

    // Removal fields (when in remove mode)
    should_remove_callback: Option<ShouldRemoveCallback>,
    lines_to_remove: Vec<(usize, usize)>,

    // Operation mode
    operation_mode: OperationMode,
}

impl UnifiedVisitor {
    /// Create a new unified visitor for collection
    pub fn new_for_collection(module_name: &str, file_path: Option<&Path>) -> Self {
        Self::new_internal(module_name.to_string(), file_path, OperationMode::Collect)
    }

    /// Create a new unified visitor for removal with a callback
    pub fn new_for_removal_with_callback(
        module_name: &str,
        file_path: Option<&Path>,
        should_remove: ShouldRemoveCallback,
    ) -> Self {
        let mut visitor = Self::new_internal(
            module_name.to_string(),
            file_path,
            OperationMode::RemoveWithCallback,
        );
        visitor.should_remove_callback = Some(should_remove);
        visitor
    }

    /// Create a new unified visitor for removal with version criteria
    pub fn new_for_removal(
        module_name: &str,
        file_path: Option<&Path>,
        before_version: Option<&str>,
        remove_all: bool,
        current_version: Option<&str>,
    ) -> Self {
        let before_ver = before_version.and_then(|s| s.parse::<Version>().ok());
        let current_ver = current_version.and_then(|s| s.parse::<Version>().ok());

        let should_remove = Box::new(move |replace_info: &ReplaceInfo| {
            replace_info.should_remove(before_ver.as_ref(), remove_all, current_ver.as_ref())
        });

        Self::new_for_removal_with_callback(module_name, file_path, should_remove)
    }

    fn new_internal(
        module_name: String,
        file_path: Option<&Path>,
        operation_mode: OperationMode,
    ) -> Self {
        Self {
            module_name,
            _file_path: file_path.map(Path::to_path_buf),
            source: String::new(),
            class_stack: Vec::new(),
            builtins: Self::get_all_builtins(),
            local_classes: HashSet::new(),
            local_functions: HashSet::new(),

            replacements: HashMap::new(),
            unreplaceable: HashMap::new(),
            imports: Vec::new(),
            inheritance_map: HashMap::new(),
            class_methods: HashMap::new(),

            should_remove_callback: None,
            lines_to_remove: Vec::new(),
            operation_mode,
        }
    }

    /// Process source code according to the operation mode
    pub fn process_source(mut self, source: String) -> Result<UnifiedResult> {
        let parsed = parse(&source, Mode::Module, "<module>")?;
        self.source = source;

        match parsed {
            ast::Mod::Module(module) => {
                for stmt in &module.body {
                    self.visit_stmt(stmt);
                }
            }
            _ => {
                // Handle other module types if needed
            }
        }

        match self.operation_mode {
            OperationMode::Collect => Ok(UnifiedResult::Collection(CollectorResult {
                replacements: self.replacements,
                unreplaceable: self.unreplaceable,
                imports: self.imports,
                inheritance_map: self.inheritance_map,
                class_methods: self.class_methods,
            })),
            OperationMode::RemoveWithCallback => {
                let cleaned_source = self.apply_removals(&self.source);
                Ok(UnifiedResult::Removal(cleaned_source))
            }
        }
    }

    /// Apply line removals to source code
    fn apply_removals(&self, source: &str) -> String {
        // Sort removal ranges by start position for efficient processing
        let mut sorted_ranges: Vec<_> = self.lines_to_remove.iter().collect();
        sorted_ranges.sort_by_key(|(start, _)| *start);

        source
            .lines()
            .enumerate()
            .filter_map(|(i, line)| {
                // Check if line should be removed using binary search for efficiency
                let should_remove = sorted_ranges
                    .binary_search_by(|&(start, end)| {
                        if i < *start {
                            std::cmp::Ordering::Greater
                        } else if i >= *end {
                            std::cmp::Ordering::Less
                        } else {
                            std::cmp::Ordering::Equal
                        }
                    })
                    .is_ok();

                if should_remove {
                    None
                } else {
                    Some(line)
                }
            })
            .collect::<Vec<_>>()
            .join("\n")
    }

    /// 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
        })
    }

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

    /// Build the full object path including module and class names
    fn build_full_path(&self, name: &str) -> String {
        if self.class_stack.is_empty() {
            format!("{}.{}", self.module_name, name)
        } else {
            format!(
                "{}.{}.{}",
                self.module_name,
                self.class_stack.join("."),
                name
            )
        }
    }

    /// 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 decorator metadata (since, remove_in, message)
    fn extract_decorator_metadata(
        &self,
        decorators: &[ast::Expr],
    ) -> (Option<String>, Option<String>, Option<String>) {
        let since = Self::extract_decorator_arg(decorators, "since");
        let remove_in = Self::extract_decorator_arg(decorators, "remove_in");
        let message = Self::extract_decorator_arg(decorators, "message");
        (since, remove_in, message)
    }

    /// Extract any argument from @replace_me decorator
    fn extract_decorator_arg(decorators: &[ast::Expr], arg_name: &str) -> Option<String> {
        decorators
            .iter()
            .filter_map(|dec| {
                if let ast::Expr::Call(call) = dec {
                    if matches!(&*call.func, ast::Expr::Name(name) if name.id.as_str() == "replace_me") {
                        return call.keywords.iter().find_map(|keyword| {
                            keyword.arg.as_ref().and_then(|arg| {
                                if arg.as_str() == arg_name {
                                    Self::extract_value(&keyword.value)
                                } else {
                                    None
                                }
                            })
                        });
                    }
                }
                None
            })
            .next()
    }

    /// Extract string value from AST expression
    fn extract_value(expr: &ast::Expr) -> Option<String> {
        match expr {
            ast::Expr::Constant(c) => match &c.value {
                ast::Constant::Str(s) => Some(s.to_string()),
                ast::Constant::Int(i) => Some(i.to_string()),
                _ => None,
            },
            ast::Expr::Tuple(tuple) => {
                let parts: Vec<String> =
                    tuple.elts.iter().filter_map(Self::extract_value).collect();
                if parts.is_empty() {
                    None
                } else {
                    Some(parts.join("."))
                }
            }
            _ => None,
        }
    }

    /// Check if statement should be processed based on operation mode
    fn should_process_statement(&self, decorators: &[ast::Expr]) -> bool {
        Self::has_replace_me_decorator(decorators)
    }

    /// Check if a replacement should be removed using the callback
    fn should_remove_replacement(&self, replace_info: &ReplaceInfo) -> bool {
        self.should_remove_callback
            .as_ref()
            .is_some_and(|callback| callback(replace_info))
    }

    /// Create ReplaceInfo from decorator metadata
    fn create_replace_info(
        &self,
        name: &str,
        decorators: &[ast::Expr],
        construct_type: ConstructType,
    ) -> ReplaceInfo {
        let full_path = self.build_full_path(name);
        let (since, remove_in, message) = self.extract_decorator_metadata(decorators);

        let mut replace_info = ReplaceInfo::new(&full_path, "", construct_type);
        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;

        replace_info
    }

    /// Visit a statement in the AST
    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) => self.visit_import(import),
            ast::Stmt::ImportFrom(import) => self.visit_import_from(import),
            ast::Stmt::Assign(assign) => self.visit_assign(assign),
            ast::Stmt::AnnAssign(ann_assign) => self.visit_ann_assign(ann_assign),

            // Handle compound statements - visit their bodies
            ast::Stmt::If(if_stmt) => {
                for stmt in &if_stmt.body {
                    self.visit_stmt(stmt);
                }
                for stmt in &if_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::Try(try_stmt) => {
                for stmt in &try_stmt.body {
                    self.visit_stmt(stmt);
                }
                for handler in &try_stmt.handlers {
                    let ast::ExceptHandler::ExceptHandler(h) = handler;
                    for stmt in &h.body {
                        self.visit_stmt(stmt);
                    }
                }
                for stmt in &try_stmt.orelse {
                    self.visit_stmt(stmt);
                }
                for stmt in &try_stmt.finalbody {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::While(while_stmt) => {
                for stmt in &while_stmt.body {
                    self.visit_stmt(stmt);
                }
                for stmt in &while_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::For(for_stmt) => {
                for stmt in &for_stmt.body {
                    self.visit_stmt(stmt);
                }
                for stmt in &for_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            ast::Stmt::With(with_stmt) => {
                for stmt in &with_stmt.body {
                    self.visit_stmt(stmt);
                }
            }
            _ => {}
        }
    }

    fn visit_function(&mut self, func: &ast::StmtFunctionDef) {
        // Track locally defined functions
        if self.class_stack.is_empty() {
            self.local_functions.insert(func.name.to_string());
        }

        if !self.should_process_statement(&func.decorator_list) {
            return;
        }

        match &self.operation_mode {
            OperationMode::Collect => {
                self.collect_function(func);
            }
            OperationMode::RemoveWithCallback => {
                let replace_info = self.create_replace_info(
                    &func.name,
                    &func.decorator_list,
                    ConstructType::Function,
                );

                if self.should_remove_replacement(&replace_info) {
                    if let Some(line_range) = self.find_function_lines(func) {
                        self.lines_to_remove.push(line_range);
                    }
                }
            }
        }
    }

    fn visit_async_function(&mut self, func: &ast::StmtAsyncFunctionDef) {
        if !self.should_process_statement(&func.decorator_list) {
            return;
        }

        match &self.operation_mode {
            OperationMode::Collect => {
                self.collect_async_function(func);
            }
            OperationMode::RemoveWithCallback => {
                let replace_info = self.create_replace_info(
                    &func.name,
                    &func.decorator_list,
                    ConstructType::AsyncFunction,
                );

                if self.should_remove_replacement(&replace_info) {
                    if let Some(line_range) = self.find_async_function_lines(func) {
                        self.lines_to_remove.push(line_range);
                    }
                }
            }
        }
    }

    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);
        self.local_classes.insert(class_name.clone());
        let bases: Vec<String> = class_def
            .bases
            .iter()
            .filter_map(|base| {
                if let ast::Expr::Name(name) = base {
                    Some(format!("{}.{}", self.module_name, name.id))
                } else {
                    None
                }
            })
            .collect();

        if !bases.is_empty() {
            self.inheritance_map.insert(full_class_name.clone(), bases);
        }

        // Check if class has @replace_me
        if self.should_process_statement(&class_def.decorator_list) {
            match &self.operation_mode {
                OperationMode::Collect => {
                    self.collect_class(class_def);
                }
                OperationMode::RemoveWithCallback => {
                    let replace_info = self.create_replace_info(
                        &class_def.name,
                        &class_def.decorator_list,
                        ConstructType::Class,
                    );

                    if self.should_remove_replacement(&replace_info) {
                        if let Some(line_range) = self.find_class_lines(class_def) {
                            self.lines_to_remove.push(line_range);
                        }
                    }
                }
            }
        }

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

    fn visit_import(&mut self, import: &ast::StmtImport) {
        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()),
                )],
            ));
        }
    }

    fn visit_import_from(&mut self, import: &ast::StmtImportFrom) {
        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));
    }

    fn visit_assign(&mut self, assign: &ast::StmtAssign) {
        // Handle module/class attribute assignments with replace_me
        if assign.targets.len() == 1 {
            if let ast::Expr::Name(name) = &assign.targets[0] {
                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")
                    {
                        match &self.operation_mode {
                            OperationMode::Collect => {
                                self.collect_attribute_assignment(name, call);
                            }
                            OperationMode::RemoveWithCallback => {
                                // Create a basic ReplaceInfo for the assignment
                                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::new(&full_name, "", construct_type);

                                // Check if it should be removed - for now, assignment removal not implemented
                                if self.should_remove_replacement(&replace_info) {
                                    // TODO: Find assignment line range for removal
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    fn visit_ann_assign(&mut self, ann_assign: &ast::StmtAnnAssign) {
        // Handle annotated assignments with replace_me
        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")
                    {
                        match &self.operation_mode {
                            OperationMode::Collect => {
                                self.collect_attribute_assignment(name, call);
                            }
                            OperationMode::RemoveWithCallback => {
                                // Create a basic ReplaceInfo for the assignment
                                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::new(&full_name, "", construct_type);

                                // Check if it should be removed - for now, assignment removal not implemented
                                if self.should_remove_replacement(&replace_info) {
                                    // TODO: Find assignment line range for removal
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    /// Extract replacement expression from function body
    fn extract_replacement_from_function(&self, func: &ast::StmtFunctionDef) -> Result<String> {
        // 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,
                    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 means remove the function completely
            return Ok("".to_string());
        }

        if body_stmts.len() == 1 {
            // Single statement - should be a return statement
            if let ast::Stmt::Return(ret_stmt) = body_stmts[0] {
                if let Some(value) = &ret_stmt.value {
                    // Get function parameter names
                    let param_names: Vec<String> = func
                        .args
                        .args
                        .iter()
                        .map(|arg| arg.def.arg.to_string())
                        .collect();

                    // Convert expression to string with parameter placeholders
                    let param_str_refs: Vec<&str> =
                        param_names.iter().map(|s| s.as_str()).collect();
                    return self.expr_to_string_with_placeholders(value, &param_str_refs);
                }
            }
        }

        // Multiple statements or no clear replacement
        Ok("".to_string())
    }

    /// Convert expression to string with parameter placeholders
    fn expr_to_string_with_placeholders(
        &self,
        expr: &ast::Expr,
        param_names: &[&str],
    ) -> Result<String> {
        match expr {
            ast::Expr::Name(name) => {
                if param_names.contains(&name.id.as_str()) {
                    Ok(format!("{{{}}}", name.id))
                } else {
                    Ok(name.id.to_string())
                }
            }
            ast::Expr::Call(call) => {
                let func_str = self.expr_to_string_with_placeholders(&call.func, param_names)?;
                let mut all_args: Vec<String> = Vec::new();

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

                // Add 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)?;
                        all_args.push(format!("{}={}", arg_name, value_str));
                    } else {
                        // **kwargs expansion
                        let value_str = self.expr_to_string_with_placeholders(&keyword.value, param_names)?;
                        all_args.push(format!("**{}", value_str));
                    }
                }

                Ok(format!("{}({})", func_str, all_args.join(", ")))
            }
            ast::Expr::Attribute(attr) => {
                let value_str = self.expr_to_string_with_placeholders(&attr.value, param_names)?;
                Ok(format!("{}.{}", value_str, attr.attr))
            }
            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 => "@",
                };
                Ok(format!("{} {} {}", left, op_str, right))
            }
            ast::Expr::UnaryOp(unaryop) => {
                let operand = self.expr_to_string_with_placeholders(&unaryop.operand, param_names)?;
                let op_str = match &unaryop.op {
                    ast::UnaryOp::Not => "not ",
                    ast::UnaryOp::UAdd => "+",
                    ast::UnaryOp::USub => "-",
                    ast::UnaryOp::Invert => "~",
                };
                Ok(format!("{}{}", op_str, operand))
            }
            ast::Expr::List(list) => {
                let mut elts = Vec::new();
                for e in &list.elts {
                    elts.push(self.expr_to_string_with_placeholders(e, param_names)?);
                }
                Ok(format!("[{}]", elts.join(", ")))
            }
            ast::Expr::Tuple(tuple) => {
                let mut elts = Vec::new();
                for e in &tuple.elts {
                    elts.push(self.expr_to_string_with_placeholders(e, param_names)?);
                }
                if elts.len() == 1 {
                    Ok(format!("({},)", elts[0]))
                } else {
                    Ok(format!("({})", elts.join(", ")))
                }
            }
            ast::Expr::Dict(dict) => {
                let mut items = Vec::new();
                for (k, v) in dict.keys.iter().zip(&dict.values) {
                    if let Some(key) = k {
                        let key_str = self.expr_to_string_with_placeholders(key, param_names)?;
                        let val_str = self.expr_to_string_with_placeholders(v, param_names)?;
                        items.push(format!("{}: {}", key_str, val_str));
                    }
                }
                Ok(format!("{{{}}}", items.join(", ")))
            }
            ast::Expr::Constant(c) => match &c.value {
                ast::Constant::Str(s) => Ok(format!("\"{}\"", s.escape_default())),
                ast::Constant::Int(i) => Ok(i.to_string()),
                ast::Constant::Float(f) => Ok(f.to_string()),
                ast::Constant::Bool(b) => Ok(if *b { "True" } else { "False" }.to_string()),
                ast::Constant::None => Ok("None".to_string()),
                ast::Constant::Ellipsis => Ok("...".to_string()),
                _ => {
                    Err(anyhow::anyhow!("Unsupported constant type in replacement expression"))
                }
            },
            ast::Expr::Starred(starred) => {
                let value = self.expr_to_string_with_placeholders(&starred.value, param_names)?;
                Ok(format!("*{}", value))
            }
            ast::Expr::Subscript(sub) => {
                let value = self.expr_to_string_with_placeholders(&sub.value, param_names)?;
                let slice = self.expr_to_string_with_placeholders(&sub.slice, param_names)?;
                Ok(format!("{}[{}]", value, slice))
            }
            // Complex expressions that we don't support
            _ => {
                Err(anyhow::anyhow!(
                    "Replacement expression contains unsupported construct (e.g., comprehensions, conditionals, comparisons). Skipping function."
                ))
            }
        }
    }

    // Collection methods (simplified from stub_collector.rs)
    fn collect_function(&mut self, func: &ast::StmtFunctionDef) {
        let full_path = self.build_full_path(&func.name);
        let (since, remove_in, message) = self.extract_decorator_metadata(&func.decorator_list);

        // Determine construct type
        let construct_type = if self.class_stack.is_empty() {
            ConstructType::Function
        } else {
            // Check for method decorators
            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
            }
        };

        // Extract replacement from function body
        let replacement_expr = match self.extract_replacement_from_function(func) {
            Ok(expr) => expr,
            Err(e) => {
                // Log warning and skip this function
                tracing::warn!(
                    "Skipping function '{}' in {}: {}",
                    full_path,
                    self._file_path
                        .as_ref()
                        .and_then(|p| p.to_str())
                        .unwrap_or("<unknown>"),
                    e
                );
                // Return early without adding to replacements
                return;
            }
        };

        // Create ReplaceInfo with the extracted replacement
        let mut replace_info = ReplaceInfo::new(&full_path, &replacement_expr, construct_type);
        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);
    }

    fn collect_async_function(&mut self, func: &ast::StmtAsyncFunctionDef) {
        let full_path = self.build_full_path(&func.name);
        let (since, remove_in, message) = self.extract_decorator_metadata(&func.decorator_list);

        let mut replace_info = ReplaceInfo::new(&full_path, "", ConstructType::AsyncFunction);
        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);
    }

    fn collect_class(&mut self, class_def: &ast::StmtClassDef) {
        let full_path = self.build_full_path(&class_def.name);
        let (since, remove_in, message) =
            self.extract_decorator_metadata(&class_def.decorator_list);

        let mut replace_info = ReplaceInfo::new(&full_path, "", ConstructType::Class);
        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);
    }

    fn collect_attribute_assignment(&mut self, name: &ast::ExprName, call: &ast::ExprCall) {
        if let Some(arg) = call.args.first() {
            // Simple string representation for the assignment value
            let replacement_expr = format!("{:?}", arg); // Basic fallback for collection

            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::new(&full_name, &replacement_expr, construct_type);
            self.replacements.insert(full_name, replace_info);
        }
    }

    // Line finding methods for removal (simplified from remover.rs)
    fn find_function_lines(&self, func: &ast::StmtFunctionDef) -> Option<(usize, usize)> {
        self.find_statement_lines_by_name("def", &func.name)
    }

    fn find_async_function_lines(
        &self,
        func: &ast::StmtAsyncFunctionDef,
    ) -> Option<(usize, usize)> {
        self.find_statement_lines_by_name("async def", &func.name)
    }

    fn find_class_lines(&self, class_def: &ast::StmtClassDef) -> Option<(usize, usize)> {
        self.find_statement_lines_by_name("class", &class_def.name)
    }

    fn find_statement_lines_by_name(&self, keyword: &str, name: &str) -> Option<(usize, usize)> {
        let lines: Vec<&str> = self.source.lines().collect();

        for (i, line) in lines.iter().enumerate() {
            if line.contains(&format!("{} {}", keyword, name)) {
                // Find the end of the statement (next def, class, or dedent)
                let indent = line.chars().take_while(|c| c.is_whitespace()).count();
                for (j, end_line) in lines[i + 1..].iter().enumerate() {
                    let end_i = i + j + 1;
                    if !end_line.trim().is_empty() {
                        let end_indent = end_line.chars().take_while(|c| c.is_whitespace()).count();
                        if end_indent <= indent && !end_line.trim_start().starts_with('#') {
                            // Include decorators before the statement
                            let start = self.find_decorator_start(&lines, i);
                            return Some((start, end_i));
                        }
                    }
                }
                // If we reach the end, include everything
                let start = self.find_decorator_start(&lines, i);
                return Some((start, lines.len()));
            }
        }
        None
    }

    fn find_decorator_start(&self, lines: &[&str], def_line: usize) -> usize {
        // Look backwards for decorators
        let mut start = def_line;
        for i in (0..def_line).rev() {
            let line = lines[i].trim();
            if line.starts_with('@') || line.is_empty() || line.starts_with('#') {
                start = i;
            } else {
                break;
            }
        }
        start
    }
}

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

    #[test]
    fn test_unified_visitor_collection() {
        let source = r#"
from dissolve import replace_me

@replace_me(since="1.0.0")
def old_function():
    return new_function()

def regular_function():
    return 42
"#;

        let visitor = UnifiedVisitor::new_for_collection("test_module", None);
        let result = visitor.process_source(source.to_string()).unwrap();

        match result {
            UnifiedResult::Collection(collection) => {
                assert!(collection
                    .replacements
                    .contains_key("test_module.old_function"));
                assert!(!collection
                    .replacements
                    .contains_key("test_module.regular_function"));
            }
            _ => panic!("Expected Collection result"),
        }
    }

    #[test]
    fn test_unified_visitor_removal_criteria() {
        use crate::core::types::Version;

        // Test the callback approach
        let should_remove_callback = Box::new(|replace_info: &ReplaceInfo| {
            if let Some(since) = &replace_info.since {
                let before_version = Version::new(2, 0, 0);
                since < &before_version
            } else {
                false
            }
        });

        let visitor = UnifiedVisitor::new_for_removal_with_callback(
            "test_module",
            None,
            should_remove_callback,
        );

        // Create a ReplaceInfo with since="1.5.0"
        let replace_info = ReplaceInfo::new("test.func", "new_func()", ConstructType::Function)
            .with_since_version(Version::new(1, 5, 0));

        assert!(visitor.should_remove_replacement(&replace_info));

        // Test that version >= 2.0.0 is NOT removed
        let replace_info2 = ReplaceInfo::new("test.func2", "new_func2()", ConstructType::Function)
            .with_since_version(Version::new(2, 1, 0));

        assert!(!visitor.should_remove_replacement(&replace_info2));
    }

    #[test]
    fn test_unified_visitor_removal() {
        let source = r#"
from dissolve import replace_me

@replace_me(since="1.0.0")
def old_function():
    return new_function()

def regular_function():
    return 42

@replace_me(since="2.0.0")
def newer_function():
    return new_api()
"#;

        let visitor = UnifiedVisitor::new_for_removal(
            "test_module",
            None,
            Some("1.5.0"), // Remove functions deprecated before 1.5.0
            false,
            None,
        );

        let result = visitor.process_source(source.to_string()).unwrap();

        match result {
            UnifiedResult::Removal(cleaned_source) => {
                // old_function (since="1.0.0") should be removed
                assert!(!cleaned_source.contains("def old_function"));
                // regular_function should remain
                assert!(cleaned_source.contains("def regular_function"));
                // newer_function (since="2.0.0") should remain
                assert!(cleaned_source.contains("def newer_function"));
            }
            _ => panic!("Expected Removal result"),
        }
    }
}