from krita import Extension
from PyQt5.QtWidgets import QDialog, QVBoxLayout, QHBoxLayout, QLabel, QLineEdit, QComboBox, QCheckBox, QPushButton
from PyQt5.QtCore import Qt, QByteArray
from math import ceil
import re

class BookletArranger(Extension):

    def __init__(self, parent):
        super().__init__(parent)

    # Krita.instance() exists, so do any setup work
    def setup(self):
        pass

    # called after setup(self)
    def createActions(self, window):
        # Create menu item in Tools > Scripts.
        action = window.createAction("bookletarranger", "Booklet Arranger")
        action.triggered.connect(self.booklet_arranger)

    def booklet_arranger(self):
        # Dialog creation.
        newDialog = QDialog()
        newDialog.setWindowTitle("Booklet Arranger")
        layout = QVBoxLayout()
        newDialog.setLayout(layout)

        # Description row.
        desc = QLabel('Arranges a set of layers into a booklet.')
        desc.setAlignment(Qt.AlignCenter)
        row0 = QHBoxLayout()
        row0.addWidget(desc)
        layout.addLayout(row0)

        # Page Information row.
        preflight_info = self.preflightInfo()
        info_text = preflight_info["InfoText"]
        self.page_count = preflight_info["PageCount"]
        stats = QLabel(info_text)
        stats.setGeometry(200, 250, 100, 50)
        stats.setStyleSheet("border: 1px solid white")
        stats.setTextFormat(Qt.RichText)
        stats.setWordWrap(True)
        rown = QHBoxLayout()
        rown.addWidget(stats)
        layout.addLayout(rown)

        # Options row. Bleed, cropmarks, outer-page offset.
        self.bleedInput = QLineEdit('24')
        self.cropmarksCheck = QCheckBox('Cropmarks')
        self.cropmarksCheck.setChecked(True)
        self.offsetOuterCheck = QCheckBox('Offset Outer Pages')
        self.offsetOuterCheck.setChecked(True)
        row2 = QHBoxLayout()
        row2.addWidget(QLabel('Bleed size:'))
        row2.addWidget(self.bleedInput)
        row2.addWidget(self.cropmarksCheck)
        row2.addWidget(self.offsetOuterCheck)
        layout.addLayout(row2)

        # Do It row.
        goButton = QPushButton("Arrange Pages")
        goButton.setIcon( Krita.instance().icon('animation_play') )
        if not preflight_info["Success"]:
            goButton.setEnabled(False)
        row4 = QHBoxLayout()
        row4.addWidget(goButton)
        layout.addLayout(row4)

        # Hook up the actions.
        goButton.clicked.connect( self.arrangeBooklet )

        # Show the dialog.
        newDialog.exec_()

    # Run some checks and generate the preflight message to brief the user.
    def preflightInfo(self):
        # - Check number of selected layers
        # - Check layer names (valid forms are "Page 2" or "2")
        # - Check layer dimensions for irregularities.
        success = True
        info_text = ''

        w = Krita.instance().activeWindow()
        v = w.activeView()
        layers = v.selectedNodes()

        # Check layer count. Should be more than 1, and a multiple of 4.
        # In the future, will detect missing pages and possibly pad with blank pages (like page2 and second last page inside cover pages).
        page_count = len(layers)
        message = ''
        if page_count < 4:
            message = " - not enough selected! Minimum of 4 pages."
            success = False
        elif page_count % 4 > 0:
            message = " - incorrect page count! Must be multiple of 4."
            success = False
        info_text += 'Pages Selected: ' + str(page_count) + message + '<br/>'

        # Check for complete sequence of page numbers. Grabs the first number in a layer name
        invalid_names = ''
        page_numbers = []

        ## Collect each layer's first number as a page number. Report on names without numbers.
        for layer in layers:
            name = layer.name()
            match = re.search('[0-9]+', name)
            if match:
                page_numbers.append(int(match[0]))
            else:
                invalid_names += name + " "
        if invalid_names:
            info_text += 'Invalid Layer Names: ' + invalid_names + '<br>'
            success = False

        ## Check for missing page numbers. Pages must fulfil the range "1 to page_count".
        highest_page = max(page_numbers)
        last_page = highest_page if highest_page > page_count else page_count
        missing_pages = ''
        for p in range(1, last_page+1):
            if p in page_numbers:
                pass
            else:
                missing_pages += str(p) + " "
        if missing_pages:
            info_text += 'Missing Page Numbers: ' + missing_pages + '<br>'
            success = False

        return {"Success": success, "InfoText": info_text, "PageCount": page_count}

    ##########
    # Slots
    ##########

    def arrangeBooklet(self, e):

        # Document dimensions.
        self.doc = Krita.instance().activeDocument()
        self.doc_width = self.doc.width()
        self.doc_height = self.doc.height()

        # Get the selected layer(s).
        w = Krita.instance().activeWindow()
        v = w.activeView()
        layers = v.selectedNodes()

        for layer in layers:
            if layer.type() == 'paintlayer':
                self.arrangePage(layer)
                didNotRun = False

        if didNotRun:
            dialog = QDialog()
            dialog.setWindowTitle("Paint Layer Required")
            layout = QVBoxLayout()
            layout.addWidget(QLabel('Booklet Arranger only works on paint layers.'))
            dialog.setLayout(layout)
            dialog.exec_()
        else:
            # Refresh the view, or the cropmarks will not be immediately shown.
            self.doc.refreshProjection()

    # Actually arranges the pages into a booklet!
    def arrangePage(self, page):
        max_sheets = self.page_count / 4
        page_name = page.name()
        match = re.search('[0-9]+', page_name)
        page_number = int(match[0])

        # Is the page in the higher or lower half? e.g. p1-16 or p17-32?
        # Calculation differs based on this.
        high = page_number > (self.page_count / 2)

        # Work out which paper sheet/side a page should be on.
        # e.g. for 32 pages: p32 p1/p2 p31 will be on sheet 1 sides A/B.
        #
        # Side number 1 = sheet 1 side A (front).
        # Side number 2 = sheet 1 side B (back).
        # Side number 3 = sheet 2 side A (front).
        # ...
        #
        if high:
            side_number = self.page_count - page_number + 1
        else:
            side_number = page_number
        current_sheet = ceil(side_number / 2)

        # Work out if a page should be on the left or right side of the sheet.
        front_side = side_number % 2 # Odds sheet/sides (1, 3, 5...) are front.
        right = True
        if (high and front_side) or (not high and not front_side):
            right = False

        # Add page to correct group (which represent paper sheets/sides).
        side_name = "p" + str(current_sheet)
        side_name += "a" if front_side else "b"
        sidegroup = self.doc.nodeByName(side_name)
        if not sidegroup:
            sidegroup = self.doc.createGroupLayer(side_name)
            root = self.doc.rootNode()
            root.addChildNode(sidegroup, None)
        page.remove()
        sidegroup.addChildNode(page, None)

        # Offset calculation. First sheet has largest offset, centre sheet has zero offset. 3 pixel step.
        if self.offsetOuterCheck.isChecked():
            offset = 3 * (max_sheets - current_sheet)
        else:
            offset = 0

        # Position page (with or without offset)
        page_pos = page.position()
        bds = page.bounds()
        x_target = self.doc_width / 2
        x_diff = x_target - bds.left()
        orig_x = x_diff
        y_target = (self.doc_height - (bds.bottom() - bds.top())) / 2
        y_diff = y_target - bds.top()
        orig_y = y_diff
        if right:
            x_diff = x_diff + offset
        else:
            x_diff = x_diff - (bds.right() + 1 - bds.left()) - offset
        page.move(int(page_pos.x() + x_diff), int(page_pos.y() + y_diff))

        # Keep pre-bleed bounds for crop marks.
        bds = page.bounds()

        # Bleed by stated amount on outward sides, bleed to middle on inward sides.
        bleed_amount = int(self.bleedInput.text())
        bleed_amounts = { "top": bleed_amount, "bottom": bleed_amount }
        if right:
            bleed_amounts["right"] = bleed_amount
            bleed_amounts["left"] = offset
        else:
            bleed_amounts["right"] = offset
            bleed_amounts["left"] = bleed_amount
        self.bleedPage(page, bleed_amounts)

        # Add crop marks. Which crop marks are created depends on which side the page is on.
        if self.cropmarksCheck.isChecked():
            black_pixel = b'\x00\x00\x00\xff'
            crop_mark_width = 2
            crop_mark_height = 48
            crop_mark = QByteArray(black_pixel * crop_mark_width * crop_mark_height)
            ## Vertical cropmarks
            x = bds.right() if right else bds.left() - 1
            ytop = bds.top() - crop_mark_height - bleed_amount
            page.setPixelData(crop_mark, x, ytop, crop_mark_width, crop_mark_height)
            ybot = bds.bottom() + 1 + bleed_amount
            page.setPixelData(crop_mark, x, ybot, crop_mark_width, crop_mark_height)
            if page_number == 1:
                # Specially mark centre on first page (always p1a right side)
                x = int(bds.left() - 1 - offset)
                page.setPixelData(crop_mark, x, ytop, crop_mark_width, crop_mark_height)
                page.setPixelData(crop_mark, x, ybot, crop_mark_width, crop_mark_height)

            ## Horizontal cropmarks
            crop_mark_width = 48
            crop_mark_height = 2
            x = bds.right() + 1 + bleed_amount if right else bds.left() - bleed_amount - crop_mark_width
            y = bds.top() - 1
            page.setPixelData(crop_mark, x, y, crop_mark_width, crop_mark_height)
            y = bds.bottom()
            page.setPixelData(crop_mark, x, y, crop_mark_width, crop_mark_height)

    def bleedPage(self, page, amounts):

        # Top
        bds = page.bounds()
        xpos = bds.left()
        ypos = bds.top()
        width = bds.width()
        height = amounts['top']
        bleed_line = page.pixelData(xpos, ypos, width, 1)
        bleed = bleed_line * height
        page.setPixelData(bleed, xpos, ypos - height, width, height)

        # Bottom
        bds = page.bounds()
        xpos = bds.left()
        ypos = bds.bottom()
        width = bds.width()
        height = amounts['bottom']
        bleed_line = page.pixelData(xpos, ypos, width, 1)
        bleed = bleed_line * height
        page.setPixelData(bleed, xpos, ypos + 1, width, height)

        # Left
        bds = page.bounds()
        xpos = bds.left()
        ypos = bds.top()
        width = int(amounts['left'])
        height = bds.height()
        bleed_line = page.pixelData(xpos, ypos, 1, height)
        for c in range(1, width + 1):
            page.setPixelData(bleed_line, (xpos - c), ypos, 1, height)

        # Right
        bds = page.bounds()
        xpos = bds.right()
        ypos = bds.top()
        width = int(amounts['right'])
        height = bds.height()
        bleed_line = page.pixelData(xpos, ypos, 1, height)
        for c in range(1, width + 1):
            page.setPixelData(bleed_line, (xpos + c), ypos, 1, height)