from PyQt6.QtCore import Qt, QSize
from PyQt6.QtGui import QAction, QIcon
from PyQt6.QtWidgets import QToolBar, QWidget, QSizePolicy


class ToolBar(QToolBar):
    def __init__(self, parent,
                 orientation: Qt.Orientation = Qt.Orientation.Horizontal,
                 style: Qt.ToolButtonStyle = Qt.ToolButtonStyle.ToolButtonTextUnderIcon,
                 icon_size: tuple[int, int] = (32, 32)) -> None:
        """
        Initialize the toolbar.

        Parameters
        ----------
        parent : QWidget
            The parent widget.
        orientation : Qt.Orientation, optional
            The toolbar's orientation, by default Qt.Orientation.Horizontal.
        style : Qt.ToolButtonStyle, optional
            The toolbar's tool button style, by default Qt.ToolButtonStyle.ToolButtonTextUnderIcon.
        icon_size : tuple[int, int], optional
            The toolbar's icon size, by default (32, 32).
        """
        super().__init__(parent)
        self.actions_call = {}
        self.setOrientation(orientation)

        self.setToolButtonStyle(style)
        self.setIconSize(QSize(icon_size[0], icon_size[1]))

    def add_button(self, key: str, text: str, icon: str, trigger_action) -> None:
        """
        Add a button to the toolbar.

        Parameters
        ----------
        key : str
            Buttons key.
        text : str
            The button's text.
        icon : str
            The button's icon.
        trigger_action : callable
            The action to be executed when the button is clicked.
        """
        self.actions_call[key] = QAction(QIcon(icon), text, self)
        self.actions_call[key].triggered.connect(trigger_action)
        self.addAction(self.actions_call[key])

    def add_separator(self) -> None:
        """
        Add a separator to the toolbar.
        """
        separator = QWidget(self)
        separator.setSizePolicy(
            QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        self.addWidget(separator)

    def enable_button(self, key):
        """
        Enable a button in the toolbar.

        Parameters
        ----------
        key : str
            The button's key.
        """
        self.actions_call[key].setEnabled(True)

    def disable_button(self, key):
        """
        Disable a button in the toolbar.

        Parameters
        ----------
        key : str
            The button's key.
        """
        self.actions_call[key].setEnabled(False)