import ast
import importlib
import inspect
import subprocess
import sys

from reflex_ai import paths
from reflex_ai.scripts import get_module_script


def import_module_from_scratch(module_name: str, root_dir: str, scratch_dir: str):
    """Import the module from the scratchpad.

    Args:
        module_name: The name of the module to import.
        root_dir: The root directory of the module.
        scratch_dir: The scratch directory of the module.

    Returns:
        The imported module.
    """
    module = importlib.import_module(module_name)
    module_file = inspect.getsourcefile(module)
    module_tmp_file = module_file.replace(root_dir, scratch_dir)
    spec = importlib.util.spec_from_file_location(
        module.__name__,
        module_tmp_file,
    )
    new_module = importlib.util.module_from_spec(spec)
    # spec.loader.exec_module(new_module)
    return new_module


def get_module_source(module_name: str) -> str:
    """Get the file path of a module in a subprocess.

    Args:
        module_name: The name of the module.

    Returns:
        The source code of the module.
    """
    path = inspect.getfile(get_module_script)
    result = subprocess.run(
        [
            sys.executable,
            path,
            module_name,
            str(paths.base_paths[0].parent),
            str(paths.tmp_root_path),
        ],
        capture_output=True,
        text=True,
        check=True,
    )
    # Capture the output
    module_file = result.stdout.strip().split("\n")[-1]
    with open(module_file) as f:
        return f.read()


def get_module_ast(module_name: str) -> ast.Module:
    """Get the AST of a module.

    Args:
        module_name: The name of the module.

    Returns:
        The AST of the module.
    """
    # Import the module
    module_source = get_module_source(module_name)

    # Parse the module source code into an AST
    module_ast = ast.parse(module_source)
    return module_ast


def find_node(tree, element_type: str, element_name: str) -> ast.AST | None:
    for node in ast.walk(tree):
        if (
            isinstance(node, ast.FunctionDef)
            and element_type == "function"
            and node.name == element_name
        ):
            return node
        elif (
            isinstance(node, ast.ClassDef)
            and element_type == "class"
            and node.name == element_name
        ):
            return node
        elif isinstance(node, ast.Assign) and element_type == "variable":
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == element_name:
                    return node
        elif isinstance(node, ast.Return) and element_type == "return":
            return node
    return None
