#!/usr/bin/python

"""
    Utility to take a snapshot of the screen
"""

import sys
from PyQt5.QtGui import QGuiApplication, QPainter, QPen
from PyQt5.QtCore import QRect, Qt, pyqtSignal, QBuffer
from PyQt5.QtWidgets import QApplication, QWidget


class RegionSelecter(QWidget):
    sigBitmap = pyqtSignal(bytes)
    def __init__(self):
        super().__init__()
        screen = QGuiApplication.primaryScreen()
        self.desktopPixmap = screen.grabWindow(0)
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.X11BypassWindowManagerHint | Qt.WindowStaysOnTopHint)
        self.setWindowState(Qt.WindowFullScreen)
        self.move(0, 0)
        self.resize(self.desktopPixmap.size())
        self.drawSelection = False

    @property
    def Selection(self):
        return QRect(self.startPoint, self.endPoint).normalized()

    def paintEvent(self, ev):
        painter = QPainter(self)
        painter.drawPixmap(0, 0, self.desktopPixmap)
        if self.drawSelection:
            painter.setPen(QPen(Qt.red))
            painter.drawRect(self.Selection)

    def mousePressEvent(self, ev):
        self.startPoint = ev.pos()
        self.endPoint = ev.pos()
        self.drawSelection = True

    def mouseMoveEvent(self, ev):
        self.update()
        self.endPoint = ev.pos()

    def mouseReleaseEvent(self, ev):
        self.endPoint = ev.pos()
        pm = self.desktopPixmap.copy(self.Selection)
        memfile = QBuffer()
        pm.save(memfile, "png")
        self.sigBitmap.emit(bytes(memfile.data()))


def takeShot():
    app = QApplication(sys.argv)
    screen = QGuiApplication.primaryScreen()
    desktopPixmap = screen.grabWindow(0)
    desktopPixmap.save('tst2.png')


def takeRegion():
    """ Take a snapshot of the current desktop """
    app = QApplication(sys.argv)
    r = RegionSelecter()
    TMP = {}
    def handleBitmap(bm):
        app.quit()
        TMP['bm'] = bm
    r.sigBitmap.connect(handleBitmap)
    r.show()
    app.exec()
    return TMP['bm']


if __name__ == '__main__':
    bm = takeRegion()
    print(bm)
    with open('tst.png', 'wb') as f:
        f.write(bm)