import os
import time
import fitz
import signal
import base64
import asyncio
import aiohttp
from openpyxl import Workbook
from django.core.management import BaseCommand

from common.mixins import LoggerMixin
from common.tools.file_tools import write_zip_file
from common.tools.pdf_to_img import PDFHandler
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}.xlsx'.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)

    # 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')
                self.cronjob_log.info('{0} [pdf to img start] [business_type={1}] [doc_id={2}]'.format(
                    self.log_base, business_type, doc.id))
                pdf_handler = PDFHandler(pdf_path, img_save_path)
                pdf_handler.extract_image()
                self.cronjob_log.info('{0} [pdf to img end] [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 pdf_handler.img_path_list]
                loop.run_until_complete(asyncio.wait(tasks))
                # loop.close()

                # 整合excel文件
                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))