#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
 Copyright 2019 by gniuk. All rights reserved.
 This file is part of the grimedit project:
 https://github.com/gniuk/grimedit
"""

import os
import sys
import time
import subprocess
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QWidget, QLabel, QMainWindow
from PyQt5.QtGui import QPixmap, QPainter, QBrush, QPen
# import ToolPanel
from PyQt5.QtWidgets import QHBoxLayout, QToolButton, QTextEdit, QFrame
from PyQt5.QtGui import QIcon, QScreen, QFont, QPolygon, QPainterPath
from PyQt5.QtCore import QSize, QLineF, QPointF

class Simed(QWidget):

    def __init__(self, image_path=None, image_content=None, phyScrH=1080):
        if image_path==None and image_content==None:
            print("No image input. Exit.")
            sys.exit(1)
        super().__init__()
        self.initUI(image_path, image_content, phyScrH)
        self.initStatus()
        self.initPanelUI()

    def initUI(self, image_path, image_content, phyScrH):
        if image_path:
            self.pixmap = QPixmap(image_path)
        else:
            self.pixmap = QPixmap()
            self.pixmap.loadFromData(image_content)

        if self.pixmap.height() + 35 > phyScrH:
            if self.pixmap.width() < 359:
                self.g_width = 359
                self.g_height = self.pixmap.height()
            else:
                self.g_width = self.pixmap.width()
                self.g_height = self.pixmap.height()
        else:
            if self.pixmap.width() < 359:
                self.g_width = 359
                self.g_height = self.pixmap.height()+35
            else:
                self.g_width = self.pixmap.width()
                self.g_height = self.pixmap.height()+35
        self.resize(self.g_width, self.g_height)
        self.setStyleSheet("background-color: transparent")

        self.setWindowTitle("Grimedit")

        self.g_border = QWidget(self)
        self.g_border.setStyleSheet("border: 2px dashed red")
        self.g_border.resize(self.pixmap.width(), self.pixmap.height())

        self.show()

    def initPanelUI(self):
        self.g_toolPanel = ToolPanel(self)
        self.g_toolPanel.setGeometry(self.rect().right()-359, self.rect().bottom()-35, 350, 35)
        self.g_toolPanel.resize(359, self.g_toolPanel.height())
        self.g_toolPanel.setParent(self)
        self.g_toolPanel.show()

    def initStatus(self):
        self.g_mousePressed = False
        self.g_mouseReleased = True
        self.g_drawStarted = False
        self.g_selectedTool = None
        self.g_isTextEditing = False
        self.startX = None
        self.startY = None
        self.g_drawShapeView = {
                              None: self.g_drawNoneView,
                              'Brush': self.g_drawBrushView,
                              'Rec': self.g_drawRecView,
                              'Elli': self.g_drawElliView,
                              'Text': self.g_drawTextView,
                              'Arrow': self.g_drawArrowView
                             }
        self.g_drawShapePixmap = {
                                'Brush': self.g_drawBrushPixmap,
                                'Rec': self.g_drawRecPixmap,
                                'Elli': self.g_drawElliPixmap,
                                'Text': self.g_drawTextPixmap,
                                'Arrow': self.g_drawArrowPixmap
                               }
        self.g_penStyle = QPen(Qt.red, 4)
        self.g_undoStack = []     # push pixmap snapshot into stack before a drawing on current pixmap
        self.g_textEdit = QTextEdit("", self)
        self.g_text = ""
        font = QFont()
        font.setPointSize(18)
        self.g_textEdit.setFont(font)
        self.g_textEdit.setStyleSheet("color:rgb(255,0,0); background:rgba(0,0,0,0%)")
        self.g_textEdit.textChanged.connect(self.g_onTextChanged)
        self.g_textEdit.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.g_textEdit.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        self.g_textEdit.setAlignment(Qt.AlignLeft)
        self.g_textEdit.setFrameStyle(QFrame.Box)

        self.g_curvepath = None

    def g_drawNoneView(self, painter):
        painter.setPen(self.g_penStyle)
        painter.drawPixmap(self.pixmap.rect(), self.pixmap)

    def g_drawLineView(self, painter):
        """Draw on mainwindow Widget, it is fake drawing, just for viewing.
        """
        painter.setPen(self.g_penStyle)
        painter.drawPixmap(self.pixmap.rect(), self.pixmap)
        if self.g_mousePressed:
            painter.drawLine(self.startX,self.startY, self.endX,self.endY)

    def g_drawLinePixmap(self):
        """Draw on pixmap, do real pixel modify.
        """
        painter = QPainter(self.pixmap) # draw on pixmap
        painter.setPen(self.g_penStyle)
        painter.drawLine(self.startX,self.startY, self.endX,self.endY)
        painter.end()
        self.update()

    def g_drawBrushView(self, painter):
        painter.setPen(self.g_penStyle)
        painter.drawPixmap(self.pixmap.rect(), self.pixmap)
        if self.g_mousePressed:
            painter.drawPath(self.g_curvepath)

    def g_drawBrushPixmap(self):
        painter = QPainter(self.pixmap) # draw on pixmap
        painter.setPen(self.g_penStyle)
        painter.drawPath(self.g_curvepath)
        painter.end()
        self.update()

    def g_calcArrowPolygons(self):
        g_line = QLineF(self.startX, self.startY, self.endX, self.endY)
        v1 = g_line.unitVector()
        v1.setLength(18.66)        # This length and the _v1 length can decide the arrow size
        g_lineLength = g_line.length()
        v2 = g_line.unitVector()
        v2.setLength(g_lineLength - 18.66)
        v1.translate(QPointF(v2.dx(), v2.dy()))
        n1 = v1.normalVector()
        n1.setLength(10.77)      # 18.66*1.732/3
        n2 = n1.normalVector().normalVector()
        _g_line = QLineF(self.endX, self.endY, self.startX, self.startY) # The reverse line to get p2
        _v1 = _g_line.unitVector()
        _v1.setLength(10)
        # return g_line.p2(), _g_line.p2(), n1.p2(), n2.p2() # This is a good mistake :)
        # return g_line.p2(), _v1.p2(), n1.p2(), n2.p2() # This is what I originally want

        # Let's do more work to simulate WxWork or WeChat
        p1=g_line.p2(); p2=_v1.p2(); p3=n1.p2(); p4=n2.p2();
        p5 = QLineF(p2,p3).center(); p6 = QLineF(p2,p4).center()
        p7 = g_line.p1()
        return p1,p2,p3,p4,p5,p6,p7

    def g_drawArrowView(self, painter):
        painter.setPen(QPen(Qt.red, 1))
        painter.setBrush(QBrush(Qt.red))
        painter.drawPixmap(self.pixmap.rect(), self.pixmap)
        if self.g_mousePressed:
            p1,p2,p3,p4,p5,p6,p7 = self.g_calcArrowPolygons()
            #       p1
            # p3_ _ /\
            #      /  \
            #     /   p4
            #    p7
            # painter.drawLine(self.startX,self.startY, self.endX,self.endY)
            # painter.drawPolygon(p1, p2, p3)
            # painter.drawPolygon(p1, p2, p4)
            # painter.drawPolygon(p7, p5, p2, p6)
            painter.drawPolygon(p1, p3, p5, p7, p6, p4)
    def g_drawArrowPixmap(self):
        painter = QPainter(self.pixmap) # draw on pixmap
        painter.setPen(QPen(Qt.red, 1))
        painter.setBrush(QBrush(Qt.red))
        p1,p2,p3,p4,p5,p6,p7 = self.g_calcArrowPolygons()
        # painter.drawLine(self.startX,self.startY, self.endX,self.endY)
        # painter.drawPolygon(p1, p2, p3)
        # painter.drawPolygon(p1, p2, p4)
        # painter.drawPolygon(p7, p5, p2, p6)
        painter.drawPolygon(p1, p3, p5, p7, p6, p4)
        painter.end()
        self.update()

    def g_drawRecView(self,painter):
        painter.setPen(self.g_penStyle)
        painter.drawPixmap(self.pixmap.rect(), self.pixmap)
        if self.g_mousePressed:
            painter.drawRect(self.startX,self.startY, self.endX-self.startX,self.endY-self.startY)

    def g_drawRecPixmap(self):
        painter = QPainter(self.pixmap) # draw on pixmap
        painter.setPen(self.g_penStyle)
        painter.drawRect(self.startX,self.startY, self.endX-self.startX,self.endY-self.startY)
        painter.end()
        self.update()

    def g_drawElliView(self,painter):
        painter.setPen(self.g_penStyle)
        painter.drawPixmap(self.pixmap.rect(), self.pixmap)
        if self.g_mousePressed:
            painter.drawEllipse(self.startX,self.startY, self.endX-self.startX,self.endY-self.startY)

    def g_drawElliPixmap(self):
        painter = QPainter(self.pixmap) # draw on pixmap
        painter.setPen(self.g_penStyle)
        painter.drawEllipse(self.startX,self.startY, self.endX-self.startX,self.endY-self.startY)
        painter.end()
        self.update()

    def g_drawTextView(self,painter):
        painter.setPen(self.g_penStyle)
        painter.drawPixmap(self.pixmap.rect(), self.pixmap)
        # the view is drawing automatically by QTextEdit's show.

    def g_drawTextPixmap(self, oldX=None, oldY=None):
        if oldX != None:
            painter = QPainter(self.pixmap) # draw on pixmap
            painter.setPen(self.g_penStyle)
            font = QFont()
            font.setPointSize(18)
            painter.setFont(font)
            painter.drawText(oldX+5, oldY+33, self.g_text) # the view of text not aligned with the draw, there may be some better method to deal with this problem.
            painter.end()
            self.update()
            self.g_textEdit.hide()

    def paintEvent(self, e):
        painter = QPainter(self)
        self.g_drawShapeView[self.g_selectedTool](painter)
        painter.end()

    def keyPressEvent(self, e):
        """Override the default key press event handler.
        """
        # print(dir(e))
        if e.key() == Qt.Key_Escape:
            self.g_quit()
        if e.modifiers() & Qt.ControlModifier:
            if e.key() == Qt.Key_Z:
                self.g_undoDraw()

    def g_undoDraw(self):
        if self.g_undoStack:
            self.pixmap = self.g_undoStack.pop()
            print("old pixmap at: {}".format(self.pixmap))
            self.update()

    def genImagePath(self):
        picDir = os.path.expanduser('~') + '/Pictures'
        try:
            # Maybe the best compatibility to mkdir in python :)
            os.mkdir(picDir)
        except:
            pass
        screentShotDir = picDir + '/ScreenShots'
        try:
            os.mkdir(screentShotDir)
        except:
            pass
        pic = "".join([screentShotDir, '/', time.strftime("%Y%m%d-%H-%M-%S"), '.png'])
        return pic

    def g_saveWithClipboard(self):
        pic = self.genImagePath()
        self.pixmap.save(pic, 'PNG')
        with open(pic, 'rb') as f:
            subprocess.Popen(['wl-copy'], stdin=f)
        self.g_quit()

    def g_saveToClipboard(self):
        pic = self.genImagePath()
        self.pixmap.save(pic, 'PNG')
        with open(pic, 'rb') as f:
            subprocess.Popen(['wl-copy'], stdin=f)
        try:
            os.remove(pic)
        except:
            pass
        self.g_quit()

    def g_quit(self):
        self.g_toolPanel.close()
        self.close()

    def g_doTextSave(self, oldX, oldY):
        self.g_text = self.g_textEdit.toPlainText()
        if len(self.g_text):
            self.g_textEdit.setText("")
            self.g_undoStack.append(QPixmap(self.pixmap))
            self.g_drawTextPixmap(oldX, oldY)
        self.g_isTextEditing = False
        self.g_textEdit.hide()

    def mousePressEvent(self, e):
        print(type(e))
        oldX = self.startX
        oldY = self.startY
        self.startX = e.x()
        self.startY = e.y()
        # print(self.startX, self.startY)
        self.g_mousePressed = True

        if self.g_selectedTool == "Text":
            if self.g_isTextEditing == False:
                self.g_isTextEditing = True
                self.g_textEdit.setGeometry(self.startX, self.startY, 60, 40)
                self.g_textEdit.show()
                # g_cursor = self.g_textEdit.textCursor()
                # self.g_textEdit.setTextCursor(g_cursor)
                # TODO: how to automatically get into TextEdit with the cursor showup?
            else:
                self.g_doTextSave(oldX, oldY)
        else:
            if self.g_selectedTool == "Brush":
                self.g_curvepath = QPainterPath()
                self.g_curvepath.moveTo(self.startX, self.startY)
            self.g_doTextSave(oldX, oldY)
            self.endX = self.startX
            self.endY = self.startY

    def mouseReleaseEvent(self, e):
        self.endX = e.x()
        self.endY = e.y()
        self.g_mousePressed = False

        # save the pixmap snapshot for undo
        if self.g_selectedTool:
            if self.endX != self.startX and self.endY != self.startY:
                self.g_undoStack.append(QPixmap(self.pixmap))
                print("modified pixmap at {}, size = {} Bytes".format(self.pixmap, self.pixmap.__sizeof__()))
                self.g_drawShapePixmap[self.g_selectedTool]()

    def mouseMoveEvent(self, e):
        # print(dir(e))
        self.endX = e.x(); self.endY = e.y()
        # print("mousebuttonpress:{},{}".format(e.MouseButtonPress, type(e.MouseButtonPress)))
        if self.g_mousePressed:
            if self.g_selectedTool == "Brush":
                self.g_curvepath.lineTo(self.endX, self.endY)
            # print(type(e),type(e.MouseButtonPress))
            self.update()

    def g_onTextChanged(self):
        print("---------- text changed ----------")
        self.g_textEdit.resize(len(self.g_textEdit.toPlainText())*18+60, int(self.g_textEdit.document().size().height())+2)

    def g_selectRect(self):
        self.g_selectedTool = 'Rec'
        print("Rectangle Tool selected")

    def g_selectLine(self):
        self.g_selectedTool = 'Line'
        print("Line Tool selected")

    def g_selectArrow(self):
        self.g_selectedTool = 'Arrow'
        print("Arrow Tool selected")

    def g_selectEllipse(self):
        self.g_selectedTool = 'Elli'
        print("Ellipse Tool selected")

    def g_selectText(self):
        self.g_selectedTool = 'Text'
        print("Text Tool selected")

    def g_selectBrush(self):
        self.g_selectedTool = 'Brush'
        print("Brush Tool selected")


class ToolPanel(QWidget):
    def __init__(self, mainWindow):
        super().__init__()
        self.mWindow = mainWindow
        self.initToolLayout()

    def initToolLayout(self):
        icondir = os.path.dirname(os.path.abspath(__file__)) + '/icons/'
        if not os.path.exists(icondir):
            icondir = '/usr/share/grimedit/icons/'
        self.rectBtn = QToolButton(self)
        self.rectBtn.setIcon(QIcon(icondir+'rect.png'))
        self.rectBtn.setFixedSize(35,35)
        self.rectBtn.setIconSize(QSize(22,22))
        self.rectBtn.clicked.connect(self.mWindow.g_selectRect)

        self.elliBtn = QToolButton(self)
        self.elliBtn.setIcon(QIcon(icondir+'ellipse.png'))
        self.elliBtn.setFixedSize(35,35)
        self.elliBtn.setIconSize(QSize(23,23))
        self.elliBtn.clicked.connect(self.mWindow.g_selectEllipse)

        self.arrowBtn = QToolButton(self)
        self.arrowBtn.setIcon(QIcon(icondir+'arrow.png'))
        self.arrowBtn.setFixedSize(35,35)
        self.arrowBtn.setIconSize(QSize(22,22))
        self.arrowBtn.clicked.connect(self.mWindow.g_selectArrow)

        self.brushBtn = QToolButton(self)
        self.brushBtn.setIcon(QIcon(icondir+'brush.png'))
        self.brushBtn.setFixedSize(35,35)
        self.brushBtn.setIconSize(QSize(24,24))
        self.brushBtn.clicked.connect(self.mWindow.g_selectBrush)
        self.mosaicBtn = QToolButton(self)
        self.mosaicBtn.setIcon(QIcon(icondir+'mosaic.png'))
        self.mosaicBtn.setFixedSize(35,35)
        self.mosaicBtn.setIconSize(QSize(22,22))

        self.textBtn = QToolButton(self)
        self.textBtn.setIcon(QIcon(icondir+'text.png'))
        self.textBtn.setFixedSize(35,35)
        self.textBtn.setIconSize(QSize(22,22))
        self.textBtn.clicked.connect(self.mWindow.g_selectText)

        self.undoBtn = QToolButton(self)
        self.undoBtn.setIcon(QIcon(icondir+'undo.png'))
        self.undoBtn.setFixedSize(35,35)
        self.undoBtn.setIconSize(QSize(23,23))
        self.undoBtn.clicked.connect(self.mWindow.g_undoDraw)

        self.saveBtn = QToolButton(self)
        self.saveBtn.setIcon(QIcon(icondir+'save.png'))
        self.saveBtn.setFixedSize(35,35)
        self.saveBtn.setIconSize(QSize(22,22))
        self.saveBtn.clicked.connect(self.mWindow.g_saveWithClipboard)

        self.cancelBtn = QToolButton(self)
        self.cancelBtn.setIcon(QIcon(icondir+'cancel.png'))
        self.cancelBtn.setFixedSize(35,35)
        self.cancelBtn.setIconSize(QSize(22,22))
        self.cancelBtn.clicked.connect(self.mWindow.g_quit)

        self.finishBtn = QToolButton(self)
        self.finishBtn.setIcon(QIcon(icondir+'finish.png'))
        self.finishBtn.setFixedSize(35,35)
        self.finishBtn.setIconSize(QSize(25,23))
        self.finishBtn.clicked.connect(self.mWindow.g_saveToClipboard)

        self.hLayout = QHBoxLayout()
        self.hLayout.addWidget(self.rectBtn)
        self.hLayout.addWidget(self.elliBtn)
        self.hLayout.addWidget(self.arrowBtn)
        self.hLayout.addWidget(self.brushBtn)
        self.hLayout.addWidget(self.mosaicBtn)
        self.hLayout.addWidget(self.textBtn)
        self.hLayout.addWidget(self.undoBtn)
        self.hLayout.addWidget(self.saveBtn)
        self.hLayout.addWidget(self.cancelBtn)
        self.hLayout.addWidget(self.finishBtn)
        self.hLayout.setSpacing(1)
        self.hLayout.setContentsMargins(0,0,0,0)

        self.setLayout(self.hLayout)
        self.setStyleSheet("background:rgba(55,55,55,50%)")


if __name__ == '__main__':
    app = QApplication([])
    phyScreenHeight = app.primaryScreen().size().height()
    print("phyScreenHeight: {}".format(phyScreenHeight))

    if len(sys.argv) > 1:
        image_path = sys.argv[1]
        image_content = None
    else:
        image_path = None
        with open(0, 'rb') as f:
            image_content = f.read()
            if not image_content:
                sys.exit(1)
    simed = Simed(image_path, image_content, phyScreenHeight)
    sys.exit(app.exec_())
