"""Documentation generation tools for AI agent integration."""

from langchain_core.tools import tool

from ieum.docs import CodeExtractor, DocGenerator, DocStyle


@tool
def generate_docstring_tool(
    file_path: str,
    function_name: str,
    style: str = "google"
) -> str:
    """Generate a docstring for a specific function or method.

    Args:
        file_path: Path to the Python file containing the function.
        function_name: Name of the function to generate docstring for.
        style: Docstring style - 'google', 'numpy', or 'sphinx' (default: google).

    Returns:
        Generated docstring in the specified style.
    """
    extractor = CodeExtractor()
    module = extractor.extract_module(file_path)

    # Find the function
    target_func = None

    # Search in module-level functions
    for func in module.functions:
        if func.name == function_name:
            target_func = func
            break

    # Search in class methods
    if not target_func:
        for cls in module.classes:
            for method in cls.methods:
                if method.name == function_name:
                    target_func = method
                    break
            if target_func:
                break

    if not target_func:
        return f"Function '{function_name}' not found in {file_path}"

    generator = DocGenerator()
    style_enum = DocStyle(style.lower())
    docstring = generator.generate_docstring(target_func, style_enum)

    return docstring


@tool
def generate_readme_tool(project_path: str, include_api: bool = True) -> str:
    """Generate a README.md file for a project.

    Args:
        project_path: Path to the project root directory.
        include_api: Whether to include API documentation section (default: True).

    Returns:
        Generated README.md content in Markdown format.
    """
    generator = DocGenerator()
    readme = generator.generate_readme(project_path, include_api=include_api)
    return readme


@tool
def generate_api_docs_tool(module_path: str, format: str = "markdown") -> str:
    """Generate API documentation for a Python module.

    Args:
        module_path: Path to the Python module file.
        format: Output format - 'markdown' or 'rst' (default: markdown).

    Returns:
        Generated API documentation in the specified format.
    """
    generator = DocGenerator()
    docs = generator.generate_api_docs(module_path, format=format)
    return docs
