from qgis.PyQt.QtCore import QPoint, QRect, Qt
from qgis.PyQt.QtGui import (
    QColor,
    QGuiApplication,
    QKeySequence,
    QMouseEvent,
    QPainter,
    QPaintEvent,
    QPen,
    QPixmap,
    QShortcut,
)
from qgis.PyQt.QtWidgets import QApplication, QWidget


class ScreenshotMakerWidget(QWidget):
    def __init__(self, parent: QWidget | None = None) -> None:
        """Widget for selection of screenshot area and save of selected screen content to QPixmap

        :param parent: widget parent, defaults to None
        :type parent: Optional[QWidget], optional
        """
        super().__init__(parent)
        # Start and end position of screenshot area selection
        self._start: QPoint | None = None
        self._end: QPoint | None = None
        self._result: QPixmap | None = None

        self._pen_size = 2
        self._pen_color = QColor(255, 0, 0)

        # Define geometry with all screen (use virtualGeometry)
        desktop_geom = QGuiApplication.primaryScreen().virtualGeometry()
        self.setGeometry(desktop_geom)

        # Flags needed for selection
        # Qt.WindowType.FramelessWindowHint : no title bar, easiest cursor position managment
        # Qt.WindowType.SplashScreen : needed to be able to use multiple screen
        self.setWindowFlags(
            Qt.WindowType.SplashScreen | Qt.WindowType.FramelessWindowHint
        )

        # Transparent Widget
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)

        # Modal widget to avoid use of other widget
        self.setWindowModality(Qt.WindowModality.ApplicationModal)

        # Cursor for selection
        self.setCursor(Qt.CursorShape.CrossCursor)

        # Shortcut for direct close and no image save
        self.shortcut_escape: QShortcut = QShortcut(
            QKeySequence(Qt.Key.Key_Escape), self
        )
        self.shortcut_escape.activated.connect(self.close)

        # Display for selection, paintEvent overriden to display current selection
        self.show()

    def paintEvent(self, _: QPaintEvent) -> None:  # noqa: N802
        """Paint event to display the current selection area.

        :param _: Paint event
        :type _: QPaintEvent
        """
        # Dark overlay on all screen content
        painter = QPainter(self)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        painter.fillRect(self.rect(), QColor(0, 0, 0, 100))  # overlay sombre

        if self._start and self._end:
            # Update selected area clear current overlay
            rect = QRect(self._start, self._end).normalized()

            painter.setCompositionMode(QPainter.CompositionMode.CompositionMode_Clear)
            painter.fillRect(rect, QColor())

            # Display rect with selection
            painter.setCompositionMode(
                QPainter.CompositionMode.CompositionMode_SourceOver
            )
            painter.setPen(QPen(self._pen_color, self._pen_size))
            painter.drawRect(rect)

        painter.end()

    def mousePressEvent(self, event: QMouseEvent) -> None:  # noqa: N802
        """Override mousePressEvent to define selection area start point

        :param event: mouve event to get position
        :type event: QMouseEvent
        """
        self._start = event.pos()
        self._end = self._start

    def mouseMoveEvent(self, event: QMouseEvent) -> None:  # noqa: N802
        """Override mouseMoveEvent to define selection area end point

        :param event: mouve event to get position
        :type event: QMouseEvent
        """
        if self._start is None:
            return
        self._end = event.pos()
        self.update()

    def mouseReleaseEvent(self, event: QMouseEvent) -> None:  # noqa: N802
        """Override mouseReleaseEvent to stop selection and save selected area content to QPixmap

        :param event: mouve event to get position
        :type event: QMouseEvent
        """
        if self._start is None:
            return
        self._end = event.pos()

        # Define selected area
        rect = QRect(self._start, self._end).normalized()

        # Remove displayed rect
        rect = QRect(
            rect.topLeft() + QPoint(self._pen_size, self._pen_size),
            rect.bottomRight() - QPoint(self._pen_size, self._pen_size),
        )

        # Convert position to global position
        top_left_global = self.mapToGlobal(rect.topLeft())
        global_rect = QRect(top_left_global, rect.size())

        # Get current selection QPixmap
        self._result = QGuiApplication.primaryScreen().grabWindow(
            0,
            global_rect.x(),
            global_rect.y(),
            global_rect.width(),
            global_rect.height(),
        )

        # Close selection widget
        self.close()

    def get_result(self) -> QPixmap | None:
        """Return selected area screenshot as QPixmap if available

        :return: selected area QPixmap
        :rtype: Optional[QPixmap]
        """
        return self._result

    def clear_result(self) -> None:
        """Clear current result"""
        self._result = None


def take_screenshot() -> QPixmap | None:
    """
    Take screenshot and returned selected area as QPixmap if available
    """
    app = QApplication.instance()
    if app is None:
        return None

    screenshot_widget = ScreenshotMakerWidget(None)

    # Wait for use selection
    while screenshot_widget.isVisible():
        app.processEvents()

    # Return result
    return screenshot_widget.get_result()
