// src/main.rs

use std::{
    fs,
    path::{Path, PathBuf},
};
use tree_sitter::{Parser, TreeCursor};
use walkdir::{DirEntry, WalkDir};

fn is_bad_entry(entry: &DirEntry) -> bool {
    entry
        .path()
        .components()
        .map(|comp| comp.as_os_str().to_str().unwrap())
        .any(|comp| matches!(comp, ".git" | "target"))
}

fn fix_file(path: &PathBuf, parser: &mut Parser) {
    let mut source_code = fs::read_to_string(path).unwrap();

    let tree = parser.parse(&source_code, None).unwrap();
    let root = tree.root_node();
    let mut cursor = root.walk();

    let mut edits = vec![];

    walk(&mut cursor, source_code.as_bytes(), &mut edits);

    println!("Checking: {}\n", path.display());
    if !edits.is_empty() {
        println!("Found {} edits", edits.len());

        for (start, end) in edits.into_iter().rev() {
            source_code.replace_range(start..end, "");
        }

        let _ = fs::write(path, source_code).unwrap();
    }
}

fn walk(cursor: &mut TreeCursor, source_code_bytes: &[u8], edits: &mut Vec<(usize, usize)>) {
    loop {
        let node = cursor.node();

        if node.kind() == "call_expression" {
            if let Some(func_expr) = node.child_by_field_name("function") {
                if func_expr.kind() == "field_expression" {
                    if let Some(method) = func_expr.child_by_field_name("field") {
                        if matches!(
                            method.utf8_text(source_code_bytes).unwrap(),
                            "new_attribute_error"
                                | "new_buffer_error"
                                | "new_eof_error"
                                | "new_errno_error"
                                | "new_import_error"
                                | "new_index_error"
                                | "new_lookup_error"
                                | "new_memory_error"
                                | "new_name_error"
                                | "new_not_implemented_error"
                                | "new_overflow_error"
                                | "new_recursion_error"
                                | "new_runtime_error"
                                | "new_system_error"
                                | "new_type_error"
                                | "new_unicode_decode_error"
                                | "new_unicode_encode_error"
                                | "new_value_error"
                                | "new_zero_division_error"
                        ) {
                            if let Some(args) = node.child_by_field_name("arguments") {
                                if args.named_child_count() >= 1 {
                                    let first_arg = args.named_child(0).unwrap();
                                    if first_arg.kind() == "call_expression" {
                                        if let Some(inner_func) =
                                            first_arg.child_by_field_name("function")
                                        {
                                            if inner_func.kind() == "field_expression" {
                                                let str_lit =
                                                    inner_func.child_by_field_name("value");
                                                let str_method =
                                                    inner_func.child_by_field_name("field");

                                                if let (Some(val), Some(method)) =
                                                    (str_lit, str_method)
                                                {
                                                    if val.kind() == "string_literal" {
                                                        let method_name = method
                                                            .utf8_text(source_code_bytes)
                                                            .unwrap();
                                                        if matches!(
                                                            method_name,
                                                            "to_owned" | "to_string"
                                                        ) {
                                                            let delete_start = val.end_byte();
                                                            let delete_end = first_arg.end_byte();
                                                            edits.push((delete_start, delete_end));
                                                        }
                                                    }
                                                }
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        if cursor.goto_first_child() {
            walk(cursor, source_code_bytes, edits);
            cursor.goto_parent();
        }

        if !cursor.goto_next_sibling() {
            break;
        }
    }
}

fn main() {
    let mut parser = Parser::new();
    parser
        .set_language(&tree_sitter_rust::LANGUAGE.into())
        .unwrap();
    let target_dir = Path::new(env!("CARGO_MANIFEST_DIR"))
        .join("../RustPython")
        .canonicalize()
        .expect("Failed to resolve path");

    WalkDir::new(&target_dir)
        .into_iter()
        .filter_entry(|e| !is_bad_entry(e))
        .filter_map(|e| e.ok())
        .map(|e| e.into_path())
        .filter(|p| p.is_file() && p.extension().map_or(false, |ext| ext == "rs"))
        .for_each(|path| fix_file(&path, &mut parser));
}