import os
import sys
import importlib
import inspect
from pathlib import Path
import shutil

def generate_markdown_handbook():
    print("📝 Starting mxPIC Markdown Generation...")

    # Define paths
    src_root = Path("mxpic/components")
    # This should point to where your Sphinx .md files are stored
    docs_root = Path("docs/source/mxpic/components") 
    
    # 1. Delete the directory and all its contents
    if docs_root.exists() and docs_root.is_dir():
        shutil.rmtree(docs_root)

    # 2. Recreate the directory
    docs_root.mkdir(parents=True, exist_ok=True)

    success_count = 0

    for py_file in src_root.rglob("*.py"):
        if py_file.name == "__init__.py":
            continue

        # Convert path to module name (e.g., mxpic.components.primitives.beam_splitters)
        rel_path = py_file.relative_to(src_root)
        module_name = "mxpic.components." + str(rel_path.with_suffix("")).replace(os.sep, ".")

        try:
            module = importlib.import_module(module_name)
        except Exception as e:
            print(f"⚠️ Could not import {module_name}: {e}")
            continue

        # Find all classes defined inside this specific module
        classes = []
        for name, obj in inspect.getmembers(module, inspect.isclass):
            if obj.__module__ == module_name:
                classes.append(name)

        # If the file has no classes, skip it
        if not classes:
            continue

        # Define where to save the .md file
        md_file_path = docs_root / rel_path.with_suffix(".md")
        md_file_path.parent.mkdir(parents=True, exist_ok=True)

        # --- WRITE THE MARKDOWN FILE ---
        with open(md_file_path, "w", encoding="utf-8") as f:
            # 1. Write the Module Header
            f.write(f"# {module_name}\n\n")
            
            # 2. Document any module-level docstrings (skipping classes)
            f.write("```{eval-rst}\n")
            f.write(f".. automodule:: {module_name}\n")
            f.write("   :no-members:\n") # This prevents duplicating the classes!
            f.write("```\n\n")

            # 3. Loop through and write each class with its image
            for class_name in classes:
                f.write(f"## {class_name}\n\n")
                
                # Point to the image path in Sphinx
                # img_path = f"{sphinx_image_root}/{rel_path.parent.as_posix()}/{class_name}.png"
                img_path = f"{class_name}.png"
                
                f.write("```{eval-rst}\n")
                # Insert the Sphinx image directive
                f.write(f".. image:: {img_path}\n")
                f.write("   :align: center\n")
                f.write("   :width: 600px\n\n")
                
                # Insert the specific class documentation
                f.write(f".. autoclass:: {module_name}.{class_name}\n")
                f.write("   :members:\n")
                f.write("   :undoc-members:\n")
                f.write("   :show-inheritance:\n")
                f.write("```\n\n")

        print(f"✅ Generated docs for: {module_name}")
        success_count += 1

    print(f"\n✨ Markdown generation complete! Updated {success_count} files.")

if __name__ == "__main__":
    generate_markdown_handbook()