bec_widgets/bec_widgets/utils/toolbars/bundles.py

from __future__ import annotations

from collections import defaultdict
from typing import TYPE_CHECKING, DefaultDict
from weakref import ReferenceType

import louie
from bec_lib.logger import bec_logger
from pydantic import BaseModel
from qtpy.QtCore import Qt
from qtpy.QtWidgets import QSizePolicy

from bec_widgets.utils.toolbars.actions import SeparatorAction, SplitterAction, ToolBarAction

DEFAULT_SIZE = 400
MAX_SIZE = 10_000_000

if TYPE_CHECKING:
    from qtpy.QtWidgets import QWidget

    from bec_widgets.utils.toolbars.connections import BundleConnection
    from bec_widgets.utils.toolbars.toolbar import ModularToolBar

logger = bec_logger.logger


class ActionInfo(BaseModel):
    action: ToolBarAction
    toolbar_bundle: ToolbarBundle | None = None

    model_config = {"arbitrary_types_allowed": True}


class ToolbarComponents:
    def __init__(self, toolbar: ModularToolBar):
        """
        Initializes the toolbar components.

        Args:
            toolbar (ModularToolBar): The toolbar to which the components will be added.
        """
        self.toolbar = toolbar

        self._components: dict[str, ActionInfo] = {}
        self.add("separator", SeparatorAction())

    def add(self, name: str, component: ToolBarAction):
        """
        Adds a component to the toolbar.

        Args:
            component (ToolBarAction): The component to add.
        """
        if name in self._components:
            raise ValueError(f"Component with name '{name}' already exists.")
        self._components[name] = ActionInfo(action=component, toolbar_bundle=None)

    def add_safe(self, name: str, component: ToolBarAction):
        """
        Adds a component to the toolbar, ensuring it does not already exist.

        Args:
            name (str): The name of the component.
            component (ToolBarAction): The component to add.
        """
        if self.exists(name):
            logger.info(f"Component with name '{name}' already exists. Skipping addition.")
            return
        self.add(name, component)

    def exists(self, name: str) -> bool:
        """
        Checks if a component exists in the toolbar.

        Args:
            name (str): The name of the component to check.

        Returns:
            bool: True if the component exists, False otherwise.
        """
        return name in self._components

    def get_action_reference(self, name: str) -> ReferenceType[ToolBarAction]:
        """
        Retrieves a component by name.

        Args:
            name (str): The name of the component to retrieve.

        """
        if not self.exists(name):
            raise KeyError(f"Component with name '{name}' does not exist.")
        return louie.saferef.safe_ref(self._components[name].action)

    def get_action(self, name: str) -> ToolBarAction:
        """
        Retrieves a component by name.

        Args:
            name (str): The name of the component to retrieve.

        Returns:
            ToolBarAction: The action associated with the given name.
        """
        if not self.exists(name):
            raise KeyError(
                f"Component with name '{name}' does not exist. The following components are available: {list(self._components.keys())}"
            )
        return self._components[name].action

    def set_bundle(self, name: str, bundle: ToolbarBundle):
        """
        Sets the bundle for a component.

        Args:
            name (str): The name of the component.
            bundle (ToolbarBundle): The bundle to set.
        """
        if not self.exists(name):
            raise KeyError(f"Component with name '{name}' does not exist.")
        comp = self._components[name]
        if comp.toolbar_bundle is not None:
            logger.info(
                f"Component '{name}' already has a bundle ({comp.toolbar_bundle.name}). Setting it to {bundle.name}."
            )
            comp.toolbar_bundle.bundle_actions.pop(name, None)
        comp.toolbar_bundle = bundle

    def remove_action(self, name: str):
        """
        Removes a component from the toolbar.

        Args:
            name (str): The name of the component to remove.
        """
        if not self.exists(name):
            raise KeyError(f"Action with ID '{name}' does not exist.")
        action_info = self._components.pop(name)
        if action_info.toolbar_bundle:
            action_info.toolbar_bundle.bundle_actions.pop(name, None)
            self.toolbar.refresh()
            action_info.toolbar_bundle = None
        if hasattr(action_info.action, "cleanup"):
            # Call cleanup if the action has a cleanup method
            action_info.action.cleanup()

    def cleanup(self):
        """
        Cleans up the toolbar components by removing all actions and bundles.
        """
        for action_info in self._components.values():
            if hasattr(action_info.action, "cleanup"):
                # Call cleanup if the action has a cleanup method
                action_info.action.cleanup()
        self._components.clear()


class ToolbarBundle:
    def __init__(self, name: str, components: ToolbarComponents):
        """
        Initializes a new bundle component.

        Args:
            bundle_name (str): Unique identifier for the bundle.
        """
        self.name = name
        self.components = components
        self.bundle_actions: DefaultDict[str, ReferenceType[ToolBarAction]] = defaultdict()
        self._connections: dict[str, BundleConnection] = {}

    def add_action(self, name: str):
        """
        Adds an action to the bundle.

        Args:
            name (str): Unique identifier for the action.
            action (ToolBarAction): The action to add.
        """
        if name in self.bundle_actions:
            raise ValueError(f"Action with name '{name}' already exists in bundle '{self.name}'.")
        if not self.components.exists(name):
            raise ValueError(
                f"Component with name '{name}' does not exist in the toolbar. Please add it first using the `ToolbarComponents.add` method."
            )
        self.bundle_actions[name] = self.components.get_action_reference(name)
        self.components.set_bundle(name, self)

    def remove_action(self, name: str):
        """
        Removes an action from the bundle.

        Args:
            name (str): The name of the action to remove.
        """
        if name not in self.bundle_actions:
            raise KeyError(f"Action with name '{name}' does not exist in bundle '{self.name}'.")
        del self.bundle_actions[name]

    def add_separator(self):
        """
        Adds a separator action to the bundle.
        """
        self.add_action("separator")

    def add_splitter(
        self,
        name: str = "splitter",
        target_widget: QWidget | None = None,
        initial_width: int = 10,
        min_width: int | None = None,
        max_width: int | None = None,
        size_policy_expanding: bool = True,
    ):
        """
        Adds a resizable splitter action to the bundle.

        Args:
            name (str): Unique identifier for the splitter action.
            target_widget (QWidget, optional): The widget whose size (width for horizontal,
                height for vertical orientation) will be controlled by the splitter. If None,
                the splitter will not control any widget.
            initial_width (int): The initial size of the splitter (width for horizontal,
                height for vertical orientation).
            min_width (int, optional): The minimum size the target widget can be resized to
                (width for horizontal, height for vertical orientation). If None, the target
                widget's minimum size hint in that orientation will be used.
            max_width (int, optional): The maximum size the target widget can be resized to
                (width for horizontal, height for vertical orientation). If None, the target
                widget's maximum size hint in that orientation will be used.
            size_policy_expanding (bool): If True, the size policy of the target_widget will be
                set to Expanding in the appropriate orientation if it is not already set.
        """

        # Resolve effective bounds
        eff_min = min_width if min_width is not None else None
        eff_max = max_width if max_width is not None else None

        is_horizontal = self.components.toolbar.orientation() == Qt.Orientation.Horizontal

        if target_widget is not None:
            # Use widget hints if bounds not provided
            if eff_min is None:
                eff_min = (
                    target_widget.minimumWidth() if is_horizontal else target_widget.minimumHeight()
                ) or 6
            if eff_max is None:
                mw = (
                    target_widget.maximumWidth() if is_horizontal else target_widget.maximumHeight()
                )
                eff_max = mw if mw and mw < MAX_SIZE else DEFAULT_SIZE  # avoid "no limit"

            # Adjust size policy if needed
            if size_policy_expanding:
                size_policy = target_widget.sizePolicy()

                if is_horizontal:
                    if size_policy.horizontalPolicy() not in (
                        QSizePolicy.Policy.Expanding,
                        QSizePolicy.Policy.MinimumExpanding,
                    ):
                        size_policy.setHorizontalPolicy(QSizePolicy.Policy.Expanding)
                        target_widget.setSizePolicy(size_policy)
                else:
                    if size_policy.verticalPolicy() not in (
                        QSizePolicy.Policy.Expanding,
                        QSizePolicy.Policy.MinimumExpanding,
                    ):
                        size_policy.setVerticalPolicy(QSizePolicy.Policy.Expanding)
                        target_widget.setSizePolicy(size_policy)

        splitter_action = SplitterAction(
            orientation="auto",
            parent=self.components.toolbar,
            initial_width=initial_width,
            min_width=eff_min,
            max_width=eff_max,
            target_widget=target_widget,
        )

        self.components.add_safe(name, splitter_action)
        self.add_action(name)

    def add_connection(self, name: str, connection):
        """
        Adds a connection to the bundle.

        Args:
            name (str): Unique identifier for the connection.
            connection: The connection to add.
        """
        if name in self._connections:
            raise ValueError(
                f"Connection with name '{name}' already exists in bundle '{self.name}'."
            )
        self._connections[name] = connection

    def remove_connection(self, name: str):
        """
        Removes a connection from the bundle.

        Args:
            name (str): The name of the connection to remove.
        """
        if name not in self._connections:
            raise KeyError(f"Connection with name '{name}' does not exist in bundle '{self.name}'.")
        self._connections[name].disconnect()
        del self._connections[name]

    def get_connection(self, name: str):
        """
        Retrieves a connection by name.

        Args:
            name (str): The name of the connection to retrieve.

        Returns:
            The connection associated with the given name.
        """
        if name not in self._connections:
            raise KeyError(f"Connection with name '{name}' does not exist in bundle '{self.name}'.")
        return self._connections[name]

    def disconnect(self):
        """
        Disconnects all connections in the bundle.
        """
        for connection in self._connections.values():
            connection.disconnect()
        self._connections.clear()