dissolve_python/

ruff_parser.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.

//! Python parser and CST manipulation using Ruff's parser.
//!
//! This module provides a Rust implementation that preserves formatting
//! and integrates with mypy for type inference.

use anyhow::{anyhow, Result};
use rustpython_ast::{ExceptHandler as AstExceptHandler, Expr as AstExpr, Ranged, Stmt as AstStmt};
use rustpython_parser::{parse, Mode};
use std::collections::HashMap;
use std::ops::Range;

use crate::core::{CollectorResult, ReplaceInfo};
use crate::types::TypeIntrospectionMethod;

/// Parse Python source code preserving all formatting information
pub struct PythonModule<'a> {
    source: &'a str,
    ast: Vec<AstStmt>,
    /// Map from byte offset to line/column for mypy integration
    position_map: HashMap<u32, (u32, u32)>,
}

impl<'a> PythonModule<'a> {
    /// Parse Python source code
    pub fn parse(source: &'a str) -> Result<Self> {
        let parsed =
            parse(source, Mode::Module, "<module>").map_err(|e| anyhow!("Parse error: {:?}", e))?;

        // Convert Mod to Suite
        let ast = match parsed {
            rustpython_ast::Mod::Module(module) => module.body,
            _ => return Err(anyhow!("Expected Module, got Expression")),
        };

        // Build position map for byte offset -> line/column conversion
        let position_map = Self::build_position_map(source);

        Ok(Self {
            source,
            ast,
            position_map,
        })
    }

    /// Build a map from byte offset to (line, column)
    fn build_position_map(source: &str) -> HashMap<u32, (u32, u32)> {
        let mut map = HashMap::new();
        let mut line = 1;
        let mut col = 0;

        for (offset, ch) in source.char_indices() {
            map.insert(offset as u32, (line, col));
            if ch == '\n' {
                line += 1;
                col = 0;
            } else {
                col += 1;
            }
        }

        // Add end position
        map.insert(source.len() as u32, (line, col));

        map
    }

    /// Get AST
    pub fn ast(&self) -> &[AstStmt] {
        &self.ast
    }

    /// Convert byte offset to line/column for mypy
    pub fn offset_to_position(&self, offset: usize) -> Option<(u32, u32)> {
        self.position_map.get(&(offset as u32)).copied()
    }

    /// Convert byte offset to line/column (alias for compatibility)
    pub fn line_col_at_offset(&self, offset: usize) -> (u32, u32) {
        self.offset_to_position(offset).unwrap_or((1, 0))
    }

    /// Get text for a range
    pub fn text_at_range(&self, range: Range<usize>) -> &str {
        &self.source[range]
    }
}

/// Collect deprecated functions using Ruff's AST
/// For now, we delegate to LibCST collector until we implement full extraction
pub fn collect_deprecated_functions(_source: &str, _module_name: &str) -> Result<CollectorResult> {
    // TODO: Complete rustpython migration
    // For now, return empty result
    Ok(CollectorResult {
        replacements: HashMap::new(),
        unreplaceable: HashMap::new(),
        imports: Vec::new(),
        inheritance_map: HashMap::new(),
        class_methods: HashMap::new(),
    })
}

/// Visitor to find and replace function calls
pub struct FunctionCallReplacer<'a> {
    replacements_info: HashMap<String, ReplaceInfo>,
    replacements: Vec<(Range<usize>, String)>,
    #[allow(dead_code)]
    source_module: &'a PythonModule<'a>,
    #[allow(dead_code)]
    type_introspection: TypeIntrospectionMethod,
    #[allow(dead_code)]
    file_path: String,
    #[allow(dead_code)]
    module_name: String,
    import_map: HashMap<String, String>, // Maps imported names to their full module paths
}

impl<'a> FunctionCallReplacer<'a> {
    pub fn new(
        replacements: HashMap<String, ReplaceInfo>,
        source_module: &'a PythonModule<'a>,
        type_introspection: TypeIntrospectionMethod,
        file_path: String,
        module_name: String,
    ) -> Self {
        let mut replacer = Self {
            replacements_info: replacements,
            replacements: Vec::new(),
            source_module,
            type_introspection,
            file_path,
            module_name,
            import_map: HashMap::new(),
        };
        // Collect imports from the module
        replacer.collect_imports();
        replacer
    }

    pub fn get_replacements(self) -> Vec<(Range<usize>, String)> {
        self.replacements
    }

    /// Collect import statements and build the import map
    fn collect_imports(&mut self) {
        for stmt in self.source_module.ast() {
            match stmt {
                AstStmt::Import(import) => {
                    // Handle: import module.submodule as alias
                    for alias in &import.names {
                        let module_name = alias.name.to_string();
                        let alias_name = alias
                            .asname
                            .as_ref()
                            .map(|id| id.to_string())
                            .unwrap_or_else(|| module_name.clone());
                        self.import_map.insert(alias_name, module_name);
                    }
                }
                AstStmt::ImportFrom(import_from) => {
                    if let Some(module) = &import_from.module {
                        let module_str = module.to_string();
                        // Handle: from module import name as alias
                        for alias in &import_from.names {
                            let imported_name = alias.name.to_string();
                            let alias_name = alias
                                .asname
                                .as_ref()
                                .map(|id| id.to_string())
                                .unwrap_or_else(|| imported_name.clone());

                            // Only track imports without aliases for now
                            // Alias tracking would require more complex import rewriting
                            if alias.asname.is_none() {
                                let full_name = format!("{}.{}", module_str, imported_name);
                                self.import_map.insert(alias_name, full_name);
                            }
                        }
                    }
                }
                _ => {}
            }
        }
    }
}

impl<'a> FunctionCallReplacer<'a> {
    pub fn visit_expr(&mut self, expr: &'a AstExpr) {
        match expr {
            AstExpr::Call(call) => {
                // Check if this is a magic method call (int, repr, bool, len, etc.)
                if let Some(replacement) = self.try_replace_magic_method_call(call) {
                    let range = expr.range();
                    let byte_range = std::ops::Range {
                        start: usize::from(range.start()),
                        end: usize::from(range.end()),
                    };
                    self.replacements.push((byte_range, replacement));
                } else if let Some(func_name) = self.get_function_name(&call.func) {
                    // Check if this is a regular function call we need to replace
                    if let Some(replace_info) = self.replacements_info.get(&func_name) {
                        // We found a function to replace
                        // Build the replacement with actual arguments
                        let replacement = self.build_replacement(replace_info, call);

                        // Get the range of the entire call expression using the Ranged trait
                        let range = expr.range();
                        let byte_range = std::ops::Range {
                            start: usize::from(range.start()),
                            end: usize::from(range.end()),
                        };

                        self.replacements.push((byte_range, replacement));
                    }
                }

                // Visit arguments recursively
                for arg in &call.args {
                    self.visit_expr(arg);
                }
            }
            AstExpr::Attribute(attr) => {
                self.visit_expr(&attr.value);
            }
            AstExpr::BinOp(binop) => {
                self.visit_expr(&binop.left);
                self.visit_expr(&binop.right);
            }
            AstExpr::UnaryOp(unaryop) => {
                self.visit_expr(&unaryop.operand);
            }
            AstExpr::IfExp(ifexp) => {
                self.visit_expr(&ifexp.test);
                self.visit_expr(&ifexp.body);
                self.visit_expr(&ifexp.orelse);
            }
            AstExpr::Dict(dict) => {
                for key in dict.keys.iter().flatten() {
                    self.visit_expr(key);
                }
                for value in &dict.values {
                    self.visit_expr(value);
                }
            }
            AstExpr::Set(set) => {
                for elt in &set.elts {
                    self.visit_expr(elt);
                }
            }
            AstExpr::ListComp(comp) => {
                self.visit_expr(&comp.elt);
                // Would need to visit generators too
            }
            AstExpr::SetComp(comp) => {
                self.visit_expr(&comp.elt);
                // Would need to visit generators too
            }
            AstExpr::DictComp(comp) => {
                self.visit_expr(&comp.key);
                self.visit_expr(&comp.value);
                // Would need to visit generators too
            }
            AstExpr::GeneratorExp(gen) => {
                self.visit_expr(&gen.elt);
                // Would need to visit generators too
            }
            AstExpr::Await(await_expr) => {
                self.visit_expr(&await_expr.value);
            }
            AstExpr::Yield(yield_expr) => {
                if let Some(value) = &yield_expr.value {
                    self.visit_expr(value);
                }
            }
            AstExpr::YieldFrom(yieldfrom) => {
                self.visit_expr(&yieldfrom.value);
            }
            AstExpr::Compare(compare) => {
                self.visit_expr(&compare.left);
                for comparator in &compare.comparators {
                    self.visit_expr(comparator);
                }
            }
            AstExpr::Lambda(lambda) => {
                self.visit_expr(&lambda.body);
            }
            AstExpr::List(list) => {
                for elt in &list.elts {
                    self.visit_expr(elt);
                }
            }
            AstExpr::Tuple(tuple) => {
                for elt in &tuple.elts {
                    self.visit_expr(elt);
                }
            }
            AstExpr::Slice(slice) => {
                if let Some(lower) = &slice.lower {
                    self.visit_expr(lower);
                }
                if let Some(upper) = &slice.upper {
                    self.visit_expr(upper);
                }
                if let Some(step) = &slice.step {
                    self.visit_expr(step);
                }
            }
            AstExpr::Subscript(subscript) => {
                self.visit_expr(&subscript.value);
                self.visit_expr(&subscript.slice);
            }
            AstExpr::Starred(starred) => {
                self.visit_expr(&starred.value);
            }
            AstExpr::NamedExpr(named) => {
                self.visit_expr(&named.target);
                self.visit_expr(&named.value);
            }
            AstExpr::Name(_) | AstExpr::Constant(_) => {
                // Leaf nodes, nothing to visit
            }
            _ => {
                // Other expression types we don't handle yet
            }
        }
    }

    pub fn visit_stmt(&mut self, stmt: &'a AstStmt) {
        match stmt {
            AstStmt::FunctionDef(func) => {
                // Visit function body
                for stmt in &func.body {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::AsyncFunctionDef(func) => {
                // Visit async function body
                for stmt in &func.body {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::ClassDef(class) => {
                // Visit class body
                for stmt in &class.body {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::Return(ret) => {
                if let Some(value) = &ret.value {
                    self.visit_expr(value);
                }
            }
            AstStmt::Delete(del) => {
                for target in &del.targets {
                    self.visit_expr(target);
                }
            }
            AstStmt::Assign(assign) => {
                self.visit_expr(&assign.value);
                for target in &assign.targets {
                    self.visit_expr(target);
                }
            }
            AstStmt::AugAssign(augassign) => {
                self.visit_expr(&augassign.target);
                self.visit_expr(&augassign.value);
            }
            AstStmt::AnnAssign(annassign) => {
                self.visit_expr(&annassign.target);
                if let Some(value) = &annassign.value {
                    self.visit_expr(value);
                }
            }
            AstStmt::For(for_stmt) => {
                self.visit_expr(&for_stmt.target);
                self.visit_expr(&for_stmt.iter);
                for stmt in &for_stmt.body {
                    self.visit_stmt(stmt);
                }
                for stmt in &for_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::AsyncFor(for_stmt) => {
                self.visit_expr(&for_stmt.target);
                self.visit_expr(&for_stmt.iter);
                for stmt in &for_stmt.body {
                    self.visit_stmt(stmt);
                }
                for stmt in &for_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::While(while_stmt) => {
                self.visit_expr(&while_stmt.test);
                for stmt in &while_stmt.body {
                    self.visit_stmt(stmt);
                }
                for stmt in &while_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::If(if_stmt) => {
                self.visit_expr(&if_stmt.test);
                for stmt in &if_stmt.body {
                    self.visit_stmt(stmt);
                }
                for stmt in &if_stmt.orelse {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::With(with_stmt) => {
                for item in &with_stmt.items {
                    self.visit_expr(&item.context_expr);
                }
                for stmt in &with_stmt.body {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::AsyncWith(with_stmt) => {
                for item in &with_stmt.items {
                    self.visit_expr(&item.context_expr);
                }
                for stmt in &with_stmt.body {
                    self.visit_stmt(stmt);
                }
            }
            AstStmt::Raise(raise_stmt) => {
                if let Some(exc) = &raise_stmt.exc {
                    self.visit_expr(exc);
                }
                if let Some(cause) = &raise_stmt.cause {
                    self.visit_expr(cause);
                }
            }
            AstStmt::Try(try_stmt) => {
                for stmt in &try_stmt.body {
                    self.visit_stmt(stmt);
                }
                for handler in &try_stmt.handlers {
                    match handler {
                        AstExceptHandler::ExceptHandler(h) => {
                            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);
                }
            }
            AstStmt::Assert(assert_stmt) => {
                self.visit_expr(&assert_stmt.test);
                if let Some(msg) = &assert_stmt.msg {
                    self.visit_expr(msg);
                }
            }
            AstStmt::Expr(expr_stmt) => {
                self.visit_expr(&expr_stmt.value);
            }
            AstStmt::Pass(_) | AstStmt::Break(_) | AstStmt::Continue(_) => {
                // Nothing to visit
            }
            _ => {
                // Other statement types we don't handle yet
            }
        }
    }

    fn get_function_name(&self, expr: &AstExpr) -> Option<String> {
        match expr {
            AstExpr::Name(name) => {
                let simple_name = name.id.to_string();

                // First check if this name was imported
                if let Some(full_name) = self.import_map.get(&simple_name) {
                    // This name was imported, use its full module path
                    if self.replacements_info.contains_key(full_name) {
                        return Some(full_name.clone());
                    }
                }

                // Check if we have a qualified version of this name
                let qualified_name = format!("{}.{}", self.module_name, simple_name);
                // Try the qualified name first, then the simple name
                #[allow(clippy::map_entry)]
                // This is not contains_key + insert, it's conditional lookup
                if self.replacements_info.contains_key(&qualified_name) {
                    Some(qualified_name)
                } else if self.replacements_info.contains_key(&simple_name) {
                    Some(simple_name)
                } else {
                    None
                }
            }
            AstExpr::Attribute(attr) => {
                // For method calls like obj.method(), we need to check all possible matches
                let method_name = attr.attr.to_string();

                // Try all keys that end with the method name
                for key in self.replacements_info.keys() {
                    if key.ends_with(&format!(".{}", method_name)) {
                        return Some(key.clone());
                    }
                }

                // Fallback to simple name
                Some(method_name)
            }
            _ => None,
        }
    }

    fn build_replacement(
        &self,
        replace_info: &ReplaceInfo,
        call: &rustpython_ast::ExprCall,
    ) -> String {
        let mut replacement = replace_info.replacement_expr.clone();

        // Strip module prefix if present
        if replacement.starts_with(&format!("{}.", self.module_name)) {
            replacement = replacement[self.module_name.len() + 1..].to_string();
        }

        // Special case: handle literal *args, **kwargs pattern
        if replacement.contains("(*args, **kwargs)") {
            // Build the actual call with all arguments
            let func_name = replacement.split('(').next().unwrap_or(&replacement);
            let mut new_call = format!("{}(", func_name);

            // Add positional arguments
            for (i, arg) in call.args.iter().enumerate() {
                if i > 0 {
                    new_call.push_str(", ");
                }
                new_call.push_str(&self.expr_to_source(arg));
            }

            // Add keyword arguments
            if !call.keywords.is_empty() {
                if !call.args.is_empty() {
                    new_call.push_str(", ");
                }
                for (i, keyword) in call.keywords.iter().enumerate() {
                    if i > 0 {
                        new_call.push_str(", ");
                    }
                    if let Some(arg_name) = &keyword.arg {
                        new_call.push_str(&format!(
                            "{}={}",
                            arg_name,
                            self.expr_to_source(&keyword.value)
                        ));
                    } else {
                        // **kwargs
                        new_call.push_str(&format!("**{}", self.expr_to_source(&keyword.value)));
                    }
                }
            }

            new_call.push(')');
            return new_call;
        }

        // Simple case: if no placeholders, just return the replacement as-is
        if !replacement.contains("{") {
            return replacement;
        }

        // Handle special cases like {self} placeholders
        if replacement.contains("{self}") || replacement.contains("{cls}") {
            // For method calls, we need to handle the receiver specially
            if let AstExpr::Attribute(attr) = call.func.as_ref() {
                let receiver = self.expr_to_source(&attr.value);
                replacement = replacement.replace("{self}", &receiver);
                replacement = replacement.replace("{cls}", &receiver);
            }
        }

        // Build a map of parameter names to argument values
        let mut arg_map = HashMap::new();

        // Get the parameters, skipping self/cls if present
        let params_start = if !replace_info.parameters.is_empty()
            && (replace_info.parameters[0].name == "self"
                || replace_info.parameters[0].name == "cls")
        {
            1
        } else {
            0
        };

        // Map positional arguments
        for (i, arg) in call.args.iter().enumerate() {
            if let Some(param) = replace_info.parameters.get(i + params_start) {
                let arg_str = self.expr_to_source(arg);
                arg_map.insert(param.name.clone(), arg_str);
            }
        }

        // Map keyword arguments
        for keyword in &call.keywords {
            if let Some(arg_name) = &keyword.arg {
                let arg_str = self.expr_to_source(&keyword.value);
                arg_map.insert(arg_name.to_string(), arg_str);
            } else {
                // Handle **kwargs
                let arg_str = self.expr_to_source(&keyword.value);
                arg_map.insert("**kwargs".to_string(), format!("**{}", arg_str));
            }
        }

        // Handle *args if present
        let _has_varargs = replace_info.parameters.iter().any(|p| p.is_vararg);
        let _has_kwargs = replace_info.parameters.iter().any(|p| p.is_kwarg);

        // Replace placeholders with actual arguments
        for (param_name, arg_value) in &arg_map {
            let placeholder = format!("{{{}}}", param_name);
            replacement = replacement.replace(&placeholder, arg_value);
        }

        // For *args and **kwargs, we need special handling
        if replacement.contains("{*args}") {
            // Calculate how many positional arguments are consumed by named parameters
            let named_param_count = replace_info
                .parameters
                .iter()
                .take_while(|p| !p.is_vararg && !p.is_kwarg && !p.is_kwonly)
                .count();

            // Skip parameters for self/cls if present
            let actual_named_count = named_param_count.saturating_sub(params_start);

            // Get the extra positional arguments that go to *args
            let extra_args: Vec<String> = call
                .args
                .iter()
                .skip(actual_named_count)
                .map(|arg| self.expr_to_source(arg))
                .collect();

            let args_str = extra_args.join(", ");

            // Replace {*args} with the extra arguments
            if !args_str.is_empty() {
                replacement = replacement.replace("{*args}", &args_str);
            } else {
                replacement = replacement.replace(", {*args}", "");
                replacement = replacement.replace("{*args}", "");
            }
        }

        if replacement.contains("{**kwargs}") {
            // Collect all keyword arguments
            let kwargs_strs: Vec<String> = call
                .keywords
                .iter()
                .map(|kw| {
                    if let Some(arg) = &kw.arg {
                        format!("{}={}", arg, self.expr_to_source(&kw.value))
                    } else {
                        format!("**{}", self.expr_to_source(&kw.value))
                    }
                })
                .collect();

            if !kwargs_strs.is_empty() {
                let kwargs_str = kwargs_strs.join(", ");
                replacement = replacement.replace("{**kwargs}", &kwargs_str);
            } else {
                replacement = replacement.replace(", {**kwargs}", "");
                replacement = replacement.replace("{**kwargs}", "");
            }
        }

        // Clean up unresolved placeholders for parameters with default values
        let mut cleaned = replacement.clone();

        // Remove placeholders for parameters that have default values and weren't provided
        // But only remove simple parameter assignments, not placeholders in complex expressions
        for param in &replace_info.parameters {
            if param.has_default && !arg_map.contains_key(&param.name) {
                let placeholder = format!("{{{}}}", param.name);

                // Only remove keyword argument patterns where the placeholder is the entire value
                // Pattern: "param={param}" - remove entire assignment
                cleaned = cleaned.replace(&format!(", {}={}", param.name, placeholder), "");
                cleaned = cleaned.replace(&format!("{}={}, ", param.name, placeholder), "");
                cleaned = cleaned.replace(&format!("{}={}", param.name, placeholder), "");

                // Do NOT remove bare placeholders that might be part of complex expressions
                // like "{cache_size} * 2" - these should remain as placeholders for later substitution
            }
        }

        // Clean up any remaining syntax issues from placeholder removal
        // Fix multiple consecutive commas
        while cleaned.contains(", ,") {
            cleaned = cleaned.replace(", ,", ",");
        }
        // Fix trailing commas before closing parenthesis
        cleaned = cleaned.replace(", )", ")");
        // Fix leading commas after opening parenthesis
        cleaned = cleaned.replace("(, ", "(");

        // Remove empty *args and **kwargs placeholders
        cleaned = cleaned.replace(", {*args}", "");
        cleaned = cleaned.replace("{*args}, ", "");
        cleaned = cleaned.replace("({*args})", "()");
        cleaned = cleaned.replace("{*args}", "");

        cleaned = cleaned.replace(", {**kwargs}", "");
        cleaned = cleaned.replace("{**kwargs}, ", "");
        cleaned = cleaned.replace("({**kwargs})", "()");
        cleaned = cleaned.replace("{**kwargs}", "");

        cleaned
    }

    fn try_replace_magic_method_call(&self, call: &rustpython_ast::ExprCall) -> Option<String> {
        // Check if this is a built-in function call that corresponds to a magic method
        if let AstExpr::Name(name) = call.func.as_ref() {
            let magic_method_map = match name.id.as_str() {
                "int" => "__int__",
                "repr" => "__repr__",
                "bool" => "__bool__",
                "len" => "__len__",
                "str" => "__str__",
                "float" => "__float__",
                "bytes" => "__bytes__",
                "hash" => "__hash__",
                _ => return None,
            };

            // For now, we need exactly one argument (the object to call the method on)
            if call.args.len() != 1 {
                return None;
            }

            let obj_expr = &call.args[0];
            let obj_str = self.expr_to_source(obj_expr);

            // Look for a replacement info for this magic method
            // We need to find if any of our replacements match this magic method pattern
            for (key, replace_info) in &self.replacements_info {
                // Check if this replacement is for a magic method that matches
                if key.ends_with(&format!(".{}", magic_method_map)) {
                    // Extract the class name from the key
                    if let Some(class_name) = key.rsplit('.').nth(1) {
                        // Basic heuristic: only apply magic method replacements if the argument
                        // is a simple name or looks like it could be of the right type.
                        // Avoid replacing complex expressions like obj.attr, obj[key], etc.
                        if !self.should_apply_magic_method_replacement(obj_expr, class_name) {
                            continue;
                        }

                        // Extract the replacement method name from the replacement expression
                        // Handle patterns like:
                        // 1. "{self}.method()" -> "obj.method()"
                        // 2. "builtin({self}.method())" -> "obj.method()"
                        let replacement_expr = &replace_info.replacement_expr;

                        // First try to find pattern: "builtin({self}.method())" or "builtin({self}.attr)"
                        // Handle both with and without module prefix
                        let patterns = [
                            format!("{}({{self}}.", name.id),                      // int({self}.
                            format!("{}.{}({{self}}.", self.module_name, name.id), // test_module.int({self}.
                        ];

                        for builtin_start in &patterns {
                            if replacement_expr.starts_with(builtin_start) {
                                if let Some(inner_part) =
                                    replacement_expr.strip_prefix(builtin_start)
                                {
                                    // Handle method calls: "builtin({self}.method())" -> "obj.method()"
                                    if inner_part.ends_with("())") {
                                        if let Some(method_name) = inner_part.strip_suffix("())") {
                                            let result = format!("{}.{}()", obj_str, method_name);
                                            return Some(result);
                                        }
                                    }
                                    // Handle attributes: "builtin({self}.attr)" -> "obj.attr"
                                    else if inner_part.ends_with(")") {
                                        if let Some(attr_name) = inner_part.strip_suffix(")") {
                                            let result = format!("{}.{}", obj_str, attr_name);
                                            return Some(result);
                                        }
                                    }
                                }
                            }
                        }

                        // Fallback to direct patterns: "{self}.method()" or "self.method()"
                        let method_expr = if replacement_expr.starts_with("{self}.") {
                            replacement_expr.strip_prefix("{self}.")
                        } else {
                            replacement_expr.strip_prefix("self.")
                        };

                        if let Some(method_part) = method_expr {
                            if let Some(method_name) = method_part.strip_suffix("()") {
                                let result = format!("{}.{}()", obj_str, method_name);
                                return Some(result);
                            }
                        }

                        // Handle function calls with {self} as parameter: "func({self})" -> "func(obj)"
                        if replacement_expr.contains("({self})") {
                            let mut result =
                                replacement_expr.replace("({self})", &format!("({})", obj_str));
                            // Strip module prefix if present
                            if result.starts_with(&format!("{}.", self.module_name)) {
                                result = result[self.module_name.len() + 1..].to_string();
                            }
                            return Some(result);
                        }

                        // Handle function calls with self as parameter: "func(self)" -> "func(obj)"
                        if replacement_expr.contains("(self)") {
                            let result =
                                replacement_expr.replace("(self)", &format!("({})", obj_str));
                            return Some(result);
                        }

                        // Handle simple expressions like "True", "False", constants, etc.
                        // If the replacement doesn't contain placeholders or method calls, use it directly
                        if !replacement_expr.contains("{") && !replacement_expr.contains("(") {
                            // This is a simple expression (constant, variable name, etc.)
                            return Some(replacement_expr.clone());
                        }
                    }
                }
            }
        }

        None
    }

    fn should_apply_magic_method_replacement(&self, obj_expr: &AstExpr, _class_name: &str) -> bool {
        // Basic heuristic to determine if we should apply a magic method replacement
        match obj_expr {
            // Simple names like 'obj' are likely candidates for replacement
            AstExpr::Name(_) => true,

            // Complex expressions like obj.attr, obj[key], func() should NOT be replaced
            // since they're likely to be of different types
            AstExpr::Attribute(_) => false,
            AstExpr::Subscript(_) => false,
            AstExpr::Call(_) => false,

            // TODO: We could implement more sophisticated type checking here
            // using the type introspection context to actually determine the type
            // of the expression and check if it matches class_name

            // For now, be conservative and only replace simple names
            _ => false,
        }
    }

    fn expr_to_source(&self, expr: &AstExpr) -> String {
        // Always reconstruct from AST to ensure proper escaping and formatting
        self.reconstruct_expr(expr)
    }

    #[allow(clippy::only_used_in_recursion)] // self is needed for recursive method calls
    fn reconstruct_expr(&self, expr: &AstExpr) -> String {
        match expr {
            AstExpr::Constant(c) => match &c.value {
                rustpython_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>()
                    )
                }
                rustpython_ast::Constant::Int(i) => i.to_string(),
                rustpython_ast::Constant::Float(f) => f.to_string(),
                rustpython_ast::Constant::Bool(b) => {
                    if *b {
                        "True".to_string()
                    } else {
                        "False".to_string()
                    }
                }
                rustpython_ast::Constant::None => "None".to_string(),
                rustpython_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)
                }
                rustpython_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)
                    }
                }
                rustpython_ast::Constant::Ellipsis => "...".to_string(),
                _ => "...".to_string(),
            },
            AstExpr::Name(name) => name.id.to_string(),
            AstExpr::Attribute(attr) => {
                format!("{}.{}", self.reconstruct_expr(&attr.value), attr.attr)
            }
            AstExpr::Subscript(subscript) => {
                format!(
                    "{}[{}]",
                    self.reconstruct_expr(&subscript.value),
                    self.reconstruct_expr(&subscript.slice)
                )
            }
            AstExpr::Slice(slice) => {
                let lower = slice
                    .lower
                    .as_ref()
                    .map(|e| self.reconstruct_expr(e))
                    .unwrap_or_default();
                let upper = slice
                    .upper
                    .as_ref()
                    .map(|e| self.reconstruct_expr(e))
                    .unwrap_or_default();
                let step = slice
                    .step
                    .as_ref()
                    .map(|e| format!(":{}", self.reconstruct_expr(e)))
                    .unwrap_or_default();
                format!("{}:{}{}", lower, upper, step)
            }
            AstExpr::BinOp(binop) => {
                let left = self.reconstruct_expr(&binop.left);
                let right = self.reconstruct_expr(&binop.right);
                let op = match binop.op {
                    rustpython_ast::Operator::Add => "+",
                    rustpython_ast::Operator::Sub => "-",
                    rustpython_ast::Operator::Mult => "*",
                    rustpython_ast::Operator::MatMult => "@",
                    rustpython_ast::Operator::Div => "/",
                    rustpython_ast::Operator::Mod => "%",
                    rustpython_ast::Operator::Pow => "**",
                    rustpython_ast::Operator::LShift => "<<",
                    rustpython_ast::Operator::RShift => ">>",
                    rustpython_ast::Operator::BitOr => "|",
                    rustpython_ast::Operator::BitXor => "^",
                    rustpython_ast::Operator::BitAnd => "&",
                    rustpython_ast::Operator::FloorDiv => "//",
                };
                format!("{} {} {}", left, op, right)
            }
            AstExpr::UnaryOp(unaryop) => {
                let operand = self.reconstruct_expr(&unaryop.operand);
                let op = match unaryop.op {
                    rustpython_ast::UnaryOp::Invert => "~",
                    rustpython_ast::UnaryOp::Not => "not ",
                    rustpython_ast::UnaryOp::UAdd => "+",
                    rustpython_ast::UnaryOp::USub => "-",
                };
                format!("{}{}", op, operand)
            }
            AstExpr::BoolOp(boolop) => {
                let op = match boolop.op {
                    rustpython_ast::BoolOp::And => " and ",
                    rustpython_ast::BoolOp::Or => " or ",
                };
                boolop
                    .values
                    .iter()
                    .map(|v| self.reconstruct_expr(v))
                    .collect::<Vec<_>>()
                    .join(op)
            }
            AstExpr::Compare(compare) => {
                let mut result = self.reconstruct_expr(&compare.left);
                for (op, comparator) in compare.ops.iter().zip(&compare.comparators) {
                    let op_str = match op {
                        rustpython_ast::CmpOp::Eq => " == ",
                        rustpython_ast::CmpOp::NotEq => " != ",
                        rustpython_ast::CmpOp::Lt => " < ",
                        rustpython_ast::CmpOp::LtE => " <= ",
                        rustpython_ast::CmpOp::Gt => " > ",
                        rustpython_ast::CmpOp::GtE => " >= ",
                        rustpython_ast::CmpOp::Is => " is ",
                        rustpython_ast::CmpOp::IsNot => " is not ",
                        rustpython_ast::CmpOp::In => " in ",
                        rustpython_ast::CmpOp::NotIn => " not in ",
                    };
                    result.push_str(op_str);
                    result.push_str(&self.reconstruct_expr(comparator));
                }
                result
            }
            AstExpr::Call(call) => {
                let func = self.reconstruct_expr(&call.func);
                // Combine positional and keyword arguments using iterators
                let args: Vec<String> = call
                    .args
                    .iter()
                    .map(|arg| self.reconstruct_expr(arg))
                    .chain(call.keywords.iter().map(|keyword| {
                        if let Some(arg_name) = &keyword.arg {
                            format!("{}={}", arg_name, self.reconstruct_expr(&keyword.value))
                        } else {
                            format!("**{}", self.reconstruct_expr(&keyword.value))
                        }
                    }))
                    .collect();

                format!("{}({})", func, args.join(", "))
            }
            AstExpr::List(list) => {
                let elements: Vec<String> =
                    list.elts.iter().map(|e| self.reconstruct_expr(e)).collect();
                format!("[{}]", elements.join(", "))
            }
            AstExpr::Tuple(tuple) => {
                let elements: Vec<String> = tuple
                    .elts
                    .iter()
                    .map(|e| self.reconstruct_expr(e))
                    .collect();
                if elements.len() == 1 {
                    format!("({},)", elements[0])
                } else {
                    format!("({})", elements.join(", "))
                }
            }
            AstExpr::Dict(dict) => {
                let items: Vec<String> = dict
                    .keys
                    .iter()
                    .zip(&dict.values)
                    .map(|(key, value)| {
                        if let Some(k) = key {
                            format!(
                                "{}: {}",
                                self.reconstruct_expr(k),
                                self.reconstruct_expr(value)
                            )
                        } else {
                            format!("**{}", self.reconstruct_expr(value))
                        }
                    })
                    .collect();
                format!("{{{}}}", items.join(", "))
            }
            AstExpr::Set(set) => {
                let elements: Vec<String> =
                    set.elts.iter().map(|e| self.reconstruct_expr(e)).collect();
                if elements.is_empty() {
                    "set()".to_string()
                } else {
                    format!("{{{}}}", elements.join(", "))
                }
            }
            AstExpr::ListComp(comp) => {
                let elt = self.reconstruct_expr(&comp.elt);

                // Handle all generators (for multiple 'for' clauses)
                let mut generators = Vec::new();
                for gen in &comp.generators {
                    let target = self.reconstruct_expr(&gen.target);
                    let iter = self.reconstruct_expr(&gen.iter);
                    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.reconstruct_expr(if_expr))
                            .collect();
                        gen_str.push_str(&format!(" if {}", conds.join(" if ")));
                    }

                    generators.push(gen_str);
                }

                format!("[{} {}]", elt, generators.join(" "))
            }
            AstExpr::SetComp(comp) => {
                let conditions = if comp.generators[0].ifs.is_empty() {
                    String::new()
                } else {
                    let conds: Vec<String> = comp.generators[0]
                        .ifs
                        .iter()
                        .map(|if_expr| self.reconstruct_expr(if_expr))
                        .collect();
                    format!(" if {}", conds.join(" if "))
                };
                format!(
                    "{{{} for {} in {}{}}}",
                    self.reconstruct_expr(&comp.elt),
                    self.reconstruct_expr(&comp.generators[0].target),
                    self.reconstruct_expr(&comp.generators[0].iter),
                    conditions
                )
            }
            AstExpr::DictComp(comp) => {
                let conditions = if comp.generators[0].ifs.is_empty() {
                    String::new()
                } else {
                    let conds: Vec<String> = comp.generators[0]
                        .ifs
                        .iter()
                        .map(|if_expr| self.reconstruct_expr(if_expr))
                        .collect();
                    format!(" if {}", conds.join(" if "))
                };
                format!(
                    "{{{}: {} for {} in {}{}}}",
                    self.reconstruct_expr(&comp.key),
                    self.reconstruct_expr(&comp.value),
                    self.reconstruct_expr(&comp.generators[0].target),
                    self.reconstruct_expr(&comp.generators[0].iter),
                    conditions
                )
            }
            AstExpr::GeneratorExp(gen) => {
                let conditions = if gen.generators[0].ifs.is_empty() {
                    String::new()
                } else {
                    let conds: Vec<String> = gen.generators[0]
                        .ifs
                        .iter()
                        .map(|if_expr| self.reconstruct_expr(if_expr))
                        .collect();
                    format!(" if {}", conds.join(" if "))
                };
                format!(
                    "({} for {} in {}{})",
                    self.reconstruct_expr(&gen.elt),
                    self.reconstruct_expr(&gen.generators[0].target),
                    self.reconstruct_expr(&gen.generators[0].iter),
                    conditions
                )
            }
            AstExpr::IfExp(ifexp) => {
                format!(
                    "{} if {} else {}",
                    self.reconstruct_expr(&ifexp.body),
                    self.reconstruct_expr(&ifexp.test),
                    self.reconstruct_expr(&ifexp.orelse)
                )
            }
            AstExpr::Lambda(lambda) => {
                let args: Vec<String> = lambda
                    .args
                    .args
                    .iter()
                    .map(|arg| arg.def.arg.to_string())
                    .collect();
                format!(
                    "lambda {}: {}",
                    args.join(", "),
                    self.reconstruct_expr(&lambda.body)
                )
            }
            AstExpr::Starred(starred) => {
                format!("*{}", self.reconstruct_expr(&starred.value))
            }
            AstExpr::Yield(yield_expr) => {
                if let Some(value) = &yield_expr.value {
                    format!("yield {}", self.reconstruct_expr(value))
                } else {
                    "yield".to_string()
                }
            }
            AstExpr::YieldFrom(yieldfrom) => {
                format!("yield from {}", self.reconstruct_expr(&yieldfrom.value))
            }
            AstExpr::Await(await_expr) => {
                format!("await {}", self.reconstruct_expr(&await_expr.value))
            }
            AstExpr::JoinedStr(joined) => {
                let mut result = String::from("f\"");
                for value in &joined.values {
                    match value {
                        AstExpr::Constant(c) => {
                            if let rustpython_ast::Constant::Str(s) = &c.value {
                                result.push_str(s);
                            }
                        }
                        AstExpr::FormattedValue(fmt) => {
                            result.push('{');
                            result.push_str(&self.reconstruct_expr(&fmt.value));
                            match fmt.conversion {
                                rustpython_ast::ConversionFlag::Str => result.push_str("!s"),
                                rustpython_ast::ConversionFlag::Repr => result.push_str("!r"),
                                rustpython_ast::ConversionFlag::Ascii => result.push_str("!a"),
                                rustpython_ast::ConversionFlag::None => {}
                            }
                            if let Some(spec) = &fmt.format_spec {
                                result.push(':');
                                match &**spec {
                                    AstExpr::Constant(c) => {
                                        if let rustpython_ast::Constant::Str(s) = &c.value {
                                            result.push_str(s);
                                        }
                                    }
                                    AstExpr::JoinedStr(_) => {
                                        // For f-strings inside format specs, we need special handling
                                        // This shouldn't normally happen in valid Python
                                        let reconstructed = self.reconstruct_expr(spec);
                                        // 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.reconstruct_expr(spec)),
                                }
                            }
                            result.push('}');
                        }
                        _ => {
                            result.push('{');
                            result.push_str(&self.reconstruct_expr(value));
                            result.push('}');
                        }
                    }
                }
                result.push('"');
                result
            }
            AstExpr::FormattedValue(fmt) => {
                // This should normally be handled within JoinedStr
                let mut result = String::new();
                result.push('{');
                result.push_str(&self.reconstruct_expr(&fmt.value));
                match fmt.conversion {
                    rustpython_ast::ConversionFlag::Str => result.push_str("!s"),
                    rustpython_ast::ConversionFlag::Repr => result.push_str("!r"),
                    rustpython_ast::ConversionFlag::Ascii => result.push_str("!a"),
                    rustpython_ast::ConversionFlag::None => {}
                }
                if let Some(spec) = &fmt.format_spec {
                    result.push(':');
                    match &**spec {
                        AstExpr::Constant(c) => {
                            if let rustpython_ast::Constant::Str(s) = &c.value {
                                result.push_str(s);
                            }
                        }
                        _ => result.push_str(&self.reconstruct_expr(spec)),
                    }
                }
                result.push('}');
                result
            }
            AstExpr::NamedExpr(named) => {
                format!(
                    "{} := {}",
                    self.reconstruct_expr(&named.target),
                    self.reconstruct_expr(&named.value)
                )
            }
        }
    }
}

/// Apply replacements to source code preserving formatting
pub fn apply_replacements(source: &str, mut replacements: Vec<(Range<usize>, String)>) -> String {
    // Sort replacements by start position (reverse order for applying)
    replacements.sort_by_key(|(range, _)| std::cmp::Reverse(range.start));

    let mut result = source.to_string();

    for (range, replacement) in replacements {
        let original_text = &source[range.clone()];
        tracing::debug!(
            "Applying replacement at {:?}: '{}' -> '{}'",
            range,
            original_text,
            replacement
        );
        result.replace_range(range, &replacement);
    }

    result
}

/// Main entry point for migrating a file using Ruff parser
pub fn migrate_file_with_ruff(
    source: &str,
    module_name: &str,
    file_path: String,
    type_introspection: TypeIntrospectionMethod,
) -> Result<String> {
    // Use the main migrate_file function from migrate_stub
    use crate::type_introspection_context::TypeIntrospectionContext;
    use std::path::Path;

    let mut type_context = TypeIntrospectionContext::new(type_introspection)?;

    let result = crate::migrate_stub::migrate_file(
        source,
        module_name,
        Path::new(&file_path),
        &mut type_context,
        std::collections::HashMap::new(), // No predefined replacements
        std::collections::HashMap::new(), // No dependency inheritance map
    );

    type_context.shutdown()?;
    result
}

/// Compatibility function for import tracking tests
/// This provides the same interface as the old ruff_parser_improved module
pub fn migrate_file_with_improved_ruff(
    source: &str,
    module_name: &str,
    file_path: String,
    type_introspection: TypeIntrospectionMethod,
) -> Result<String> {
    // For now, delegate to the standard migration function
    // In the future, this could include additional import tracking functionality
    migrate_file_with_ruff(source, module_name, file_path, type_introspection)
}

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

    #[test]
    fn test_parse_simple() {
        let source = "x = 1\ny = 2";
        let module = PythonModule::parse(source).unwrap();
        assert_eq!(module.ast().len(), 2);
    }

    #[test]
    fn test_position_mapping() {
        let source = "x = 1\ny = 2";
        let module = PythonModule::parse(source).unwrap();

        // First line, first column
        assert_eq!(module.offset_to_position(0), Some((1, 0)));

        // Second line start
        assert_eq!(module.offset_to_position(6), Some((2, 0)));
    }
}