menu

class MenuDrawer:
    """Drawer for QMenu using native QPainter calls (no QSS).

    Paints the menu panel/border (PE_PanelMenu/PE_FrameMenu) and each
    item row (CE_MenuItem).
    """

    _WIDGET_CLS = "QMenu"
    # Horizontal and vertical padding of the box around shortcut text
    _SHORTCUT_PADDING = (5, 2)

    def __init__(self, style_inst: AYONStyle) -> None:
        self.style_inst = style_inst
        self.model = style_inst.model

    @property
    def base_class(self) -> dict:
        return {"QMenu": QMenu}

    # ------------------------------------------------------------------
    # Registration
    # ------------------------------------------------------------------

    def register_drawers(self) -> dict:
        return {
            enum_to_str(
                QStyle.PrimitiveElement,
                QStyle.PrimitiveElement.PE_PanelMenu,
                "QMenu",
            ): self.draw_panel,
            enum_to_str(
                QStyle.PrimitiveElement,
                QStyle.PrimitiveElement.PE_FrameMenu,
                "QMenu",
            ): self.draw_frame,
            enum_to_str(
                QStyle.ControlElement,
                QStyle.ControlElement.CE_MenuItem,
                "QMenu",
            ): self.draw_menu_item,
            enum_to_str(
                QStyle.ControlElement,
                QStyle.ControlElement.CE_MenuEmptyArea,
                "QMenu",
            ): do_nothing,
        }

    def register_sizers(self) -> dict:
        return {
            enum_to_str(
                QStyle.ContentsType,
                QStyle.ContentsType.CT_MenuItem,
                "QMenu",
            ): self.menu_item_size,
        }

    def register_metrics(self) -> dict:
        pm = QStyle.PixelMetric
        metrics_map = {
            pm.PM_MenuPanelWidth: self.get_metric,
            pm.PM_MenuHMargin: self.get_metric,
            pm.PM_MenuVMargin: self.get_metric,
            pm.PM_SmallIconSize: self.get_metric,
            pm.PM_MenuButtonIndicator: self.get_metric,
        }
        return {enum_to_str(pm, k, "QMenu"): v for k, v in metrics_map.items()}

    # ------------------------------------------------------------------
    # Metrics
    # ------------------------------------------------------------------

    def _base_style(self, widget: QWidget | None = None):
        """Return the base style dict, context-bound to *widget*.

        This is used for general operations like metrics calculation and
        panel/frame painting, where the state/variant-specific styles are
        not relevant.
        """
        variant = (
            widget.property("variant")
            if widget is not None
            else None
        ) or "default"
        style = self.model.get_style(self._WIDGET_CLS, variant, "base")
        style.set_context(widget)
        return style

    def get_metric(
        self,
        metric: QStyle.PixelMetric,
        opt: QStyleOption | None = None,
        widget: QWidget | None = None,
    ) -> int:
        pm = QStyle.PixelMetric
        style = self._base_style(widget)
        if metric == pm.PM_MenuPanelWidth:
            return int(style.get("border-width", 1))
        if metric in (pm.PM_MenuHMargin, pm.PM_MenuVMargin):
            pp = style.get("panel-padding", [4, 4])
            if isinstance(pp, (list, tuple)):
                return int(pp[0] if metric == pm.PM_MenuHMargin else pp[1])
            return int(pp)
        if metric == pm.PM_SmallIconSize:
            return int(style.get("icon-size", 16))
        if metric == pm.PM_MenuButtonIndicator:
            return int(style.get("icon-size", 16))
        return 0

    # ------------------------------------------------------------------
    # Sizing and Layout Helper
    # ------------------------------------------------------------------

    def _compute_layout(
        self,
        option: QStyleOptionMenuItem,
        style: dict[str, Any],
        contents_size: QSize | None = None,
    ) -> MenuItemLayout:
        """Compute the metrics, dimensions and layouts for a menu item.

        Args:
            option: The style option for the menu item.
            style: The resolved style dict.
            contents_size: Optional contents size passed from size query.

        Returns:
            The computed layout information for the menu item.
        """
        ip = style.get("item-padding", [6, 6])
        if isinstance(ip, (list, tuple)):
            pad_h, pad_v = int(ip[0]), int(ip[1])
        else:
            pad_h = pad_v = int(ip)

        icon_size = int(style.get("icon-size", 16))
        item_spacing = int(style.get("item-spacing", 4))

        fm = option.fontMetrics
        text_h = QFontMetrics(self.model.base_font).height()
        if text_h <= 0:
            text_h = (
                contents_size.height()
                if contents_size
                else option.rect.height()
            )

        min_h = int(style.get("min-item-height", 0))
        row_h = max(text_h + pad_v * 2, icon_size + pad_v * 2, min_h)

        text = option.text or ""
        label, _, shortcut = text.partition("\t")
        if doc := _rich_text_document(label, option.font):
            text_w = math.ceil(doc.idealWidth())
        elif fm:
            text_w = fm.horizontalAdvance(label)
        else:
            text_w = contents_size.width() if contents_size else 0

        sc_w = getattr(
            option,
            "reservedShortcutWidth",
            getattr(option, "tabWidth", 0),
        )
        if shortcut and sc_w == 0 and fm:
            sc_w = fm.horizontalAdvance(shortcut)
        if sc_w:
            # Shortcut is drawn with smaller font inside a padded box
            scale = float(style.get("shortcut-font-scale", 0.85))
            sc_w = int(sc_w * scale) + self._SHORTCUT_PADDING[0] * 2

        icon_gutter = icon_size + item_spacing
        is_submenu = (
            option.menuItemType == QStyleOptionMenuItem.MenuItemType.SubMenu
        )
        arrow_w = icon_size if is_submenu else 0
        # The arrow glyph only covers the middle of its icon, the blank
        # part of the icon around it already acts as padding. Use less
        # padding next to it so the arrow is not left far from the edge.
        arrow_pad = pad_h // 2

        total_w = (
            icon_gutter
            + pad_h
            + text_w
            + (pad_h + sc_w if sc_w else 0)
            + (pad_h + arrow_w + arrow_pad if arrow_w else pad_h)
        )

        return MenuItemLayout(
            pad_h=pad_h,
            pad_v=pad_v,
            icon_size=icon_size,
            item_spacing=item_spacing,
            icon_gutter=icon_gutter,
            label_text=label,
            shortcut_text=shortcut,
            text_w=text_w,
            text_h=text_h,
            sc_w=sc_w,
            arrow_w=arrow_w,
            arrow_pad=arrow_pad,
            row_h=row_h,
            total_w=total_w,
        )

    def menu_item_size(
        self,
        contents_type: QStyle.ContentsType,
        option: QStyleOption | None,
        contents_size: QSize,
        widget: QWidget | None = None,
    ) -> QSize:
        """Compute the bounding size of a single menu item row."""
        if not isinstance(option, QStyleOptionMenuItem):
            return QSize(contents_size.width(), contents_size.height())

        style = self._base_style(widget)

        # Separator: early exit
        sep_type = QStyleOptionMenuItem.MenuItemType.Separator
        if option.menuItemType == sep_type:
            sep_h = int(style.get("separator-height", 1))
            return QSize(contents_size.width(), sep_h)

        layout = self._compute_layout(option, style, contents_size)
        return QSize(
            layout.total_w,
            layout.row_h,
        )

    # ------------------------------------------------------------------
    # Primitive painting: panel & frame
    # ------------------------------------------------------------------

    def draw_panel(
        self,
        option: QStyleOption,
        painter: QPainter,
        widget: QWidget | None = None,
    ) -> None:
        """Fill the menu background with the panel colour."""
        painter.save()
        painter.setRenderHint(QPainter.RenderHint.Antialiasing, True)
        style = self._base_style(widget)
        radius = int(style.get("border-radius", 6))
        painter.setBrush(QBrush(QColor(style["background-color"])))
        painter.setPen(Qt.PenStyle.NoPen)
        painter.drawRoundedRect(option.rect, radius, radius)
        painter.restore()

    def draw_frame(
        self,
        option: QStyleOption,
        painter: QPainter,
        widget: QWidget | None = None,
    ) -> None:
        """Stroke the menu border."""
        painter.save()
        painter.setRenderHint(QPainter.RenderHint.Antialiasing, True)
        style = self._base_style(widget)
        radius = int(style.get("border-radius", 6))
        bw = int(style.get("border-width", 1))
        pen = QPen(QColor(style["border-color"]))
        pen.setWidth(bw)
        painter.setPen(pen)
        painter.setBrush(Qt.BrushStyle.NoBrush)
        # Inset by half the pen width so the stroke is fully inside the rect.
        inset = bw / 2.0
        inset_rect = QRectF(option.rect).adjusted(inset, inset, -inset, -inset)
        painter.drawRoundedRect(inset_rect, radius, radius)
        painter.restore()

    # ------------------------------------------------------------------
    # Control painting: individual item rows
    # ------------------------------------------------------------------

    def draw_menu_item(
        self,
        option: QStyleOption,
        painter: QPainter,
        widget: QWidget | None = None,
    ) -> None:
        """Paint a single QMenu row (normal, separator, submenu)."""
        if not isinstance(option, QStyleOptionMenuItem):
            return

        painter.save()
        painter.setRenderHint(QPainter.RenderHint.Antialiasing, True)

        # Handle separators and return
        sep_type = QStyleOptionMenuItem.MenuItemType.Separator
        if option.menuItemType == sep_type:
            self._draw_separator(option, painter, widget)
            painter.restore()
            return

        # Resolve state and fetch styles
        is_enabled = bool(option.state & QStyle.StateFlag.State_Enabled)
        is_selected = bool(option.state & QStyle.StateFlag.State_Selected)
        state = (
            "disabled"
            if not is_enabled
            else "hover"
            if is_selected and is_enabled
            else "base"
        )

        # Resolve variant from action property, if available
        # NOTE: This is a bit of a hack. I could have created a QAction section
        # in ayon_style.json and fetched the variant from there.
        action = None
        if isinstance(widget, QMenu):
            action = widget.actionAt(option.rect.center())
        variant = (
            (action.property("variant") if action else None)
            or (widget.property("variant") if widget else None)
            or "default"
        )

        style = self.model.get_style(
            self._WIDGET_CLS, variant=variant, state=state
        )
        style.set_context(widget)

        layout = self._compute_layout(option, style)
        item_radius = int(style.get("item-radius", 4))
        rect = option.rect

        # --- Selection background ---
        bg = QColor(style.get("background-color", "#424a57"))
        painter.setBrush(QBrush(bg))
        painter.setPen(Qt.PenStyle.NoPen)
        painter.drawRoundedRect(rect, item_radius, item_radius)

        opacity = float(style.get("opacity", 1.0))
        painter.setOpacity(opacity)

        # --- Left gutter (icon or check mark) ---
        x = rect.left() + layout.pad_h
        cy = rect.center().y()
        # Icons are drawn slightly towards the edge, away from the label
        icon_x = x - 2

        check_type = option.checkType
        not_checkable = QStyleOptionMenuItem.CheckType.NotCheckable
        if check_type != not_checkable:
            check_color = QColor(style.get("color", "#f4f5f5"))
            check_style = (
                action.property("check-style") if action else None
            )
            if option.checked or check_style != "checkmark":
                if check_style == "checkmark":
                    icon_name = "check"
                elif option.checked:
                    icon_name = "check_box"
                else:
                    icon_name = "check_box_outline_blank"
                check_icon = get_icon(
                    icon_name,
                    color=check_color,
                    fill=False,
                )
                check_rect = QRect(
                    icon_x,
                    cy - layout.icon_size // 2,
                    layout.icon_size,
                    layout.icon_size,
                )
                check_icon.paint(painter, check_rect)
        elif not option.icon.isNull():
            icon_rect = QRect(
                icon_x,
                cy - layout.icon_size // 2,
                layout.icon_size,
                layout.icon_size,
            )
            mode = QIcon.Mode.Disabled if not is_enabled else QIcon.Mode.Normal
            option.icon.paint(
                painter, icon_rect, Qt.AlignmentFlag.AlignCenter, mode
            )

        x += layout.icon_gutter

        # --- Text (label + shortcut) ---
        text_color = QColor(style.get("color", "#f4f5f5"))
        painter.setPen(QPen(text_color))

        right_margin = (
            layout.pad_h + layout.arrow_w + layout.arrow_pad
            if layout.arrow_w
            else layout.pad_h
        )
        # 'rect.right()' is the last pixel inside the rect, one short of
        # its right edge - without the '+ 1' the label is a pixel narrower
        # than the width measured for it and loses its last character.
        label_rect = QRect(
            x,
            rect.top(),
            rect.right()
            + 1
            - x
            - right_margin
            - (layout.sc_w + layout.pad_h if layout.sc_w else 0),
            rect.height(),
        )
        if doc := _rich_text_document(layout.label_text, painter.font()):
            self._draw_rich_text(painter, doc, label_rect, text_color)
        else:
            # Not clipped: with fractional font metrics (display scaling)
            # the text can still be a pixel wider than measured, the
            # padding next to it has room for that.
            painter.drawText(
                label_rect,
                Qt.AlignmentFlag.AlignVCenter
                | Qt.AlignmentFlag.AlignLeft
                | Qt.TextFlag.TextDontClip,
                layout.label_text,
            )

        if layout.shortcut_text:
            sc_bg_color = QColor(
                style.get(
                    "shortcut-background-color",
                    style.get("background-color", "#2b3036"),
                )
            )
            # Shortcut is a hint, draw it smaller than the label inside
            # a rounded box aligned to the right side of the row.
            sc_font = QFont(painter.font())
            scale = float(style.get("shortcut-font-scale", 0.85))
            if sc_font.pixelSize() > 0:
                sc_font.setPixelSize(
                    max(1, round(sc_font.pixelSize() * scale))
                )
            else:
                sc_font.setPointSizeF(sc_font.pointSizeF() * scale)
            painter.setFont(sc_font)

            sc_fm = QFontMetrics(sc_font)
            sc_pad_h, sc_pad_v = self._SHORTCUT_PADDING
            sc_rect = QRect(
                0,
                0,
                sc_fm.horizontalAdvance(layout.shortcut_text)
                + sc_pad_h * 2,
                sc_fm.height() + sc_pad_v * 2,
            )
            sc_rect.moveCenter(rect.center())
            sc_rect.moveRight(rect.right() - right_margin)
            painter.setBrush(QBrush(sc_bg_color))
            painter.setPen(Qt.PenStyle.NoPen)
            painter.drawRoundedRect(sc_rect, 4, 4)

            sc_color = QColor(
                style.get(
                    "shortcut-color",
                    style.get("color", "#8b9198"),
                )
            )
            sc_color.setAlphaF(style.get("shortcut-opacity", 0.6))
            painter.setPen(QPen(sc_color))
            painter.drawText(
                sc_rect,
                Qt.AlignmentFlag.AlignVCenter | Qt.AlignmentFlag.AlignHCenter,
                layout.shortcut_text,
            )

        # --- Submenu arrow ---
        if layout.arrow_w > 0:
            arrow_color = QColor(style.get("color", "#f4f5f5"))
            arrow_icon = get_icon("chevron_right", color=arrow_color)
            arrow_rect = QRect(
                rect.right() - layout.arrow_pad - layout.arrow_w,
                cy - layout.icon_size // 2,
                layout.icon_size,
                layout.icon_size,
            )
            arrow_icon.paint(painter, arrow_rect)

        painter.restore()

    @staticmethod
    def _draw_rich_text(
        painter: QPainter,
        doc: QTextDocument,
        rect: QRect,
        color: QColor,
    ) -> None:
        """Draw a rich text label vertically centered in *rect*."""
        height = doc.size().height()

        painter.save()
        painter.translate(
            rect.left(),
            rect.top() + (rect.height() - height) / 2.0,
        )

        # paint
        context = QAbstractTextDocumentLayout.PaintContext()
        context.palette.setColor(QPalette.ColorRole.Text, color)
        doc.documentLayout().draw(painter, context)

        painter.restore()

    def _draw_separator(
        self,
        option: QStyleOptionMenuItem,
        painter: QPainter,
        widget: QWidget | None = None,
    ) -> None:
        """Draw a thin horizontal separator line."""
        style = self._base_style(widget)
        color = QColor(
            style.get("separator-color", style.get("border-color", "#41474d"))
        )
        sep_h = int(style.get("separator-height", 1))

        cy = option.rect.center().y()
        y = cy - sep_h // 2
        painter.setPen(Qt.PenStyle.NoPen)
        painter.setBrush(QBrush(color))
        painter.drawRect(
            option.rect.left(),
            y,
            option.rect.width(),
            sep_h,
        )

@dataclass
class MenuItemLayout:
    """Dataclass to hold the results of menu item layout calculations."""

    pad_h: int
    pad_v: int
    icon_size: int
    item_spacing: int
    icon_gutter: int
    label_text: str
    shortcut_text: str
    text_w: int
    text_h: int
    sc_w: int
    arrow_w: int
    arrow_pad: int
    row_h: int
    total_w: int