"""
@file: generate_mkdocs.py
@breif: Used to automatically generate MkDocs documentation for Python libraries.
@author: Wu Maojia
@update: 2025.10.3
"""
import os
import ast
import yaml
import shutil


def extract_classes(file_path: str):
    """
    Extract the names of all classes in a Python file.

    Parameters:
        file_path (str): Path to the Python file.

    Returns:
        class_names (list): List of class names.
    """
    class_names = []
    with open(file_path, 'r', encoding='utf-8') as f:
        node = ast.parse(f.read(), filename=file_path)
    for child in node.body:
        if isinstance(child, ast.ClassDef):
            class_names.append(child.name)
    return class_names


def generate_api_docs(root_folder: str, output_folder: str, index_file: str, mkdocs_file: str, ex_home_file: str, ex_assets_folder: str):
    """
    Automatically generate Markdown files for API documentation, update the homepage, and modify mkdocs.yml.

    Parameters:
        root_folder (str): Path to the root directory of the library.
        output_folder (str): Path to the output directory for the generated documentation.
        index_file (str): Path to the homepage file.
        mkdocs_file (str): Path to the mkdocs.yml file.
    """
    nav_structure = {}
    tutorials_structure = {}

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Process Python source code for API documentation
    for root, dirs, files in os.walk(root_folder):
        relative_root = os.path.relpath(root, root_folder)  # Get the relative path of the current directory
        for file in files:
            if file.endswith('.py') and not file.startswith('__'):
                file_path = os.path.join(root, file)
                class_names = extract_classes(file_path)
                if not class_names:
                    print(f"Warning: No class found in {file_path}. Skipping...")
                    continue

                module_path = file_path.replace('.py', '').replace('/', '.').replace('\\', '.')
                # Create a corresponding output directory that mirrors the input directory's structure
                output_dir = os.path.join(output_folder, 'api', relative_root, file.replace('.py', ''))
                os.makedirs(output_dir, exist_ok=True)

                # Generate Markdown files for each class
                for class_name in class_names:
                    class_md_path = os.path.join(output_dir, f"{class_name}.md")
                    with open(class_md_path, 'w', encoding='utf-8') as f:
                        f.write(f"# {class_name}\n\n::: {module_path}.{class_name}\n")

                # Update navigation structure
                current_nav = nav_structure
                for part in relative_root.split(os.sep):
                    if part:
                        current_nav = current_nav.setdefault(part, {})
                current_nav.setdefault(file.replace('.py', ''), []).extend([{
                    class_name: os.path.relpath(os.path.join(output_dir, f"{class_name}.md"), output_folder)
                } for class_name in class_names])

    # Process tutorials folder
    tutorials_folder = 'tutorials'
    if os.path.exists(tutorials_folder):
        # Copy tutorials folder to docs
        tutorials_output = os.path.join(output_folder, 'tutorials')
        if os.path.exists(tutorials_output):
            shutil.rmtree(tutorials_output)
        shutil.copytree(tutorials_folder, tutorials_output)
        
        # Build tutorials navigation structure
        for root, dirs, files in os.walk(tutorials_folder):
            relative_root = os.path.relpath(root, tutorials_folder)
            for file in files:
                if file.endswith('.md'):
                    file_path = os.path.join(root, file)
                    relative_path = os.path.relpath(file_path, tutorials_folder)
                    output_path = os.path.join('tutorials', relative_path).replace('\\', '/')
                    
                    # Build nested structure
                    current_nav = tutorials_structure
                    parts = relative_path.split(os.sep)
                    
                    # Handle nested directories
                    for part in parts[:-1]:  # All parts except the filename
                        if part:
                            current_nav = current_nav.setdefault(part, {})
                    
                    # Add the markdown file
                    file_name_without_ext = os.path.splitext(parts[-1])[0]
                    if isinstance(current_nav, dict):
                        current_nav[file_name_without_ext] = output_path
                    else:
                        current_nav = {file_name_without_ext: output_path}

    # Copy assets folder if exists
    if os.path.exists(ex_assets_folder):
        assets_output = os.path.join(output_folder, ex_assets_folder)
        if os.path.exists(assets_output):
            shutil.rmtree(assets_output)
        shutil.copytree(ex_assets_folder, assets_output)

    # Update homepage
    with open(ex_home_file, 'r', encoding='utf-8') as f:
        ex_home_content = f.read()

    with open(index_file, 'w', encoding='utf-8') as f:
        f.write(ex_home_content)

    # Build the nav section of mkdocs.yml
    nav = [{"Home": "index.md"}]
    
    # Add tutorials to navigation if they exist
    if tutorials_structure:
        nav.append({"Tutorials": build_nav_structure(tutorials_structure)})
    
    # Add API reference to navigation
    nav.append({"API Reference": build_nav_structure(nav_structure)})

    print("\nGenerated nav for mkdocs.yml:")
    print(yaml.dump({"nav": nav}, allow_unicode=True, sort_keys=False))

    # If mkdocs.yml exists, automatically update its nav section
    if os.path.exists(mkdocs_file):
        with open(mkdocs_file, 'r', encoding='utf-8') as f:
            mkdocs_config = yaml.unsafe_load(f)

        mkdocs_config['nav'] = nav

        with open(mkdocs_file, 'w', encoding='utf-8') as f:
            yaml.dump(mkdocs_config, f, allow_unicode=True, sort_keys=False)


def build_nav_structure(nav_dict: dict) -> list:
    """
    Build navigation structure from dictionary.
    
    Parameters:
        nav_dict (dict): Navigation dictionary
        
    Returns:
        list: Navigation structure for mkdocs
    """
    result = []
    for key, value in sorted(nav_dict.items()):
        if isinstance(value, dict):
            # It's a category with subitems
            result.append({key: build_nav_structure(value)})
        elif isinstance(value, list):
            # It's a list of class items (for API reference)
            class_items = []
            for item in value:
                if isinstance(item, dict):
                    for class_name, doc_path in item.items():
                        class_items.append({class_name: doc_path})
            result.append({key: class_items})
        else:
            # It's a direct file reference (for tutorials)
            result.append({key: value})
    return result


if __name__ == '__main__':
    # Example usage
    generate_api_docs(
        root_folder='src/python_motion_planning',  # Code directory
        output_folder='docs/',  # Directory for the generated documentation
        index_file='docs/index.md',  # Path to the homepage file
        mkdocs_file='mkdocs.yml',  # Path to the mkdocs.yml file
        ex_home_file='README.md',   # Extern homepage file
        ex_assets_folder='assets/'  # Extern assets folder
    )

# import os
# import ast
# import yaml
# import shutil


# def extract_classes(file_path: str):
#     """
#     Extract the names of all classes in a Python file.

#     Parameters:
#         file_path (str): Path to the Python file.

#     Returns:
#         class_names (list): List of class names.
#     """
#     class_names = []
#     with open(file_path, 'r', encoding='utf-8') as f:
#         node = ast.parse(f.read(), filename=file_path)
#     for child in node.body:
#         if isinstance(child, ast.ClassDef):
#             class_names.append(child.name)
#     return class_names


# def generate_api_docs(root_folder: str, output_folder: str, index_file: str, mkdocs_file: str, ex_home_file: str, ex_assets_folder: str):
#     """
#     Automatically generate Markdown files for API documentation, update the homepage, and modify mkdocs.yml.

#     Parameters:
#         root_folder (str): Path to the root directory of the library.
#         output_folder (str): Path to the output directory for the generated documentation.
#         index_file (str): Path to the homepage file.
#         mkdocs_file (str): Path to the mkdocs.yml file.
#     """
#     nav_structure = {}

#     if not os.path.exists(output_folder):
#         os.makedirs(output_folder)

#     for root, dirs, files in os.walk(root_folder):
#         relative_root = os.path.relpath(root, root_folder)  # Get the relative path of the current directory
#         for file in files:
#             if file.endswith('.py') and not file.startswith('__'):
#                 file_path = os.path.join(root, file)
#                 class_names = extract_classes(file_path)
#                 if not class_names:
#                     print(f"Warning: No class found in {file_path}. Skipping...")
#                     continue

#                 module_path = file_path.replace('.py', '').replace('/', '.').replace('\\', '.')
#                 # Create a corresponding output directory that mirrors the input directory's structure
#                 output_dir = os.path.join(output_folder, 'api', relative_root, file.replace('.py', ''))
#                 os.makedirs(output_dir, exist_ok=True)

#                 # Generate Markdown files for each class
#                 for class_name in class_names:
#                     class_md_path = os.path.join(output_dir, f"{class_name}.md")
#                     with open(class_md_path, 'w', encoding='utf-8') as f:
#                         f.write(f"# {class_name}\n\n::: {module_path}.{class_name}\n")

#                 # Update navigation structure
#                 current_nav = nav_structure
#                 for part in relative_root.split(os.sep):
#                     if part:
#                         current_nav = current_nav.setdefault(part, {})
#                 current_nav.setdefault(file.replace('.py', ''), []).extend([{
#                     class_name: os.path.relpath(os.path.join(output_dir, f"{class_name}.md"), output_folder)
#                 } for class_name in class_names])

#     if os.path.exists(ex_assets_folder):
#         shutil.copytree(ex_assets_folder, os.path.join(output_folder, ex_assets_folder), dirs_exist_ok=True)

#     with open(ex_home_file, 'r', encoding='utf-8') as f:
#         ex_home_content = f.read()

#     # Generate the content for the homepage
#     with open(index_file, 'w', encoding='utf-8') as f:
#         f.write(ex_home_content)
#         # f.write("\n\n# Documentation Contents\n\n")
#         # def write_nav(current_nav, level=2):
#         #     for category, subcategories in sorted(current_nav.items()):
#         #         f.write(f"{'#' * level} {category.capitalize()}\n\n")
#         #         if isinstance(subcategories, dict):
#         #             write_nav(subcategories, level + 1)
#         #         else:
#         #             for item in sorted(subcategories, key=lambda x: list(x.keys())[0]):
#         #                 class_name = list(item.keys())[0]
#         #                 doc_path = item[class_name].replace('\\', '/')
#         #                 f.write(f"- [{class_name}]({doc_path})\n")
#         #             f.write("\n")
#         # write_nav(nav_structure)

#     # Build the nav section of mkdocs.yml
#     nav = [{"Home": "index.md"}, {"API Reference": []}]
#     def build_nav(current_nav) -> list:
#         if isinstance(current_nav, dict):
#             return [{category: build_nav(subcategories)} for category, subcategories in sorted(current_nav.items())]
#         else:
#             return current_nav
#     nav[1]["API Reference"].extend(build_nav(nav_structure))

#     print("\nGenerated nav for mkdocs.yml:")
#     print(yaml.dump({"nav": nav}, allow_unicode=True, sort_keys=False))

#     # If mkdocs.yml exists, automatically update its nav section
#     if os.path.exists(mkdocs_file):
#         with open(mkdocs_file, 'r', encoding='utf-8') as f:
#             mkdocs_config = yaml.unsafe_load(f)

#         mkdocs_config['nav'] = nav

#         with open(mkdocs_file, 'w', encoding='utf-8') as f:
#             yaml.dump(mkdocs_config, f, allow_unicode=True, sort_keys=False)


# if __name__ == '__main__':
#     # Example usage
#     generate_api_docs(
#         root_folder='src/python_motion_planning',  # Code directory
#         output_folder='docs/',  # Directory for the generated documentation
#         index_file='docs/index.md',  # Path to the homepage file
#         mkdocs_file='mkdocs.yml',  # Path to the mkdocs.yml file
#         ex_home_file='README.md',   # Extern homepage file
#         ex_assets_folder='assets/'  # Extern assets folder
#     )
