#!/usr/bin/env python
# coding=utf-8
'''
Author: ShenMeng
Date: 2021-09-02 17:04:46
FilePath: \kmoke\smart_qtwidgets\number_plain_text_edit.py
LastEditTime: 2021-09-03 09:39:53
'''

from qtpy import QtWidgets, QtGui, QtCore


class NumberBar(QtWidgets.QWidget):
    def __init__(self, parent=None):
        super(NumberBar, self).__init__(parent)
        self._parent = parent

    def sizeHint(self):
        return QtCore.QSize(self._parent.number_bar_width(), 0)

    def paintEvent(self, event):
        self._parent.number_bar_paint(event)


class NumberPlainTextEdit(QtWidgets.QPlainTextEdit):
    def __init__(self, parent=None):
        super(NumberPlainTextEdit, self).__init__(parent)
        self.number_bar = NumberBar(self)
        self.blockCountChanged.connect(self.update_number_bar_width)
        self.updateRequest.connect(self.update_number_bar)
        self.cursorPositionChanged.connect(self.highlight_current_line)
        self.setLineWrapMode(QtWidgets.QPlainTextEdit.NoWrap)
        self.__background = [255, 255, 255, 10]
        self.__color = [255, 255, 255, 165]
        self.__current_color = [255, 255, 255, 20]
        self.__space = 12
        self.update_number_bar_width()

    def set_number_background(self, color):
        self.__background = color

    def set_number_color(self, color):
        self.__color = color

    def set_current_color(self, color):
        self.__current_color = color

    def set_space(self, num):
        self.__space = num
        self.update_number_bar_width()

    def number_bar_width(self):
        num = 1
        max_value = max(1, self.blockCount())
        while max_value >= 10:
            max_value /= 10
            num += 1
        return self.fontMetrics().width('g') * num + self.__space * 2

    def update_number_bar_width(self):
        self.setViewportMargins(self.number_bar_width(), 0, 0, 0)

    def update_number_bar(self, rect, y):
        if y:
            self.number_bar.scroll(0, y)
        else:
            self.number_bar.update(0, rect.y(), self.number_bar.width(),
                                   rect.height())
        if rect.contains(self.viewport().rect()):
            self.update_number_bar_width()

    def highlight_current_line(self):
        extra_selections = []
        if not self.isReadOnly():
            # selection = QtWidgets.QPlainTextEdit.ExtraSelection()
            selection = QtWidgets.QTextEdit.ExtraSelection()
            lineColor = QtGui.QColor(*self.__current_color).lighter(160)
            selection.format.setBackground(lineColor)
            selection.format.setProperty(
                QtGui.QTextFormat.FullWidthSelection, True)
            selection.cursor = self.textCursor()
            selection.cursor.clearSelection()
            extra_selections.append(selection)
        self.setExtraSelections(extra_selections)

    def number_bar_paint(self, event):
        painter = QtGui.QPainter(self.number_bar)
        painter.fillRect(event.rect(), QtGui.QColor(*self.__background))
        block = self.firstVisibleBlock()
        block_number = block.blockNumber()
        top = self.blockBoundingGeometry(block).translated(
            self.contentOffset()).top()
        bottom = top + self.blockBoundingRect(block).height()

        height = self.fontMetrics().height()
        while block.isValid() and (top <= event.rect().bottom()):
            if block.isVisible() and (bottom >= event.rect().top()):
                number = str(block_number + 1)
                painter.setPen(QtGui.QColor(*self.__color))
                painter.drawText(0, top,
                                 self.number_bar.width() - self.__space,
                                 height, QtCore.Qt.AlignRight, number)
            block = block.next()
            top = bottom
            bottom = top + self.blockBoundingRect(block).height()
            block_number += 1

    def resizeEvent(self, event):
        super(NumberPlainTextEdit, self).resizeEvent(event)
        _rect = self.contentsRect()
        self.number_bar.setGeometry(
            QtCore.QRect(_rect.left(), _rect.top(), self.number_bar_width(),
                         _rect.height()))