#!/usr/bin/env python3
# -*- coding:UTF-8 -*-
# Author: nigo
import enum
from comformation import ComformationUI
from comformation import ComformationTemplateUI
from comformation import MatchOneUI
from comformation import MethodUI
from comformation import ComformationAppUI
# from comformations import Ui_Dialog
# from template import Ui_Form
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
import sys
import os
import re
from docx import Document
from docxcompose.composer import Composer
from docx.table import _Cell
from docx.oxml import OxmlElement
from docx.oxml.ns import qn
from openpyxl import Workbook, load_workbook
import pandas as pd


class ComformWidget(QWidget, ComformationAppUI):
    def __init__(self):
        super().__init__()
        self.index = 0
        self.setupUi(self)
        self.stack = QStackedLayout(self.widget) # self是嵌入到哪里
        self.one = ComformWord()
        self.two = MatchOneWidget()
        self.three = MethodWidget()
        self.stack.addWidget(self.one)
        self.stack.addWidget(self.two)
        self.stack.addWidget(self.three)
        # self.stack.setCurrentIndex(0)
        self.pushButton_next.clicked.connect(self.show_panel)
        self.pushButton_preview.clicked.connect(self.show_panel)
        self.pushButton_close.clicked.connect(self.close)

    def show_panel(self):
        """根据按钮切换显示页面"""
        plus_dict = {
            'pushButton_next': 1,
            'pushButton_preview': -1,
        }
        button_name = self.sender().objectName()
        self.plus = plus_dict[button_name]
        # if self.index >= -self.plus and self.index <= 1 - self.plus:
        if self.check_status():
            self.index += self.plus
            self.stack.setCurrentIndex(self.index)

    def check_status(self):
        """翻页条件检验"""
        if self.index == 0 and self.plus == 1:
            if os.path.exists(self.one.label_excel.text()) and os.path.exists(
                    self.one.label_word.text()):
                self.two.input_path(self.one.label_excel.text(),
                                    self.one.label_word.text())
                self.path_excel = self.one.label_excel.text()
                self.path_word = self.one.label_word.text()
                return True
            else:
                QMessageBox.information(self, '提示', '请填写正确模版路径',
                                        QMessageBox.Yes)
                return False
        elif self.index == 0 and self.plus == -1:
            return False
        elif self.index == 1 and self.plus == -1:
            return True
        elif self.index == 1 and self.plus == 1:
            self.three.wb = self.two.wb
            self.wb = self.two.wb
            self.nontable = self.two.get_relation()
            self.three.relation = self.two.relation
            self.relation = self.two.relation
            self.three.init()
            return True
        elif self.index == 2 and self.plus == -1:
            return True
        elif self.index == 2 and self.plus == 1:
            self.method, self.blank, self.group_code = self.three.get_info()
            self.out_dir = QFileDialog.getExistingDirectory(
                self, '生成询证函文件保存目录', '../')
            if self.out_dir:
                self.is_id, self.prefix, self.prefix_num = self.three.get_prefix(
                )
                self.generate_mail()
        else:
            return False

    def generate_mail_test(self):
        dfs = []
        for sheet in self.wb.sheetnames:
            if sheet != '表格0':
                data = self.wb[sheet].values
                columns = next(data)
                df = pd.DataFrame(data, columns=columns)
                df = df.rename(columns={df.columns[0]: 'ID'})
                df = df.drop(df.index[0])
                df = df.set_index(df.columns[0])
                if not df.empty:
                    dfs.append(df)
        for df in dfs:
            print(df)

    def generate_doc(self, df, num):
        """生成一封函证"""
        df = df.reset_index(drop=True)
        document = Document(self.path_word)
        i = 0
        for paragraph in document.paragraphs:
            while re.findall(r'\{.+?\}', paragraph.text) and i < 20:
                i += 1
                flag = '0.' + str(i)
                if flag in self.nontable.keys():
                    code = self.nontable[flag]
                else:
                    code = None
                if code != 'id': # 判断是否是聚合字段
                    if code in df.columns:
                        repl = df.loc[0, code]
                    elif code in self.df0.columns and not self.df0.empty:
                        repl = self.df0.loc[0, code]
                    else:
                        continue
                elif code == None:
                    continue
                else:
                    if self.prefix:
                        self.prefix = str(self.prefix)
                    else:
                        self.prefix = ''
                    if self.prefix_num in [0, 1]:
                        repl = self.prefix + str(num)
                    else:
                        if self.prefix_num:
                            repl = self.prefix + str(num).zfill(
                                self.prefix_num)
                        else:
                            repl = ''
                paragraph.text = re.sub(r'\{.*?\}',
                                        repl,
                                        paragraph.text,
                                        count=1)
        i = 0
        style = {"sz": 1, 'color': '000000', 'val': 'single'} # 设置斜线样式
        for table in document.tables:
            i += 1
            flag = str(i) + '.' + '1'
            rows = len(table.rows)
            columns = len(table.columns)
            if flag in df.columns:
                sub_df = df[df[flag] != '']
                sub_df = sub_df.reset_index(drop=True)
                data_num = len(sub_df)
                while len(table.rows) - 1 < data_num:
                    table.add_row()
                for k in range(0, len(table.rows)):
                    for j in range(0, len(table.columns)):
                        if k + 1 <= data_num:
                            key = str(i) + '.' + str(j + 1)
                            if key in df.columns:
                                text = sub_df.loc[k, key]
                                if isinstance(text, (int, float)):
                                    table.cell(k + 1,
                                               j).text = "{:,}".format(text)
                                else:
                                    table.cell(k + 1, j).text = str(text)
            for k in range(0, len(table.rows)):
                for j in range(0, len(table.columns)):
                    if '表格' + str(k + 1) in self.wb.sheetnames:
                        if table.cell(k, j).text == '':
                            if self.blank == '无':
                                table.cell(k, j).text = '无'
                            elif self.blank == '左斜线':
                                set_cell_border(table.cell(k, j),
                                                tl2br=style)
                            elif self.blank == '右斜线':
                                set_cell_border(table.cell(k, j),
                                                tr2bl=style)
            document.tables[i - 1] = table
        return document

    def insert_row(self, table, ix):
        """指定位置插入行"""
        tbl = table._tbl
        successor = tbl.tr_lst[ix]
        tr = tbl._new_tr()
        for gridCol in tbl.tblGrid.gridCol_lst:
            tc = tr.add_tc()
            tc.width = gridCol.w
        successor.addprevious(tr)
        return table.rows[ix]

    def format_row(self, table):
        """保持添加的最后一行格式和前面一致"""
        columns = len(table.columns)
        for i in range(columns):
            table.cell(-1, i).paragraphs[0].style = table.cell(
                -2, i).paragraphs[0].style

    def merge_docs(self, docs):
        """合并word文档"""
        merged_document = Composer(Document())
        for index, sub_doc in enumerate(docs):
            if index < len(docs) - 1:
                sub_doc.add_page_break()
            merged_document.append(sub_doc)
            # for element in sub_doc.element.body:
            # merged_document.element.body.append(element)
        file_name = os.path.join(self.out_dir, '生成文档.docx')
        merged_document.save(file_name)

    def generate_mail(self):
        """批量生成询证函"""
        dfs = []
        self.df0 = pd.DataFrame() # 创建空的dataframe 代表表格0（防止没有表格0时报错）
        for sheet in self.wb.sheetnames:
            # converters = {}
            # for column in self.relation[sheet].values():
            # converters[column] = str
            df = pd.read_excel(self.path_excel, sheet)
            df = df.fillna('')
            df = df.drop(df.index[0])
            if sheet != '表格0':
                df = df.rename(columns={df.columns[0]: 'ID'})
                df = df.set_index(df.columns[0])
                if not df.empty:
                    dfs.append(df)
            else:
                self.df0 = df.reset_index(drop=True) # 表格0

        df_all = dfs[0]
        if self.group_code in df_all.columns:
            for index, df in enumerate(dfs):
                if index != 0:
                    try:
                        df_all = pd.concat([df_all, df], axis=1, sort=False)
                    except:
                        QMessageBox.information(
                            self, '提示', self.wb.sheetnames[index] +
                            ':ID存在重复项，建议Excel模版里将表合并在一张表中,保证ID唯一',
                            QMessageBox.Ok)
                        return None
            df_all = df_all.fillna('')
            distinct_list = list(set(
                df_all[self.group_code].tolist())) # 获取聚合字段唯一列表
            docs = []
            for index, code in enumerate(distinct_list):
                df = df_all[df_all[self.group_code] == code]
                file_name = str(code) + '.docx'
                doc = self.generate_doc(df, index + 1)
                if self.method == '单独文件':
                    out_path = os.path.join(self.out_dir, file_name)
                    doc.save(out_path)
                else:
                    docs.append(doc)
            if docs:
                self.merge_docs(docs)
            QMessageBox.information(self, '提示', '生成完毕，请查看文件保存位置',
                                    QMessageBox.Ok)

        else:
            QMessageBox.information(self, '提示', '您所选择的聚合字段没有数据，请重新选择',
                                    QMessageBox.Ok)


class MethodWidget(QWidget, MethodUI):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.comboBox_num.addItems(['合并一个', '单独文件'])
        self.comboBox_blank.addItems(['空白', '无', '左斜线', '右斜线'])
        self.comboBox_table.currentIndexChanged.connect(self.combobox_changed)
        self.radioButton.toggled.connect(self.hide)
        self.lineEdit.setHidden(True)
        self.label_2.setHidden(True)
        self.label_3.setHidden(True)
        self.spinBox.setHidden(True)
        self.relation = None
        self.wb = None
        self.doc = None

    def get_prefix(self):
        if self.radioButton.isChecked():
            return True, self.lineEdit.text(), self.spinBox.value()
        else:
            return False, None, None

    def hide(self):
        """显示或隐藏编号控件"""
        if self.radioButton.isChecked():
            self.lineEdit.setHidden(False)
            self.label_2.setHidden(False)
            self.label_3.setHidden(False)
            self.spinBox.setHidden(False)
        else:
            self.lineEdit.setHidden(True)
            self.label_2.setHidden(True)
            self.label_3.setHidden(True)
            self.spinBox.setHidden(True)

    def init(self):
        if self.comboBox_table.currentText():
            pass
        else:
            self.comboBox_table.addItems(self.wb.sheetnames)
        ws = self.wb[self.wb.sheetnames[0]]
        self.comboBox_cell.addItems(self.get_row_list(ws, 2))

    def get_info(self):
        method = self.comboBox_num.currentText()
        blank = self.comboBox_blank.currentText()
        group_table = self.comboBox_table.currentText()
        group_cell = self.comboBox_cell.currentText()
        group_code = self.relation[group_table][group_cell]
        return method, blank, group_code

    def get_row_list(self, ws, row):
        return [ws.cell(row, i).value for i in range(1, ws.max_column + 1)]

    def combobox_changed(self):
        table_name = self.comboBox_table.currentText()
        ws = self.wb[table_name]
        self.comboBox_cell.clear()
        self.comboBox_cell.addItems(self.get_row_list(ws, 2))


class MatchOneWidget(QWidget, MatchOneUI):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.path_excel = ''
        self.path_word = ''
        self.wb = None

    def input_path(self, *arg):
        self.path_excel, self.path_word = arg
        self.load_excel()
        # self.tableWidget.cellChanged.connect(self.cellchanged)

    def cellchanged(self):
        for i in range(self.tableWidget.rowCount()):
            table_name = self.tableWidget.cellWidget(i, 0).currentText()
            if table_name != self.table_sheets[i]:
                self.table_sheets[i] = table_name
                ws = self.wb[table_name]
                comBoxCell = QComboBox()
                comBoxCell.addItems(self.get_row_list(ws, 2))
                self.tableWidget.setCellWidget(i, 1, comBoxCell)
                comBoxCell.currentIndexChanged.connect(self.update_code)
        self.update_code()

    def get_relation(self):
        relation = {}
        if self.info:
            for index, value in enumerate(self.info):
                relation[value] = self.tableWidget.item(index, 2).text()
        return relation

    def update_code(self):
        for i in range(self.tableWidget.rowCount()):
            if self.tableWidget.cellWidget(
                    i, 0) and self.tableWidget.cellWidget(i, 1):
                table_name = self.tableWidget.cellWidget(i, 0).currentText()
                name = self.tableWidget.cellWidget(i, 1).currentText()
                code = self.relation[table_name][name]
                item = QTableWidgetItem(code)
                self.tableWidget.setItem(i, 2, item)

    def get_row_list(self, ws, row):
        return [ws.cell(row, i).value for i in range(1, ws.max_column + 1)]

    def load_excel(self):
        self.wb = load_workbook(self.path_excel)
        self.relation = {}
        for sheet_name in self.wb.sheetnames:
            ws = self.wb[sheet_name]
            names = self.get_row_list(ws, 2)
            codes = self.get_row_list(ws, 1)
            tmp = {}
            for index, name in enumerate(names):
                tmp[name] = codes[index]
            self.relation[sheet_name] = tmp
        if '表格0' in self.wb.sheetnames:
            ws = self.wb['表格0']
            row1 = self.get_row_list(ws, 1)
            row2 = self.get_row_list(ws, 2)
            self.info = row1
            self.table_sheets = [
                self.wb.sheetnames[0] for i in range(len(row1))
            ]
            self.tableWidget.setColumnCount(3)
            self.tableWidget.setRowCount(len(row1))
            self.tableWidget.setHorizontalHeaderLabels(['取值表', '取值字段', '编号'])
            self.tableWidget.setVerticalHeaderLabels(row2)
            self.tableWidget.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
            self.tableWidget.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)

            for i in range(len(row2)):
                comBoxTable = QComboBox()
                comBoxTable.addItems(self.wb.sheetnames)
                comBoxCell = QComboBox()
                comBoxCell.addItems(
                    self.get_row_list(self.wb[self.wb.sheetnames[0]], 2))
                comBoxCell.currentIndexChanged.connect(self.update_code)
                comBoxCell.setCurrentIndex(i)
                comBoxTable.currentIndexChanged.connect(self.cellchanged)
                self.tableWidget.setCellWidget(i, 1, comBoxCell)
                self.tableWidget.setCellWidget(i, 0, comBoxTable)
                code = self.relation[self.wb.sheetnames[0]][
                    comBoxCell.currentText()]
                item = QTableWidgetItem(code)
                self.tableWidget.setItem(i, 2, item)
                self.tableWidget.update()
        else:
            self.info = None

        # data = self.wb['表格0'].values
        # columns = next(data)
        # df = pd.DataFrame(data, columns=columns)
        # print(df)


class ComformWord(QDialog, ComformationUI):
    def __init__(self, parent=None):
        super(ComformWord, self).__init__(parent)
        self.setupUi(self)
        self.setWindowTitle("询证函批量制作")
        self.pushButton_excel.clicked.connect(self.onSelectExcel)
        self.pushButton_word.clicked.connect(self.onSelectWord)

    def get_path(self):
        """输出excel,word路径"""
        return self.label_excel.text(), self.label_word.text()

    def onSelectExcel(self):
        reply = QMessageBox.information(self, "提问", "是否已准备好Excel数据模版",
                                        QMessageBox.Yes | QMessageBox.No,
                                        QMessageBox.No)
        if reply == 16384:
            path = QFileDialog.getOpenFileName(self, "打开Excel文件", './',
                                               'xlsx(*.xlsx)')[0]
            if path:
                self.label_excel.setText(path)
        else:
            word_file_name = self.label_word.text()
            if os.path.exists(word_file_name) and re.findall(
                    'docx?$', word_file_name):
                template_ui = CreateTemplate()
                template_ui.get_word_path(word_file_name)
                template_ui.get_table_list()
                template_ui.exec_()
            else:
                QMessageBox.information(self, "提示",
                                        "请先选择word模版路径，系统将根据word模版制作Excel模版",
                                        QMessageBox.Yes, QMessageBox.Yes)

    def onSelectWord(self):
        path = QFileDialog.getOpenFileName(self, "打开word文件", './',
                                           'docx(*.docx)')[0]
        if path:
            self.label_word.setText(path)


class CreateTemplate(QDialog, ComformationTemplateUI):
    def __init__(self):
        super(CreateTemplate, self).__init__()
        self.setupUi(self)
        self.listView_in.clicked.connect(self.show_table)
        self.listView_out.clicked.connect(self.show_table)
        self.pushButton_add.clicked.connect(self.add_one)
        self.pushButton_add_all.clicked.connect(self.add_all)
        self.pushButton_minus.clicked.connect(self.minus_one)
        self.pushButton_minus_all.clicked.connect(self.minus_all)
        self.pushButton_close.clicked.connect(self.close)
        self.pushButton_next.clicked.connect(self.generate)
        self.list_model_out = QStringListModel()
        self.list_model_out.setStringList([])
        self.tables_in = []
        self.tables_out = []

    def get_word_path(self, path):
        self.path = path
        self.document = Document(self.path)
        self.tables_in = [
            '表格' + str(i + 1) for i in range(len(self.document.tables))
        ]

    def sort_table_name(self, tables):
        table_numbers = []
        table_relation = {}
        sorted_table = []
        for table in tables:
            num = table.split('表格')[1]
            table_numbers.append(int(num))
            table_relation[int(num)] = table
        table_numbers.sort()
        for table_number in table_numbers:
            sorted_table.append(table_relation[table_number])
        return sorted_table

    def update_listview(self):
        """刷新两个列表"""
        self.list_model_out.setStringList(self.tables_out)
        self.listView_out.setModel(self.list_model_out)
        self.list_model_in.setStringList(self.tables_in)
        self.listView_in.setModel(self.list_model_in)

    def add_one(self):
        """增加需要输出的表格"""
        row = self.list_model_out.rowCount() # 右列的目前的最大值
        index_out = self.listView_out.currentIndex()
        index_in = self.listView_in.currentIndex()
        if index_in.data() not in self.tables_out and index_in.data():
            self.tables_out.append(index_in.data())
            self.tables_out = self.sort_table_name(self.tables_out)
            self.tables_in.remove(index_in.data())
            self.update_listview()

    def add_all(self):
        """选择所有为需要输出的表格"""
        self.tables_out = self.tables_in + self.tables_out
        self.tables_out = self.sort_table_name(self.tables_out)
        self.tables_in = []
        self.update_listview()

    def minus_one(self):
        """去除需要输出的表格"""
        index_out = self.listView_out.currentIndex()
        index_in = self.listView_in.currentIndex()
        if index_out.data() not in self.tables_in and index_out.data():
            self.tables_in.append(index_out.data())
            self.tables_in = self.sort_table_name(self.tables_in)
            self.tables_out.remove(index_out.data())
            self.update_listview()

    def minus_all(self):
        """去除所有为需要输出的表格"""
        self.tables_in = self.tables_in + self.tables_out
        self.tables_in = self.sort_table_name(self.tables_in)
        self.tables_out = []
        self.update_listview()

    def get_table_list(self):
        self.list_model_in = QStringListModel()
        self.list_model_in.setStringList(self.tables_in)
        self.listView_in.setModel(self.list_model_in)

    def show_table(self, index):
        """展示选中的表格具体表头内容"""
        table_name = index.data()
        k = int(table_name.split('表格')[1]) - 1
        tables = self.document.tables
        column_list = []
        for i in range(len(tables[k].columns)):
            cell_text = tables[k].cell(0, i).text
            column_list.append(cell_text)
        self.tableWidget.setColumnCount(len(column_list))
        self.tableWidget.setRowCount(2)
        self.tableWidget.setHorizontalHeaderLabels(column_list)
        self.tableWidget.setVerticalHeaderLabels(['1', '2'])
        self.tableWidget.setHorizontalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        self.tableWidget.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)

    def word_info(self):
        """提取word中非表格关键字段"""
        paragraphs = self.document.paragraphs
        infos = []
        for paragraph in paragraphs:
            result = re.findall(r'(?<={).*?(?=})', paragraph.text)
            if result:
                infos += result
        return infos

    def generate(self):
        """生成Excel表格"""
        if self.tables_out:
            excel_path = QFileDialog.getSaveFileName(self, '生成模版保存位置',
                                                     './数据模版.xlsx',
                                                     'xlsx(*.xlsx)')[0]
            if excel_path:
                wb = Workbook()
                for table_name in self.tables_out:
                    index = int(table_name.split('表格')[1]) - 1
                    ws = wb.create_sheet(table_name, -1)
                    for i in range(
                            len(self.document.tables[index].columns) + 1):
                        column = i + 1
                        ws.cell(1,
                                column).value = str(index + 1) + '.' + str(i)
                        if i == 0:
                            ws.cell(2, column).value = 'ID'
                        ws.cell(2, column +
                                1).value = self.document.tables[index].cell(
                                    0, i).text
                infos = self.word_info()
                ws = wb.create_sheet('表格0', 0)
                for index, name in enumerate(infos):
                    ws.cell(1, index + 1).value = '0.' + str(index + 1)
                    ws.cell(2, index + 1).value = name
                del wb['Sheet']
                wb.save(excel_path)
                wb.close()
        else:
            QMessageBox.information(self, '提示', '请先添加需要的表格来生成模版',
                                    QMessageBox.Yes)


def set_cell_border(cell: _Cell, **kwargs):
    """
    Set cell`s border
    Usage:
    set_cell_border(
        cell,
        top={"sz": 12, "val": "single", "color": "#FF0000", "space": "0"},
        bottom={"sz": 12, "color": "#00FF00", "val": "single"},
        start={"sz": 24, "val": "dashed", "shadow": "true"},
        end={"sz": 12, "val": "dashed"},
        tr2bl={}
        tl2br={}
    )
    """
    tc = cell._tc
    tcPr = tc.get_or_add_tcPr()

    # check for tag existnace, if none found, then create one
    tcBorders = tcPr.first_child_found_in("w:tcBorders")
    if tcBorders is None:
        tcBorders = OxmlElement('w:tcBorders')
        tcPr.append(tcBorders)

    # list over all available tags
    for edge in ('start', 'top', 'end', 'bottom', 'insideH', 'insideV',
                 'tr2bl', 'tl2br'):
        edge_data = kwargs.get(edge)
        if edge_data:
            tag = 'w:{}'.format(edge)

            # check for tag existnace, if none found, then create one
            element = tcBorders.find(qn(tag))
            if element is None:
                element = OxmlElement(tag)
                tcBorders.append(element)

            # looks like order of attributes is important
            for key in ["sz", "val", "color", "space", "shadow"]:
                if key in edge_data:
                    element.set(qn('w:{}'.format(key)), str(edge_data[key]))


if __name__ == "__main__":
    app = QApplication(sys.argv)
    dialog_ui = ComformWord()
    dialog_ui.show()
    # dialog_ui = CreateTemplate()
    # dialog_ui.show()
    sys.exit(app.exec_())