from enum import Enum

from PyQt6.QtCore import Qt, QRect, QSize, QPoint, QMargins
from PyQt6.QtWidgets import QLayout


class JustifyMode(Enum):
    NONE = 0
    SPACE_BETWEEN = 1
    SPACE_AROUND = 2
    SPACE_EVENLY = 3
    CENTER = 4


class LastRowMode(Enum):
    LEFT = 0
    JUSTIFY = 1
    MATCH_PREVIOUS = 2
    CENTER = 3


class FlowLayout(QLayout):
    def __init__(
        self,
        parent=None,
        justify_mode=JustifyMode.SPACE_BETWEEN,
        last_row_mode=LastRowMode.LEFT,
    ):
        super().__init__(parent)

        if parent is not None:
            self.setContentsMargins(QMargins(0, 0, 0, 0))

        self._items = []
        self.justify_mode = justify_mode
        self.last_row_mode = last_row_mode

    # ------------------------------------------------------------------
    # QLayout API
    # ------------------------------------------------------------------

    def addItem(self, item):
        self._items.append(item)

    def count(self):
        return len(self._items)

    def itemAt(self, index):
        if 0 <= index < len(self._items):
            return self._items[index]
        return None

    def takeAt(self, index):
        if 0 <= index < len(self._items):
            return self._items.pop(index)
        return None

    def expandingDirections(self):
        return Qt.Orientation(0)

    def hasHeightForWidth(self):
        return True

    def heightForWidth(self, width):
        return self._do_layout(QRect(0, 0, width, 0), True)

    def setGeometry(self, rect):
        super().setGeometry(rect)
        self._do_layout(rect, False)

    def sizeHint(self):
        return self.minimumSize()

    def minimumSize(self):
        size = QSize()

        for item in self._items:
            size = size.expandedTo(item.minimumSize())

        margins = self.contentsMargins()

        size += QSize(
            margins.left() + margins.right(),
            margins.top() + margins.bottom(),
        )

        return size

    # ------------------------------------------------------------------
    # Internal
    # ------------------------------------------------------------------

    def _build_rows(self, available_width):
        rows = []

        current_row = []
        current_width = 0

        base_spacing = max(0, self.spacing())

        for item in self._items:
            width = item.sizeHint().width()

            if not current_row:
                current_row.append(item)
                current_width = width
                continue

            projected = current_width + base_spacing + width

            if projected > available_width:
                rows.append(current_row)

                current_row = [item]
                current_width = width
            else:
                current_row.append(item)
                current_width = projected

        if current_row:
            rows.append(current_row)

        return rows

    def _layout_row(
        self,
        row,
        rect,
        y,
        test_only,
        justify_mode,
        spacing_override=None,
    ):
        widths = [item.sizeHint().width() for item in row]
        heights = [item.sizeHint().height() for item in row]

        total_width = sum(widths)
        row_height = max(heights)

        count = len(row)
        gaps = max(0, count - 1)

        available_width = rect.width()

        x = rect.x()

        used_spacing = self.spacing()

        # --------------------------------------------------------------
        # Explicit spacing override (MATCH_PREVIOUS)
        # --------------------------------------------------------------

        if spacing_override is not None:
            used_spacing = spacing_override

            for item, w, h in zip(row, widths, heights):
                if not test_only:
                    item.setGeometry(QRect(int(round(x)), y, w, h))

                x += w + used_spacing

            return row_height, used_spacing

        # --------------------------------------------------------------
        # LEFT / NONE
        # --------------------------------------------------------------

        if justify_mode == JustifyMode.NONE:
            used_spacing = max(0, self.spacing())

            for item, w, h in zip(row, widths, heights):
                if not test_only:
                    item.setGeometry(QRect(int(round(x)), y, w, h))

                x += w + used_spacing

            return row_height, used_spacing

        # --------------------------------------------------------------
        # CENTER
        # --------------------------------------------------------------

        if justify_mode == JustifyMode.CENTER:
            used_spacing = max(0, self.spacing())

            content_width = total_width

            if count > 1:
                content_width += gaps * used_spacing

            x += max(0, (available_width - content_width) / 2)

            for item, w, h in zip(row, widths, heights):
                if not test_only:
                    item.setGeometry(QRect(int(round(x)), y, w, h))

                x += w + used_spacing

            return row_height, used_spacing

        # --------------------------------------------------------------
        # SPACE_*
        # --------------------------------------------------------------

        free_space = max(0, available_width - total_width)

        if justify_mode == JustifyMode.SPACE_BETWEEN:
            if gaps > 0:
                used_spacing = free_space / gaps
            else:
                used_spacing = 0

        elif justify_mode == JustifyMode.SPACE_AROUND:
            used_spacing = free_space / count if count else 0
            x += used_spacing / 2

        elif justify_mode == JustifyMode.SPACE_EVENLY:
            used_spacing = free_space / (count + 1) if count else 0
            x += used_spacing

        for item, w, h in zip(row, widths, heights):
            if not test_only:
                item.setGeometry(QRect(int(round(x)), y, w, h))

            x += w + used_spacing

        return row_height, used_spacing

    def _do_layout(self, rect, test_only):
        rows = self._build_rows(rect.width())

        y = rect.y()

        vertical_spacing = max(0, self.spacing())

        previous_spacing = None

        for row_index, row in enumerate(rows):
            is_last_row = row_index == len(rows) - 1

            spacing_override = None
            row_mode = self.justify_mode

            if is_last_row:

                if self.last_row_mode == LastRowMode.LEFT:
                    row_mode = JustifyMode.NONE

                elif self.last_row_mode == LastRowMode.CENTER:
                    row_mode = JustifyMode.CENTER

                elif self.last_row_mode == LastRowMode.JUSTIFY:
                    row_mode = self.justify_mode

                elif (
                    self.last_row_mode == LastRowMode.MATCH_PREVIOUS
                    and previous_spacing is not None
                ):
                    spacing_override = previous_spacing
                    row_mode = JustifyMode.NONE

            row_height, used_spacing = self._layout_row(
                row=row,
                rect=rect,
                y=y,
                test_only=test_only,
                justify_mode=row_mode,
                spacing_override=spacing_override,
            )

            previous_spacing = used_spacing

            y += row_height + vertical_spacing

        if rows:
            y -= vertical_spacing

        return y - rect.y()