from __future__ import annotations
from typing import Any, Optional

import ast
import importlib.util
from pathlib import Path

from sieve_utils.collections import set_nested_attr


def file_has_func(file_path: Path, func_name: str) -> bool:
    """Check if a Python file has a function with a specific name.

    Args:
        file_path (Path): Path to the Python file.
        func_name (str): Name of the function to check for.

    Returns:
        bool: True if the function exists, False otherwise.
    """
    with open(file_path, 'r') as file:
        tree = ast.parse(file.read())

    return any(
        isinstance(node, ast.FunctionDef) and node.name == func_name
        for node in tree.body
    )


def find_scripts_in_package(
        package_dir: str,
        files_to_ignore: Optional[list[str]] = None,
        script_func_name: str = "main"
) -> dict[str, Any]:
    """Recursively find all Python scripts in a package and return them in a format ready to pass to `fire.Fire()`.

    Args:
        package_dir (str): Path to the package directory.
        files_to_ignore (Optional[list[str]], optional): List of file paths to ignore. Defaults to None.
        script_func_name (str, optional): Name of the function to treat a script's entrypoint. Defaults to "main".

    Returns:
        dict: Dictionary of commands where the keys are the top-level command names and the values are objects (
        functions, classes, etc.) that are compatible with `fire.Fire()`.

    Example:
        Given the following package structure:
        ```
        package/
        ├── __init__.py
        ├── module1.py
        ├── module2.py
        └── subpackage1/
            ├── __init__.py
            ├── module3.py
            ├── module4.py
            └── subpackage2/
                ├── __init__.py
                └── module5.py
        ```
        The following commands dictionary will be returned:
        ```
        {
            "module1": module1.main,
            "module2": module2.main,
            "subpackage1": {
                "module3": module3.main,
                "module4": module4.main,
                "subpackage2": {
                    "module5": module5.main
                }
            }
        }
        Now you can pass this dictionary to `fire.Fire()` to create a CLI for the package. For example:
        ```
        import fire

        commands = find_scripts_in_package("package")
        fire.Fire(commands)
        ```

        This will create a CLI with the following structure:
        ```
        Usage: cli.py <COMMAND> [ARGS]

        Available commands:
            - module1
            - module2
            - subpackage1
            ```
        You can then run a specific script by name:
        ```
        python cli.py module1
        ```
        Scripts nested in subpackages can be accessed intuitively:
        ```
        python cli.py subpackage1 module3
        ```
        This works for arbitrarily nested scripts:
        ```
        python cli.py subpackage1 subpackage2 module5
        ```
    """
    package_dir = Path(package_dir)
    package_name = package_dir.name

    if files_to_ignore is None:
        files_to_ignore = []

    commands = {}

    for script in package_dir.rglob("*.py"):
        # Skip private and init files
        if script.name.startswith("_") or script.stem == "__init__":
            continue  # Skip private and init files

        # Skip files to ignore
        for file_to_ignore in files_to_ignore:
            if script.resolve() == Path(file_to_ignore).resolve():
                continue

        # Skip files without a main function
        if not file_has_func(script, script_func_name):
            continue

        # Import module
        relative_path = script.relative_to(package_dir).with_suffix("")
        module_name = ".".join([package_name] + list(relative_path.parts))
        try:
            module = importlib.import_module(module_name)
        except ModuleNotFoundError as e:
            print(f"Error: Could not import module {module_name}. {e}")
            continue
        except Exception as e:
            print(f"Unexpected error while importing {module_name}: {e}")
            continue

        # Add main function of module to the commands dictionary
        top_level_command = relative_path.parts[0]
        if len(relative_path.parts) == 1:
            commands[top_level_command] = module.main
        else:
            if top_level_command not in commands:
                commands[top_level_command] = type(top_level_command, (object,), {})()
            set_nested_attr(commands[top_level_command], relative_path.parts[1:], module.main)

    return commands
