"""
Generate `aegis/stub.py` from `_aegis/full_stub.py`.

This script extracts only the necessary function stubs with their docstrings
from the full implementation file `_aegis/full_stub.py`. It filters functions
based on enabled feature flags and writes the resulting stubs into
`aegis/stub.py`.

Designed to be run manually when updating the public-facing stub interface
based on feature flags or function changes.
"""

import ast
from importlib import resources
from pathlib import Path

from .aegis_config import has_feature

# Feature flags to control which functions get included in the stub
IS_PREDICTIONS_ENABLED = has_feature("ALLOW_AGENT_PREDICTIONS")
IS_MESSAGES_ENABLED = has_feature("ALLOW_AGENT_MESSAGES")
IS_DYNAMIC_SPAWNING_ENABLED = has_feature("ALLOW_DYNAMIC_SPAWNING")
IS_ABILITIES_ENABLED = has_feature("ALLOW_AGENT_ABILITIES")

# Grouped function names by feature flag
PREDICT_FUNCTIONS = {"predict", "read_pending_predictions"}
MESSAGE_FUNCTINS = {"read_messages", "send_message"}
SPAWN_FUNCTIONS = {"get_spawns", "spawn_agent"}
ABILITIES = {"drone_scan"}


def should_include_function(name: str) -> bool:
    """
    Determine whether a function should be included in the stub output.

    Filters based on the feature flags and which group the function belongs to.

    Args:
        name (str): The name of the function.

    Returns:
        bool: True if the function should be included, False otherwise.

    """
    if name in PREDICT_FUNCTIONS and not IS_PREDICTIONS_ENABLED:
        return False
    if name in MESSAGE_FUNCTINS and not IS_MESSAGES_ENABLED:
        return False
    if name in SPAWN_FUNCTIONS and not IS_DYNAMIC_SPAWNING_ENABLED:
        return False
    return not (name in ABILITIES and not IS_ABILITIES_ENABLED)


def build_header() -> str:
    """
    Generate the import/header section for the stub file.

    Returns:
        str: The header string to be prepended to the stub.

    """
    imports: list[str] = []
    if IS_PREDICTIONS_ENABLED:
        imports.append("import numpy as np")
        imports.append("from numpy.typing import NDArray")

    static_imports = """
from . import (
    CellInfo,
    Direction,
    Location,
    Message,
    Rubble,
    Survivor,
)
"""
    optional_imports = "\n".join(imports)

    return f'''"""
Autogenerated from `_aegis/full_stub.py`.

Do not modify manually.
"""
# ruff: noqa: F401
# pyright: reportReturnType=false
# pyright: reportUnusedImport=false
# pyright: reportUnusedParameter=false
{optional_imports}
{static_imports}
'''


def format_return(returns: ast.expr | None) -> str:
    """
    Format the return type annotation of a function.

    Args:
        returns (ast.expr | None): The return type node from the AST.

    Returns:
        str: The return type as a string, or empty if not specified.

    """
    if returns is None:
        return ""
    return ast.unparse(returns)


def format_args(args: ast.arguments) -> str:
    """
    Format the argument list of a function into a string.

    Args:
        args (ast.arguments): The argument list node from the AST.

    Returns:
        str: A string representation of the function arguments.

    """
    res: list[str] = []

    total_args = args.args
    defaults = args.defaults
    num_defaults = len(defaults)
    num_args = len(total_args)

    for i, arg in enumerate(total_args):
        has_default = i >= num_args - num_defaults
        default_value = defaults[i - (num_args - num_defaults)] if has_default else None

        arg_str = arg.arg
        if arg.annotation:
            arg_str += f": {ast.unparse(arg.annotation)}"
        if default_value is not None:
            arg_str += f" = {ast.unparse(default_value)}"

        res.append(arg_str)

    # *args
    if args.vararg:
        vararg_str = f"*{args.vararg.arg}"
        if args.vararg.annotation:
            vararg_str += f": {ast.unparse(args.vararg.annotation)}"
        res.append(vararg_str)

    return ", ".join(res)


def extract_stub_code(tree: ast.Module) -> str:
    """
    Extract function stubs with docstrings from a parsed AST.

    Args:
        tree (ast.Module): The AST of the full stub source.

    Returns:
        str: The body of the stub file with function definitions and docstrings.

    """
    lines: list[str] = []

    function_nodes = [
        node
        for node in tree.body
        if isinstance(node, ast.FunctionDef) and should_include_function(node.name)
    ]

    for i, node in enumerate(function_nodes):
        docstring = ast.get_docstring(node)
        signature = f"def {node.name}({format_args(node.args)}) -> {format_return(node.returns)}:"
        lines.append(signature)

        if docstring:
            if "\n" in docstring:
                lines.append(f'{" " * 4}"""')
                indented_lines = [
                    f"    {line}" if line.strip() else ""
                    for line in docstring.splitlines()
                ]
                lines.extend(indented_lines)
                lines.append('\n    """')
            else:
                lines.append(f'{" " * 4}"""{docstring}"""')

        if i != len(function_nodes) - 1:
            lines.append("\n")

    return "\n".join(lines)


def main() -> None:
    """
    Entry point for the stub generation script.

    Reads `_aegis/full_stub.py`, filters and formats functions based on feature flags,
    and writes the result to `aegis/stub.py`.

    Reports errors in case the source file cannot be read or the output file cannot be written.
    """
    try:
        stub_source = resources.read_text("_aegis", "full_stub.py")
    except FileNotFoundError:
        print("Error: `_aegis/full_stub.py` not found.")
        return
    except OSError as e:
        print(f"I/O error while reading full_stub.py: {e}")
        return

    tree = ast.parse(stub_source)
    header = build_header()
    stub_code = extract_stub_code(tree)

    content = header + "\n" + stub_code + "\n"

    output_path = Path(__file__).resolve().parent.parent / "aegis" / "stub.py"
    try:
        _ = output_path.write_text(content, encoding="utf-8")
        print(f"Successfully wrote stub to {output_path}")
    except (OSError, PermissionError) as e:
        print(f"Error writing stub.py: {e}")
