import os
import time
import fitz
import signal
import base64
import asyncio
import aiohttp
import locale
from PIL import Image
from io import BytesIO
from openpyxl import Workbook
from openpyxl.styles import numbers
from openpyxl.utils import get_column_letter

from django.core.management import BaseCommand
from common.mixins import LoggerMixin
from common.tools.file_tools import write_zip_file
from apps.doc.models import DocStatus, HILDoc, AFCDoc
from apps.doc import consts
from settings import conf
from apps.doc.edms import EDMS, rh


class Command(BaseCommand, LoggerMixin):

    def __init__(self):
        super().__init__()
        self.log_base = '[doc ocr process]'
        # 处理文件开关
        self.switch = True
        # 数据目录
        self.data_dir = conf.DATA_DIR
        # pdf页面转图片
        self.zoom_x = 2.0
        self.zoom_y = 2.0
        self.trans = fitz.Matrix(self.zoom_x, self.zoom_y).preRotate(0)  # zoom factor 2 in each dimension
        # ocr相关
        self.ocr_url = conf.OCR_URL
        self.ocr_header = {
            'X-Auth-Token': conf.OCR_TOKEN,
            'Content-Type': 'application/json'
        }
        # EDMS web_service_api
        self.edms = EDMS(conf.EDMS_USER, conf.EDMS_PWD)
        # 优雅退出信号：15
        signal.signal(signal.SIGTERM, self.signal_handler)

    def signal_handler(self, sig, frame):
        self.switch = False  # 停止处理文件

    def get_doc_info(self):
        task_str, is_priority = rh.dequeue()
        if task_str is None:
            self.cronjob_log.info('{0} [get_doc_info] [queue empty]'.format(self.log_base))
            return None, None

        business_type, doc_id_str = task_str.split(consts.SPLIT_STR)
        doc_id = int(doc_id_str)
        doc_class = HILDoc if business_type == consts.HIL_PREFIX else AFCDoc
        # doc_info = doc_class.objects.filter(id=doc_id, status=DocStatus.INIT.value).values(
        #     'id', 'metadata_version_id', 'application_id', 'document_name', 'document_scheme').first()
        doc = doc_class.objects.filter(id=doc_id).first()
        if doc is None:
            self.cronjob_log.warn('{0} [get_doc_info] [doc not exist] [task_str={1}] [is_priority={2}]'.format(
                self.log_base, task_str, is_priority))
            return None, None
        elif doc.status != DocStatus.INIT.value:
            self.cronjob_log.warn('{0} [get_doc_info] [doc status error] [task_str={1}] [is_priority={2}] '
                                  '[doc_status={3}]'.format(self.log_base, task_str, is_priority, doc.status))
            return None, None
        doc.status = DocStatus.PROCESSING.value
        doc.save()
        self.cronjob_log.info('{0} [get_doc_info] [success] [task_str={1}] [is_priority={2}]'.format(
            self.log_base, task_str, is_priority))
        return doc, business_type

    def pdf_download(self, doc, business_type):
        if doc is None:
            return None, None, None
        # TODO EDMS下载pdf
        doc_data_path = os.path.join(self.data_dir, business_type, str(doc.id))
        os.makedirs(doc_data_path, exist_ok=True)
        pdf_path = os.path.join(doc_data_path, '{0}.pdf'.format(doc.id))
        if not doc.application_id.startswith(consts.FIXED_APPLICATION_ID_PREFIX):
            self.edms.download(pdf_path, doc.metadata_version_id)

        excel_path = os.path.join(doc_data_path, '{0}.xls'.format(doc.id))
        self.cronjob_log.info('{0} [pdf download success] [business_type={1}] [doc_id={2}] [pdf_path={3}]'.format(
            self.log_base, business_type, doc.id, pdf_path))
        return doc_data_path, excel_path, pdf_path

    @staticmethod
    def append_sheet(wb, sheets_list, img_name):
        for i, sheet in enumerate(sheets_list):
            ws = wb.create_sheet('{0}_{1}'.format(img_name, i))
            cells = sheet.get('cells')
            for cell in cells:
                c1 = cell.get('start_column')
                # c2 = cell.get('end_column')
                r1 = cell.get('start_row')
                # r2 = cell.get('end_row')
                label = cell.get('words')
                ws.cell(row=r1+1, column=c1+1, value=label)

    @staticmethod
    def get_ocr_json(img_path):
        with open(img_path, "rb") as f:
            base64_data = base64.b64encode(f.read())
        return {'imgBase64': base64_data.decode('utf-8')}

    async def fetch_ocr_result(self, img_path):
        async with aiohttp.ClientSession(
                headers=self.ocr_header, connector=aiohttp.TCPConnector(ssl=False)
        ) as session:
            json_data = self.get_ocr_json(img_path)
            async with session.post(self.ocr_url, json=json_data) as response:
                return await response.json()

    async def img_ocr_excel(self, wb, img_path):
        res = await self.fetch_ocr_result(img_path)
        self.cronjob_log.info('{0} [fetch ocr result success] [img={1}] [res={2}]'.format(self.log_base, img_path, res))
        sheets_list = res.get('result').get('res')
        img_name = os.path.basename(img_path)
        self.append_sheet(wb, sheets_list, img_name)

    def proof(self, ws):
        # 找到金额、余额列
        amount_col = overage_col = None
        for i in ws[1]:
            if i.value in consts.AMOUNT_COL_TITLE_SET:
                amount_col = i.column
                amount_col_letter = get_column_letter(amount_col)
            elif i.value in consts.OVERAGE_COL_TITLE_SET:
                overage_col = i.column
                overage_col_letter = get_column_letter(overage_col)
        if amount_col is None or overage_col is None:
            return
        # 文本转数值
        for col_tuple in ws.iter_cols(min_row=2, min_col=amount_col, max_col=overage_col):
            for c in col_tuple:
                try:
                    c.value = locale.atof(c.value)
                    c.number_format = numbers.FORMAT_NUMBER_00
                except Exception:
                    continue
        # 增加核对结果列
        proof_col_letter = get_column_letter(ws.max_column + 1)
        for c in ws[proof_col_letter]:
            if c.row == 1:
                c.value = consts.PROOF_COL_TITLE
            elif c.row == 2:
                continue
            else:
                c.value = '=IF({3}{0}=SUM({2}{0},{3}{1}), "{4}", "{5}")'.format(
                    c.row, c.row - 1, amount_col_letter, overage_col_letter, *consts.PROOF_RES)

    def wb_process(self, wb, excel_path):
        locale.setlocale(locale.LC_NUMERIC, 'en_US.UTF-8')
        for ws in wb.worksheets:
            if ws.title == 'Sheet':
                ws.title = consts.META_SHEET_TITLE
            else:
                self.proof(ws)
        wb.save(excel_path)  # TODO no sheet (res always [])

    @staticmethod
    def getimage(pix):
        if pix.colorspace.n != 4:
            return pix
        tpix = fitz.Pixmap(fitz.csRGB, pix)
        return tpix

    def recoverpix(self, doc, item):
        x = item[0]  # xref of PDF image
        s = item[1]  # xref of its /SMask
        is_rgb = True if item[5] == 'DeviceRGB' else False

        # RGB
        if is_rgb:
            if s == 0:
                return doc.extractImage(x)
            # we need to reconstruct the alpha channel with the smask
            pix1 = fitz.Pixmap(doc, x)
            pix2 = fitz.Pixmap(doc, s)  # create pixmap of the /SMask entry

            # sanity check
            if not (pix1.irect == pix2.irect and pix1.alpha == pix2.alpha == 0 and pix2.n == 1):
                pix2 = None
                return self.getimage(pix1)

            pix = fitz.Pixmap(pix1)  # copy of pix1, alpha channel added
            pix.setAlpha(pix2.samples)  # treat pix2.samples as alpha value
            pix1 = pix2 = None  # free temp pixmaps
            return self.getimage(pix)

        # CMYK
        pix1 = fitz.Pixmap(doc, x)
        pix = fitz.Pixmap(pix1)  # copy of pix1, alpha channel added

        if s != 0:
            pix2 = fitz.Pixmap(doc, s)  # create pixmap of the /SMask entry

            # sanity check
            if not (pix1.irect == pix2.irect and pix1.alpha == pix2.alpha == 0 and pix2.n == 1):
                pix2 = None
                return self.getimage(pix1)

            pix.setAlpha(pix2.samples)  # treat pix2.samples as alpha value

        pix1 = pix2 = None  # free temp pixmaps

        pix = fitz.Pixmap(fitz.csRGB, pix)  # GRAY/CMYK to RGB
        return self.getimage(pix)

    @staticmethod
    def get_img_data(pix):
        if type(pix) is dict:  # we got a raw image
            ext = pix["ext"]
            img_data = pix["image"]
        else:  # we got a pixmap
            ext = 'png'
            img_data = pix.getPNGData()
        return ext, img_data

    @staticmethod
    def split_il(il):
        img_il_list = []
        start = 0
        length = len(il)
        for i in range(length):
            if i == start:
                if i == length - 1:
                    img_il_list.append(il[start: length])
                continue
            elif i == length - 1:
                img_il_list.append(il[start: length])
                continue
            if il[i][2] != il[i - 1][2]:
                img_il_list.append(il[start: i])
                start = i
            elif il[i][3] != il[i - 1][3]:
                img_il_list.append(il[start: i + 1])
                start = i + 1
        return img_il_list

    # TODO 细化文件状态，不同异常状态采取不同的处理
    # TODO 调用接口重试
    def handle(self, *args, **kwargs):
        sleep_second = int(conf.SLEEP_SECOND)
        max_sleep_second = int(conf.MAX_SLEEP_SECOND)

        while self.switch:
            # 1. 从队列获取文件信息
            doc, business_type = self.get_doc_info()

            try:
                # 2. 从EDMS获取PDF文件
                doc_data_path, excel_path, pdf_path = self.pdf_download(doc, business_type)

                # 队列为空时的处理
                if pdf_path is None:
                    time.sleep(sleep_second)
                    sleep_second = min(max_sleep_second, sleep_second+5)
                    continue

                sleep_second = int(conf.SLEEP_SECOND)

                # 3.PDF文件提取图片
                img_save_path = os.path.join(doc_data_path, 'img')
                os.makedirs(img_save_path, exist_ok=True)
                img_path_list = []
                with fitz.Document(pdf_path) as pdf:
                    self.cronjob_log.info('{0} [pdf_path={1}] [metadata={2}]'.format(
                        self.log_base, pdf_path, pdf.metadata))
                    # xref_list = []  # TODO 图片去重 特殊pdf：如电子发票
                    for pno in range(pdf.pageCount):
                        il = pdf.getPageImageList(pno)
                        il.sort(key=lambda x: x[0])
                        img_il_list = self.split_il(il)
                        del il

                        if len(img_il_list) > 3:  # 单页无规律小图过多时，使用页面转图片
                            page = pdf.loadPage(pno)
                            pm = page.getPixmap(matrix=self.trans, alpha=False)
                            save_path = os.path.join(img_save_path, 'page_{0}_img_0.png'.format(page.number))
                            pm.writePNG(save_path)
                            img_path_list.append(save_path)
                            self.cronjob_log.info('{0} [page to img success] [pdf_path={1}] [page={2}]'.format(
                                self.log_base, pdf_path, page.number))
                        else:  # 提取图片
                            for img_index, img_il in enumerate(img_il_list):
                                if len(img_il) == 1:  # 当只有一张图片时， 简化处理
                                    pix = self.recoverpix(pdf, img_il[0])
                                    ext, img_data = self.get_img_data(pix)
                                    save_path = os.path.join(img_save_path, 'page_{0}_img_{1}.{2}'.format(
                                        pno, img_index, ext))
                                    with open(save_path, "wb") as f:
                                        f.write(img_data)
                                    img_path_list.append(save_path)
                                    self.cronjob_log.info(
                                        '{0} [extract img success] [pdf_path={1}] [page={2}] [img_index={3}]'.format(
                                            self.log_base, pdf_path, pno, img_index))
                                else:  # 多张图片，竖向拼接
                                    height_sum = 0
                                    im_list = []
                                    width = img_il[0][2]
                                    for img in img_il:
                                        # xref = img[0]
                                        # if xref in xref_list:
                                        #     continue
                                        height = img[3]
                                        pix = self.recoverpix(pdf, img)
                                        ext, img_data = self.get_img_data(pix)

                                        # xref_list.append(xref)

                                        im = Image.open(BytesIO(img_data))
                                        im_list.append((height, im, ext))
                                        height_sum += height

                                    save_path = os.path.join(img_save_path, 'page_{0}_img_{1}.{2}'.format(
                                        pno, img_index, im_list[0][2]))
                                    res = Image.new(im_list[0][1].mode, (width, height_sum))
                                    h_now = 0
                                    for h, m, _ in im_list:
                                        res.paste(m, box=(0, h_now))
                                        h_now += h
                                    res.save(save_path)
                                    img_path_list.append(save_path)
                                    self.cronjob_log.info(
                                        '{0} [extract img success] [pdf_path={1}] [page={2}] [img_index={3}]'.format(
                                            self.log_base, pdf_path, pno, img_index))
                    self.cronjob_log.info('{0} [pdf to img success] [business_type={1}] [doc_id={2}]'.format(
                        self.log_base, business_type, doc.id))

                write_zip_file(img_save_path, os.path.join(doc_data_path, '{0}_img.zip'.format(doc.id)))

                # 4.图片调用算法判断是否为银行流水, 图片调用算法OCR为excel文件
                wb = Workbook()
                loop = asyncio.get_event_loop()
                tasks = [self.img_ocr_excel(wb, img_path) for img_path in img_path_list]
                loop.run_until_complete(asyncio.wait(tasks))
                # loop.close()

                # 整合excel文件
                # self.wb_process(wb, excel_path)
                wb.save(excel_path)
            except Exception as e:
                doc.status = DocStatus.PROCESS_FAILED.value
                doc.save()
                self.cronjob_log.error('{0} [process failed] [business_type={1}] [doc_id={2}] [err={3}]'.format(
                    self.log_base, business_type, doc.id, e))
            else:
                try:
                    # 5.上传至EDMS
                    self.edms.upload(excel_path, doc, business_type)
                except Exception as e:
                    doc.status = DocStatus.UPLOAD_FAILED.value
                    doc.save()
                    self.cronjob_log.error('{0} [upload failed] [business_type={1}] [doc_id={2}] [err={3}]'.format(
                        self.log_base, business_type, doc.id, e))
                else:
                    doc.status = DocStatus.COMPLETE.value
                    doc.save()
                    self.cronjob_log.info('{0} [doc process complete] [business_type={1}] [doc_id={2}]'.format(
                        self.log_base, business_type, doc.id))

        self.cronjob_log.info('{0} [stop safely]'.format(self.log_base))