import importlib
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import nazca as nd

from generate_handbook import DEFAULT_PACKAGE_ROOT, DEFAULT_SRC_ROOT, public_top_level_members


PALETTE = {
    "WG": "blue",
    "SLAB": "cyan",
    "M1": "#FFD700",
    "M2": "silver",
    "DEEP_TRENCH": "black",
    (1, 0): "darkred",
    (2, 0): "green",
    (1111, 0): "green",
    (63, 30): "#FFD700",
}


def apply_mxpic_colors() -> None:
    """Apply the mxPIC color palette to Nazca when the layers are available."""
    print("Applying mxPIC layer colors...")
    for layer_id, color in PALETTE.items():
        try:
            nd.set_layercolor(layer=layer_id, color=color)
        except Exception:
            continue


def remove_generated_pngs(img_root: Path) -> None:
    """Remove generated images while preserving Markdown and other docs files."""
    if not img_root.exists():
        return

    for png_file in img_root.rglob("*.png"):
        png_file.unlink()


def module_name_for(py_file: Path, src_root: Path, package_root: str) -> str:
    relative_path = py_file.relative_to(src_root).with_suffix("")
    return f"{package_root}.{'.'.join(relative_path.parts)}"


def generate_image_for_class(target_dir: Path, class_name: str, component_class: type) -> str:
    """Instantiate one component class and save its cell image."""
    nd.clear_layout()
    instance = component_class()
    cell = getattr(instance, "cell", None)
    if cell is None:
        return "skipped"

    target_dir.mkdir(parents=True, exist_ok=True)
    image_path = target_dir / f"{class_name}.png"
    nd.export_plt(path="", title=class_name, topcells=[cell])
    plt.savefig(image_path, bbox_inches="tight", dpi=300)
    plt.close()
    return "generated"


def generate_component_images(
    img_root: Path = Path("images/components"),
    src_root: Path = DEFAULT_SRC_ROOT,
    package_root: str = DEFAULT_PACKAGE_ROOT,
) -> dict[str, int]:
    print("Starting mxPIC component image generation...")

    img_root = Path(img_root)
    src_root = Path(src_root)
    counts = {"generated": 0, "skipped": 0, "failed": 0}

    if not src_root.exists():
        print(f"Source directory not found: {src_root}")
        counts["failed"] += 1
        return counts

    img_root.mkdir(parents=True, exist_ok=True)
    remove_generated_pngs(img_root)
    apply_mxpic_colors()

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

        module_name = module_name_for(py_file, src_root, package_root)
        try:
            module = importlib.import_module(module_name)
        except Exception as error:
            print(f"Skipped module import {module_name}: {error}")
            counts["skipped"] += 1
            continue

        class_names, _ = public_top_level_members(py_file)
        for class_name in class_names:
            component_class = getattr(module, class_name, None)
            if component_class is None:
                print(f"Skipped missing class {module_name}.{class_name}")
                counts["skipped"] += 1
                continue

            try:
                target_dir = img_root / py_file.relative_to(src_root).parent
                result = generate_image_for_class(target_dir, class_name, component_class)
            except Exception as error:
                plt.close()
                print(f"Failed image for {module_name}.{class_name}: {error}")
                counts["failed"] += 1
                continue

            counts[result] += 1
            if result == "generated":
                print(f"Generated: {target_dir / f'{class_name}.png'}")
            else:
                print(f"Skipped class without cell: {module_name}.{class_name}")

    print(
        "Image generation complete. "
        f"Generated: {counts['generated']} | "
        f"Skipped: {counts['skipped']} | "
        f"Failures: {counts['failed']}"
    )
    return counts


if __name__ == "__main__":
    generate_component_images(img_root=Path("docs/source/mxpic/components"))