ast.rs

use oxc::{
    allocator::{Allocator, Box as OXCBox, Vec as OXCVec},
    ast::ast::{
        Argument, Expression, IdentifierName, IdentifierReference, NewExpression, ObjectExpression,
        ObjectProperty, ObjectPropertyKind, Program, PropertyKey, PropertyKind, Statement,
        VariableDeclarator,
    },
    codegen::Codegen,
    parser::{ParseOptions, Parser},
    span::{Atom, SourceType, Span},
};

use std::{cell::Cell, fs, path::Path};

pub fn source_to_ast<'a>(file_path: &str, allocator: &'a Allocator) -> Result<Program<'a>, String> {
    if !Path::new(file_path).exists() {
        return Err(format!("File does not exist"));
    }

    let source_text = fs::read_to_string(file_path)
        .map_err(|e| format!("Failed to read the JavaScript file: {}", e))?
        .into_boxed_str();
    let source_text: &'a str = Box::leak(source_text);

    let source_type = SourceType::from_path(file_path)
        .map_err(|e| format!("Failed to determine source type: {:?}", e))?;

    let parser = Parser::new(allocator, source_text, source_type).with_options(ParseOptions {
        parse_regular_expression: true,
        ..ParseOptions::default()
    });

    let parse_result = parser.parse();
    Ok(parse_result.program)
}

pub fn is_module_imported_from_ast<'a>(
    file_path: &str,
    module_name: &str,
    allocator: &'a Allocator,
) -> Result<bool, String> {
    let program = source_to_ast(file_path, allocator)?;

    for node in program.body {
        if let Statement::ImportDeclaration(import_decl) = node {
            if import_decl.source.value == module_name {
                return Ok(true);
            }
        }
    }

    Ok(false)
}

pub fn insert_import_to_ast<'a>(
    file_path: &str,
    import_lines: &str,
    allocator: &'a Allocator,
) -> Result<String, String> {
    let mut program = source_to_ast(file_path, allocator)?;

    for import_line in import_lines.lines() {
        let import_source_type = SourceType::default();
        let parser =
            Parser::new(allocator, import_line, import_source_type).with_options(ParseOptions {
                parse_regular_expression: true,
                ..ParseOptions::default()
            });

        let parsed_result = parser.parse();
        if let Some(errors) = parsed_result.errors.first() {
            return Err(format!("Failed to parse import line: {:?}", errors));
        }

        let new_import = parsed_result
            .program
            .body
            .into_iter()
            .find(|node| matches!(node, Statement::ImportDeclaration(_)))
            .ok_or_else(|| "No import declaration found in parsed import line".to_string())?;

        if program.body.iter().any(|node| {
            matches!(
                (node, &new_import),
                (
                    Statement::ImportDeclaration(existing_import),
                    Statement::ImportDeclaration(new_import_node)
                ) if existing_import.source.value == new_import_node.source.value
            )
        }) {
            continue; // Skip duplicate imports
        }

        let position = program
            .body
            .iter()
            .rposition(|node| matches!(node, Statement::ImportDeclaration(_)))
            .map(|index| index + 1)
            .unwrap_or(0);

        program.body.insert(position, new_import);
    }

    let codegen = Codegen::new();
    let generated_code = codegen.build(&program).code;

    Ok(generated_code)
}

pub fn remove_import_from_ast<'a>(
    file_path: &str,
    module_name: &str,
    allocator: &'a Allocator,
) -> Result<Program<'a>, String> {
    // Parse the source file into AST
    let mut program = source_to_ast(file_path, allocator)?;

    // Find and remove the specified import declaration
    let initial_length = program.body.len();

    program.body.retain(|node| {
        if let Statement::ImportDeclaration(import_decl) = node {
            import_decl.source.value != module_name
        } else {
            true
        }
    });

    // Check if any import was removed
    if program.body.len() == initial_length {
        return Err(format!(
            "Import module '{}' not found in the AST",
            module_name
        ));
    }

    // Re-generate the source text from the updated AST
    let codegen = Codegen::new();
    let updated_source = codegen.build(&program).code;

    // Update the source_text field in the program
    let updated_source: &'a str = allocator.alloc_str(&updated_source);
    program.source_text = updated_source;

    // Return the modified AST
    Ok(program)
}
// TODO: check is there LiveView hook in the js file

#[allow(dead_code)]
pub fn find_live_socket_node_from_ast<'a>(program: &'a Program<'a>) -> Result<bool, bool> {
    if program.body.iter().any(|node| {
        if let Statement::VariableDeclaration(var_decl) = node {
            var_decl.declarations.iter().any(|decl| {
                decl.id
                    .get_identifier()
                    .map_or(false, |ident| ident == "liveSocket")
            })
        } else {
            false
        }
    }) {
        Ok(true)
    } else {
        Err(false)
    }
}

#[allow(dead_code)]
pub fn extend_hook_object_to_ast<'a>(
    file_path: &str,
    allocator: &'a Allocator,
) -> Result<String, String> {
    let mut program = source_to_ast(file_path, allocator)?;

    let hooks_result = find_hooks_property(&mut program, allocator)?;
    if let Some(hooks_property) = hooks_result {
        if let Expression::ObjectExpression(obj_expr) = hooks_property {
            let new_property = ObjectPropertyKind::ObjectProperty(OXCBox::new_in(
                ObjectProperty {
                    span: Span::default(),
                    kind: PropertyKind::Init,
                    key: PropertyKey::StaticIdentifier(OXCBox::new_in(
                        IdentifierName {
                            span: Span::default(),
                            name: Atom::from("OXCTestHook"),
                        },
                        allocator,
                    )),
                    value: Expression::Identifier(OXCBox::new_in(
                        IdentifierReference {
                            span: Span::default(),
                            name: Atom::from("OXCTestHook"),
                            reference_id: Cell::new(None),
                        },
                        allocator,
                    )),
                    method: false,
                    shorthand: true,
                    computed: false,
                },
                allocator,
            ));

            obj_expr.properties.push(new_property);
        }
    } else {
        return Err("liveSocket not found in the AST".to_string());
    }

    let codegen = Codegen::new();
    let generated_code = codegen.build(&program).code;
    Ok(generated_code)
}

fn find_hooks_property<'short, 'long>(
    program: &'short mut Program<'long>,
    allocator: &'short Allocator,
) -> Result<Option<&'short mut Expression<'long>>, String> {
    let live_socket_found = program.body.iter_mut().any(|node| {
        if let Statement::VariableDeclaration(var_decl) = node {
            var_decl
                .declarations
                .iter()
                .any(|decl| is_target_variable(decl, "liveSocket"))
        } else {
            false
        }
    });

    if !live_socket_found {
        return Err("liveSocket not found.".to_string());
    }

    let result = program.body.iter_mut().find_map(|node| {
        let var_decl = match node {
            Statement::VariableDeclaration(var_decl) => var_decl,
            _ => return None,
        };

        var_decl.declarations.iter_mut().find_map(|decl| {
            let obj_expr = get_new_expression(decl)?;

            obj_expr.arguments.iter_mut().find_map(|arg| {
                let obj_expr_inner = get_object_expression(arg)?;
                let mut iter = obj_expr_inner.properties.iter_mut();
                if let Some(expr) = iter.find_map(|prop| get_property_by_key(prop, "hooks")) {
                    return Some(expr);
                }

                let new_hooks_property = ObjectPropertyKind::ObjectProperty(OXCBox::new_in(
                    ObjectProperty {
                        span: Span::default(),
                        kind: PropertyKind::Init,
                        key: PropertyKey::StaticIdentifier(OXCBox::new_in(
                            IdentifierName {
                                span: Span::default(),
                                name: Atom::from("hooks"),
                            },
                            allocator,
                        )),
                        value: Expression::ObjectExpression(OXCBox::new_in(
                            ObjectExpression {
                                span: Span::default(),
                                properties: OXCVec::new_in(allocator),
                                trailing_comma: None,
                            },
                            allocator,
                        )),
                        method: false,
                        shorthand: false,
                        computed: false,
                    },
                    allocator,
                ));

                iter.push(new_hooks_property)
                // .find_map(|prop| get_property_by_key(prop, "hooks"))
            })
        })
    });

    Ok(result)
}

fn is_target_variable(decl: &VariableDeclarator, name: &str) -> bool {
    decl.id
        .kind
        .get_binding_identifier()
        .map_or(false, |binding| binding.name == name)
}

fn get_new_expression<'short, 'long>(
    decl: &'short mut VariableDeclarator<'long>,
) -> Option<&'short mut NewExpression<'long>> {
    match decl.init.as_mut()? {
        Expression::NewExpression(expr) => Some(expr),
        _ => None,
    }
}

fn get_object_expression<'short, 'long>(
    arg: &'short mut Argument<'long>,
) -> Option<&'short mut ObjectExpression<'long>> {
    arg.as_expression_mut().and_then(|expr| match expr {
        Expression::ObjectExpression(boxed_obj_expr) => Some(boxed_obj_expr.as_mut()),
        _ => None,
    })
}

fn get_property_by_key<'short, 'long>(
    property: &'short mut ObjectPropertyKind<'long>,
    key_name: &str,
) -> Option<&'short mut Expression<'long>> {
    match property {
        ObjectPropertyKind::ObjectProperty(prop) => match &prop.key {
            PropertyKey::StaticIdentifier(key) if key.as_ref().name == key_name => {
                Some(&mut prop.value)
            }
            _ => None,
        },
        _ => None,
    }
}

// TODO: delete an object from hook
#[cfg(test)]
mod tests {
    use super::*;
    // use std::path::Path;

    fn create_allocator<'a>() -> &'a Allocator {
        let allocator = Box::new(Allocator::default());
        Box::leak(allocator)
    }

    #[test]
    fn test_parse_and_display_ast() {
        let test_file_path = "test_assets/test.js";

        // Ensure the test directory exists
        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }

        // Write a test JavaScript file
        let js_content = r#"
               import { foo } from 'module-name';
               import bar from 'another-module';

               console.log('Testing AST parsing');
           "#;

        fs::write(test_file_path, js_content).expect("Failed to write test file");

        // Check if the test file exists
        assert!(
            Path::new(test_file_path).exists(),
            "Test file does not exist"
        );

        // Test the AST parsing
        let allocator = create_allocator();

        match source_to_ast(test_file_path, allocator) {
            Ok(ast) => {
                println!("{:#?}", ast.body);
                assert!(!ast.body.is_empty(), "AST body should not be empty");
            }
            Err(e) => panic!("Error while parsing AST: {}", e),
        }

        // Cleanup: Remove the test file
        fs::remove_file(test_file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_is_module_imported_from_ast() {
        let test_file_path = "test_assets/sample.js";

        // Ensure the test file exists
        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }

        // Write a test JavaScript file
        let js_content = r#"
               import { foo } from 'module-name';
               import bar from 'another-module';
           "#;

        fs::write(test_file_path, js_content).expect("Failed to write test file");

        // Check if the test file exists
        assert!(
            Path::new(test_file_path).exists(),
            "Test file does not exist"
        );

        // Test the function with a valid module
        let allocator = create_allocator();
        match is_module_imported_from_ast(test_file_path, "module-name", allocator) {
            Ok(true) => println!("Module 'module-name' is imported as expected."),
            Ok(false) => panic!("Module 'module-name' should be imported but was not detected."),
            Err(e) => panic!("Error while checking module: {}", e),
        }

        // Test the function with another valid module
        match is_module_imported_from_ast(test_file_path, "another-module", allocator) {
            Ok(true) => println!("Module 'another-module' is imported as expected."),
            Ok(false) => panic!("Module 'another-module' should be imported but was not detected."),
            Err(e) => panic!("Error while checking module: {}", e),
        }

        // Test the function with a non-existent module
        match is_module_imported_from_ast(test_file_path, "non-existent-module", allocator) {
            Ok(true) => panic!("Module 'non-existent-module' should not be imported."),
            Ok(false) => println!("Module 'non-existent-module' is correctly not imported."),
            Err(e) => panic!("Error while checking module: {}", e),
        }

        // Cleanup: Remove the test file
        fs::remove_file(test_file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_insert_duplicate_import() {
        println!("Running test: test_insert_duplicate_import");
        let test_file_path = "test_assets/duplicate_import.js";

        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }

        let js_content = r#"
            import { foo } from "module-name";
            console.log("Duplicate import test");
        "#;

        fs::write(test_file_path, js_content).expect("Failed to write test file");

        let duplicate_import = r#"import { foo } from "module-name";"#;
        let allocator = create_allocator();
        let result = insert_import_to_ast(test_file_path, duplicate_import, allocator);

        match result {
            Ok(updated_content) => {
                println!("Updated Content:\n{}", updated_content);
                // Ensure the duplicate import is not added
                let import_count = updated_content.matches(duplicate_import).count();
                assert_eq!(
                    import_count, 1,
                    "Duplicate import should not be added, but it was found multiple times"
                );
            }
            Err(e) => panic!("Unexpected error: {}", e),
        }

        fs::remove_file(test_file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_insert_import_to_ast_with_existing_imports() {
        println!("Running test: test_insert_import_to_ast_with_existing_imports");
        let test_file_path = "test_assets/with_imports.js";

        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }

        let js_content = r#"
            import bar from "another-module";
            console.log("Some imports here!");
        "#;

        fs::write(test_file_path, js_content).expect("Failed to write test file");

        let new_import = r#"import { foo } from "module-name";"#;
        let allocator = create_allocator();
        let result = insert_import_to_ast(test_file_path, new_import, allocator);

        match result {
            Ok(updated_content) => {
                println!(
                    "Updated Content::test_insert_import_to_ast_with_existing_imports:::\n{}",
                    updated_content
                );
                let lines: Vec<&str> = updated_content.lines().collect();
                let last_import_position =
                    lines.iter().rposition(|&line| line.starts_with("import"));
                assert_eq!(
                    lines[last_import_position.unwrap() + 1],
                    "console.log(\"Some imports here!\");",
                    "New import should be added after the last import"
                );
            }
            Err(e) => panic!("Error while inserting import: {}", e),
        }

        fs::remove_file(test_file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_insert_multiple_imports() {
        println!("Running test: test_insert_multiple_imports");
        let test_file_path = "test_assets/multiple_imports.js";

        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }

        let js_content = r#"
            console.log("Starting with no imports!");
        "#;

        fs::write(test_file_path, js_content).expect("Failed to write test file");

        let imports = vec![
            r#"import { foo } from "module-one";"#,
            r#"import bar from "module-two";"#,
            r#"import * as namespace from "module-three";"#,
            r#"import something, { foo as bar } from "module-four";"#,
        ];

        let allocator = create_allocator();
        for import in &imports {
            let result = insert_import_to_ast(test_file_path, import, allocator);
            match result {
                Ok(updated_content) => {
                    println!(
                        "Updated Content::test_insert_multiple_imports::\n{}",
                        updated_content
                    );
                    fs::write(test_file_path, &updated_content).expect("Failed to update file");
                    assert!(
                        updated_content.contains(import),
                        "Import not added: {}",
                        import
                    );
                }
                Err(e) => panic!("Error while inserting import '{}': {}", import, e),
            }
        }

        let updated_content =
            fs::read_to_string(test_file_path).expect("Failed to read updated file");
        for import in &imports {
            assert!(
                updated_content.contains(import),
                "Final content missing import: {}",
                import
            );
        }

        fs::remove_file(test_file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_insert_import_to_ast_with_alert_only() {
        let test_file_path = "test_assets/alert_only.js";

        // Ensure the test directory exists
        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }

        // Write a test JavaScript file with only an alert
        let js_content = r#"
            alert('Hello, world!');
        "#;

        fs::write(test_file_path, js_content).expect("Failed to write test file");

        // Insert a new import
        let new_import = r#"import { foo } from "module-name";"#;
        let allocator = create_allocator();
        let result = insert_import_to_ast(test_file_path, new_import, allocator);

        match result {
            Ok(updated_content) => {
                println!("Updated Content:\n{}", updated_content);
                assert!(updated_content.contains(new_import), "New import not added");
                assert!(
                    updated_content.starts_with(new_import),
                    "New import should be at the top"
                );
            }
            Err(e) => panic!("Error while inserting import: {}", e),
        }

        // Cleanup: Remove the test file
        fs::remove_file(test_file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_remove_import_from_ast() {
        let test_file_path = "test_assets/remove_import.js";

        // Ensure the test directory exists
        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }

        // Write a test JavaScript file
        let js_content = r#"
                import { foo } from "module-name";
                import bar from "another-module";

                console.log("Testing remove import");
            "#;

        fs::write(test_file_path, js_content).expect("Failed to write test file");

        // Check if the test file exists
        assert!(
            Path::new(test_file_path).exists(),
            "Test file does not exist"
        );

        // Test the function to remove an existing module
        let allocator = create_allocator();
        match remove_import_from_ast(test_file_path, "module-name", allocator) {
            Ok(updated_ast) => {
                println!(
                    "Updated AST after removing 'module-name': {:#?}",
                    updated_ast.body
                );

                // Verify the removed module does not exist in the AST
                let module_exists = updated_ast.body.iter().any(|node| {
                    if let Statement::ImportDeclaration(import_decl) = node {
                        import_decl.source.value == "module-name"
                    } else {
                        false
                    }
                });

                assert!(
                    !module_exists,
                    "The module 'module-name' was not removed from the AST"
                );
            }
            Err(e) => panic!("Error while removing import: {}", e),
        }

        // Test the function with a non-existent module
        match remove_import_from_ast(test_file_path, "non-existent-module", allocator) {
            Ok(_) => panic!("Non-existent module should not return Ok"),
            Err(e) => {
                println!("Correctly handled non-existent module: {}", e);
                assert!(
                    e.contains("not found"),
                    "Error message does not indicate missing module"
                );
            }
        }

        // Cleanup: Remove the test file
        fs::remove_file(test_file_path).expect("Failed to remove test file");
    }

    fn setup_test_file(file_path: &str, content: &str) {
        if !Path::new("test_assets").exists() {
            fs::create_dir("test_assets").expect("Failed to create test_assets directory");
        }
        fs::write(file_path, content).expect("Failed to write test file");
    }

    #[test]
    fn test_find_live_socket_variable() {
        let file_path = "test_assets/test_live_socket.js";

        // Set up a test JavaScript file
        let js_content = r#"
            const someVar = 42;
            let liveSocket = new LiveSocket("/live", Socket, {
              hooks: { ...Hooks, CopyMixInstallationHook },
              longPollFallbackMs: 2500,
              params: { _csrf_token: csrfToken },
            });
            const anotherVar = "hello";
        "#;
        setup_test_file(file_path, js_content);

        let allocator = create_allocator();
        let program = source_to_ast(file_path, allocator).expect("Failed to parse AST");

        // Test the function
        let result = find_live_socket_node_from_ast(&program);
        println!("Result for test_find_live_socket_variable: {:?}", result);

        assert_eq!(result, Ok(true));

        // Cleanup
        fs::remove_file(file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_find_live_socket_variable_not_found() {
        let file_path = "test_assets/test_no_live_socket.js";

        // Set up a test JavaScript file
        let js_content = r#"
            const someVar = 42;
            const anotherVar = "hello";
        "#;
        setup_test_file(file_path, js_content);

        let allocator = create_allocator();
        let program = source_to_ast(file_path, allocator).expect("Failed to parse AST");

        // Test the function
        let result = find_live_socket_node_from_ast(&program);
        println!(
            "Result for test_find_live_socket_variable_not_found: {:?}",
            result
        );

        assert_eq!(result, Err(false));

        // Cleanup
        fs::remove_file(file_path).expect("Failed to remove test file");
    }

    #[test]
    fn test_extend_hook_object_to_ast() {
        let file_path = "/Users/shahryar/Documents/Programming/Elixir/mishka/assets/js/app.js";

        assert!(
            fs::metadata(file_path).is_ok(),
            "JavaScript file does not exist at the given path."
        );

        let allocator = create_allocator();
        match extend_hook_object_to_ast(file_path, allocator) {
            Ok(ast) => {
                println!("Hook object extended successfully. ==> {}", ast);
            }
            Err(e) => {
                eprintln!("Error: {}", e);
                panic!("Failed to extend hook object.");
            }
        }
    }
}