"""Group generated node classes into module files and write them out."""

from __future__ import annotations

import ast
from pathlib import Path

from .config import Disposition, TreeTypeConfig
from .emit import generate_file_header, generate_node_class
from .introspect import introspect_node
from .model import NodeInfo


def _module_all(path: Path) -> list[str]:
    """The ``__all__`` names declared in a generated module, read without
    importing it (used to re-export the separately-generated ``assets.py``)."""
    if not path.exists():
        return []
    tree = ast.parse(path.read_text(encoding="utf-8"))
    for node in tree.body:
        if isinstance(node, ast.Assign) and any(
            isinstance(t, ast.Name) and t.id == "__all__" for t in node.targets
        ):
            value = ast.literal_eval(node.value)
            return list(value)
    return []


class ModuleWriter:
    def __init__(self, module_name: str, output_dir: Path):
        self.module_name = module_name
        self.output_dir = output_dir
        self.node_list: list[NodeInfo] = list()

    @property
    def filename(self):
        return f"{self.module_name}.py"

    def write(self, config: TreeTypeConfig):
        module_path = self.output_dir / self.filename
        with open(module_path, "w") as file:
            file.write(self.generate_content(config))

    def generate_content(self, config: TreeTypeConfig) -> str:
        print(f"Generating {self.filename} with {len(self.nodes)} nodes...")
        string = generate_file_header(self.nodes, config)
        string += "\n\n"

        for node_info in self.nodes:
            class_code = generate_node_class(node_info, config)
            string += class_code
            string += "\n\n"

        return string

    @property
    def nodes(self):
        return sorted(self.node_list, key=lambda node: node.name)

    def add_node(self, node_info: NodeInfo):
        self.node_list.append(node_info)


def _enum_properties_match(a: NodeInfo, b: NodeInfo) -> bool:
    """Check whether two NodeInfo objects have identical enum properties.

    Returns True if every ENUM property in both nodes has the same set of
    allowed values, meaning one can safely re-export the other's class.
    """

    def enum_signature(node_info: NodeInfo) -> dict[str, list[str]]:
        return {
            p.identifier: [e.identifier for e in p.enum_items]
            for p in node_info.properties
            if p.prop_type == "ENUM"
        }

    return enum_signature(a) == enum_signature(b)


class ModulesHandler:
    def __init__(
        self,
        config: TreeTypeConfig,
        gn_registry: dict[str, tuple[str, str, NodeInfo]] | None = None,
    ):
        self.config = config
        self.output_dir = config.output_dir
        self.modules: dict[str, ModuleWriter] = {}
        # Track class names to detect duplicates across modules
        self._class_names: dict[str, NodeInfo] = {}
        # Nodes re-exported from the geometry module: maps class_name → gn module name
        self._gn_reexports: dict[str, str] = {}
        # Registry from the geometry handler: maps bl_idname → (class_name, module_name)
        # Used to auto-detect nodes that were already generated for geometry.
        self._gn_registry = gn_registry or {}

    def write(self):
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self._delete_stale_modules()
        self.write_modules()
        self.write_init()

    def _delete_stale_modules(self):
        """Delete auto-generated .py files that are no longer produced by this run."""
        # Prefix-match so files stamped by the old generate.py header still count.
        auto_header = "# Auto-generated by"
        # assets.py is produced by the separate `nodebpy.assets` step (which runs
        # before this one) and re-exported by generate_init; it isn't one of our
        # module writers, so guard it explicitly or we'd delete it as "stale".
        active_files = {writer.filename for writer in self.modules.values()} | {
            "__init__.py",
            "assets.py",
        }
        for path in self.output_dir.glob("*.py"):
            if path.name in active_files:
                continue
            try:
                first_line = path.read_text(encoding="utf-8").split("\n", 1)[0]
            except OSError:
                continue
            if first_line.startswith(auto_header):
                print(f"  Deleting stale module: {path.name}")
                path.unlink()

    def write_modules(self):
        for module_name, module_writer in self.modules.items():
            module_writer.write(self.config)

    def generate_init(self):
        string = "# Auto-generated by `python -m gen` — do not edit manually.\n\n"
        all_names = []
        if self.config.manually_defined:
            imports = ",\n    ".join(self.config.manually_defined)
            string += f"from .manual import (\n    {imports},\n)\n"
            all_names += list(self.config.manually_defined)

        # Emit re-exports from geometry, grouped by source module
        if self._gn_reexports:
            gn_by_module: dict[str, list[str]] = {}
            for cls_name, gn_module in self._gn_reexports.items():
                gn_by_module.setdefault(gn_module, []).append(cls_name)
            for gn_module, class_names in sorted(gn_by_module.items()):
                class_names.sort()
                all_names += class_names
                imports = ",\n    ".join(class_names)
                string += f"from ..geometry.{gn_module} import (\n    {imports},\n)\n"

        for writer in self.modules.values():
            nodes = writer.nodes
            class_names = [node.class_name_for_config(self.config) for node in nodes]
            all_names += class_names
            imports = ",\n    ".join(class_names)
            string += f"from .{writer.module_name} import (\n    {imports},\n)\n"

        # Asset-backed node groups live in a sibling assets.py generated by
        # `python -m nodebpy.assets`; re-export them so they sit alongside the
        # built-in nodes (e.g. g.SmoothByAngle()).
        asset_names = _module_all(self.output_dir / "assets.py")
        if asset_names:
            all_names += asset_names
            imports = ",\n    ".join(sorted(asset_names))
            string += f"from .assets import (\n    {imports},\n)\n"

        all_names.sort()
        all_items = ",\n    ".join(f'"{name}"' for name in all_names)
        string += f"\n__all__ = (\n    {all_items},\n)\n"
        return string

    def write_init(self):
        filepath = self.output_dir / "__init__.py"
        with open(filepath, "w") as file:
            file.write(self.generate_init())

    def add_node(self, node_info: NodeInfo | None) -> None:
        if not node_info:
            return
        cls_name = node_info.class_name_for_config(self.config)
        if cls_name in self._class_names:
            existing = self._class_names[cls_name]
            print(
                f"  Duplicate class name '{cls_name}': "
                f"{node_info.bl_idname} vs {existing.bl_idname} — skipping {node_info.bl_idname}"
            )
            return
        self._class_names[cls_name] = node_info

        # If this exact node (by bl_idname) was already generated for geometry,
        # re-export it — unless its enum properties differ in this tree type.
        if self._gn_registry and node_info.bl_idname in self._gn_registry:
            gn_cls_name, gn_module, gn_node_info = self._gn_registry[
                node_info.bl_idname
            ]
            if _enum_properties_match(gn_node_info, node_info):
                self._gn_reexports[gn_cls_name] = gn_module
                print(f"  Re-exporting from geometry.{gn_module}: {gn_cls_name}")
                return
            else:
                print(
                    f"  Enum mismatch for {node_info.bl_idname} — "
                    f"generating separate class for {self.config.output_dir_name}"
                )

        module_name = node_info.module_name
        if module_name not in self.modules:
            self.modules[module_name] = ModuleWriter(module_name, self.output_dir)
        self.modules[module_name].add_node(node_info)

    def build_registry(
        self, all_nodes: list[type]
    ) -> dict[str, tuple[str, str, NodeInfo]]:
        """Build a bl_idname → (class_name, module_name, node_info) mapping.

        Used to pass the geometry handler's registry to subsequent handlers
        so they can auto-detect shared nodes by bl_idname and compare
        properties to ensure compatibility.
        """
        registry: dict[str, tuple[str, str, NodeInfo]] = {}
        for cls_name, node_info in self._class_names.items():
            if cls_name in self.config.manually_defined:
                module = "manual"
            else:
                module = node_info.module_name
            registry[node_info.bl_idname] = (cls_name, module, node_info)

        # Also include manually-defined nodes — they were skipped during
        # add_node so they're not in _class_names, but subsequent tree-type
        # handlers still need to detect them for re-export. introspect_node is
        # cached, so re-introspecting here is free.
        for node_type in all_nodes:
            if self.config.disposition(node_type) is not Disposition.MANUAL:
                continue
            if node_type.__name__ in registry:
                continue
            node_info = introspect_node(node_type, self.config.tree_type)
            if node_info:
                cls_name = node_info.class_name_for_config(self.config)
                registry[node_info.bl_idname] = (cls_name, "manual", node_info)

        return registry

    def count_nodes(self) -> int:
        return sum(len(module.nodes) for module in self.modules.values())
