import os
import time
import json
import shutil
import base64
import signal
import requests
import traceback
from datetime import datetime
from django.core.management import BaseCommand
from multiprocessing import Process
from openpyxl import load_workbook, Workbook

from settings import conf
from common.mixins import LoggerMixin
from common.tools.pdf_to_img import PDFHandler, PDFBuild
from apps.doc import consts
from apps.doc.exceptions import OCR1Exception, OCR2Exception


class Command(BaseCommand, LoggerMixin):

    def __init__(self):
        super().__init__()
        self.log_base = '[folder f3 process]'
        # input folder
        self.input_dir = conf.F3_DIR
        # 处理文件开关
        self.switch = True
        # 睡眠时间
        self.sleep_time = float(conf.SLEEP_SECOND_FOLDER)
        # 输出结果
        self.wb_name = 'result.xlsx'

        self.field_map = {
            # sheet_name, key_field, side_field_order, src_field_order
            consts.IC_CLASSIFY: (consts.IC_CN_NAME, '有效期限', consts.IC_FIELD_ORDER_3, consts.IC_FIELD_ORDER_2),
        }

        self.pdf_name_map = {
            consts.IC_CLASSIFY: consts.IC_CN_NAME,
            consts.BL_CLASSIFY: consts.BL_CN_NAME,
        }

        # ocr url
        self.ocr_url = conf.OCR_URL_FOLDER
        self.ocr_url_2 = conf.OCR2_URL_FOLDER

        # 优雅退出信号：15
        signal.signal(signal.SIGTERM, self.signal_handler)

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

    @staticmethod
    def license1_process(ocr_data, ocr_res, classify):
        # 类别：'0'身份证， '1'居住证
        license_data = ocr_data.get('data')
        if not license_data:
            return
        if isinstance(license_data, dict):
            license_data.pop('base64_img', '')
        if classify == consts.IC_CLASSIFY:
            id_card_dict = {}
            card_type = license_data.get('type', '')
            is_ic = card_type.startswith('身份证')
            is_info_side = card_type.endswith('信息面')
            id_card_dict['类别'] = '0' if is_ic else '1'
            if is_ic:
                field_map = consts.IC_MAP_0 if is_info_side else consts.IC_MAP_1
            else:
                field_map = consts.RP_MAP_0 if is_info_side else consts.RP_MAP_1
            for write_field, search_field in field_map:
                id_card_dict[write_field] = license_data.get('words_result', {}).get(search_field, {}).get('words', '')
            if not is_info_side:
                start_time = license_data.get('words_result', {}).get('签发日期', {}).get('words', '')
                end_time = license_data.get('words_result', {}).get('失效日期', {}).get('words', '')
                id_card_dict['有效期限'] = '{0}-{1}'.format(start_time, end_time)
            # for id_card_dict in license_data:
            #     try:
            #         id_card_dict.pop('base64_img')
            #     except Exception as e:
            #         continue
            ocr_res.setdefault(classify, []).append(id_card_dict)

    def license2_process(self, ocr_data, ocr_res, classify, img_path):
        if classify != consts.BL_CLASSIFY:
            return

        pid, _, _, _, _, _ = consts.LICENSE_CLASSIFY_MAPPING.get(classify)
        file_data = ocr_data.get('section_img')
        if file_data is None:
            with open(img_path, 'rb') as f:
                base64_data = base64.b64encode(f.read())
                # 获取解码后的base64值
                file_data = base64_data.decode()
        json_data_2 = {
            "pid": str(pid),
            "filedata": file_data
        }

        for times in range(consts.RETRY_TIMES):
            try:
                start_time = time.time()
                ocr_2_response = requests.post(self.ocr_url_2, data=json_data_2)
                if ocr_2_response.status_code != 200:
                    raise OCR2Exception('ocr_2 status code: {0}'.format(ocr_2_response.status_code))
            except Exception as e:
                self.folder_log.warn(
                    '{0} [ocr_2 failed] [times={1}] [img_path={2}] [error={3}]'.format(
                        self.log_base, times, img_path, traceback.format_exc()))
            else:
                ocr_res_2 = json.loads(ocr_2_response.text)
                end_time = time.time()
                speed_time = int(end_time - start_time)
                self.folder_log.info(
                    '{0} [ocr_2 success] [img={1}] [speed_time={2}]'.format(
                        self.log_base, img_path, speed_time))

                if ocr_res_2.get('ErrorCode') in consts.SUCCESS_CODE_SET:
                    if pid == consts.BC_PID:
                        ocr_res.append(ocr_res_2)
                    else:
                        # 营业执照等
                        for result_dict in ocr_res_2.get('ResultList', []):
                            res_dict = {}
                            for field_dict in result_dict.get('FieldList', []):
                                res_dict[field_dict.get('chn_key', '')] = field_dict.get('value', '')
                            ocr_res.append(res_dict)
                break

    def wb_process(self, ocr_res, output_dir):
        excel_path = os.path.join(output_dir, self.wb_name)

        try:
            if os.path.exists(excel_path):
                wb = load_workbook(excel_path)
            else:
                wb = Workbook()

            for c, res_list in ocr_res.items():
                sheet_name, key_field, side_field_order, src_field_order = self.field_map.get(c)
                if sheet_name in wb.sheetnames:
                    ws = wb.get_sheet_by_name(sheet_name)
                    ws.append((None,))
                else:
                    ws = wb.create_sheet(sheet_name)

                for res in res_list:
                    if key_field is not None and key_field in res:
                        field_order = side_field_order
                    else:
                        field_order = src_field_order
                    for search_field, write_field in field_order:
                        field_value = res.get(search_field, '')
                        if isinstance(field_value, list):
                            ws.append((write_field, *field_value))
                        else:
                            ws.append((write_field, field_value))
                    ws.append((None,))

            if 'Sheet' in wb.sheetnames and len(wb.sheetnames) > 1:
                wb.remove(wb.get_sheet_by_name('Sheet'))
            wb.save(excel_path)
        except Exception as e:
            self.folder_log.error('{0} [wb build error] [path={1}] [error={2}]'.format(
                self.log_base, excel_path, traceback.format_exc()))

    def ocr_process(self, img_path, all_ocr_res, img_res):
        if os.path.exists(img_path):
            with open(img_path, 'rb') as f:
                base64_data = base64.b64encode(f.read())
                # 获取解码后的base64值
                file_data = base64_data.decode()
            json_data = {
                "file": file_data,
                "channel": consts.AFC_PREFIX,
            }

            for times in range(consts.RETRY_TIMES):
                try:
                    start_time = time.time()
                    ocr_response = requests.post(self.ocr_url, json=json_data)
                    if ocr_response.status_code != 200:
                        raise OCR1Exception('{0} ocr status code: {1}'.format(self.log_base, ocr_response.status_code))
                except Exception as e:
                    self.folder_log.warn('{0} [ocr failed] [times={1}] [img_path={2}] [error={3}]'.format(
                        self.log_base, times, img_path, traceback.format_exc()))
                else:
                    ocr_res = ocr_response.json()
                    end_time = time.time()
                    speed_time = int(end_time - start_time)
                    self.folder_log.info('{0} [ocr success] [img={1}] [speed_time={2}]'.format(
                        self.log_base, img_path, speed_time))

                    if isinstance(ocr_res, dict):
                        if ocr_res.get('code') == 1:
                            data_list = ocr_res.get('data', [])
                            if isinstance(data_list, list):
                                for ocr_data in data_list:
                                    classify = ocr_data.get('classify')
                                    img_res.setdefault(classify, set()).add(img_path)
                                    if classify in consts.LICENSE_CLASSIFY_SET_1:
                                        self.license1_process(ocr_data, all_ocr_res, classify)
                                    elif classify in consts.LICENSE_CLASSIFY_SET_2:
                                        self.license2_process(ocr_data, all_ocr_res, classify, img_path)
                    break
            else:
                self.folder_log.warn('{0} [ocr failed] [img_path={1}]'.format(self.log_base, img_path))

    def img_process(self, img_res, output_dir):
        for classify, img_path_set in img_res.items():
            pdf_path = os.path.join(output_dir, '{0}.pdf'.format(self.pdf_name_map.get(classify, '其他')))
            pdf_build = PDFBuild(pdf_path)
            pdf_build.insert_img(img_path_set)

    def images_process(self, img_path_list, output_dir):
        ocr_res = dict()
        img_res = dict()
        for img_path in img_path_list:
            self.ocr_process(img_path, ocr_res, img_res)
        self.wb_process(ocr_res, output_dir)
        self.img_process(img_res, output_dir)

    def pdf_process(self, name, path, img_output_dir, output_dir):
        pdf_handler = PDFHandler(path, os.path.join(img_output_dir, name))

        try:
            self.folder_log.info('{0} [pdf to img start] [path={1}]'.format(self.log_base, path))
            pdf_handler.extract_image()
            self.folder_log.info('{0} [pdf to img end] [path={1}]'.format(self.log_base, path))
        except Exception as e:
            self.folder_log.error('{0} [pdf to img error] [path={1}] [error={2}]'.format(
                self.log_base, path, traceback.format_exc()))
            raise e
        else:
            self.images_process(pdf_handler.img_path_list, output_dir)

    def folder_process(self, unique_folder_name, folder_path, main_output_dir):
        output_dir = os.path.join(main_output_dir, unique_folder_name)

        img_output_dir = os.path.join(output_dir, 'image')
        failed_output_dir = os.path.join(output_dir, 'failed')
        os.makedirs(output_dir, exist_ok=True)
        os.makedirs(img_output_dir, exist_ok=True)
        os.makedirs(failed_output_dir, exist_ok=True)

        os_error_filename_set = set()

        list_dir = os.listdir(folder_path)
        if not list_dir and len(os_error_filename_set) == 0:
            self.folder_log.info('{0} [folder empty, completed] [path={1}]'.format(self.log_base, folder_path))
            return

        all_file_set = set(list_dir)
        true_file_set = all_file_set - os_error_filename_set
        if len(true_file_set) == 0 and len(os_error_filename_set) > 0:
            true_file_set.add(os_error_filename_set.pop())
        for name in true_file_set:
            time.sleep(5)
            path = os.path.join(folder_path, name)

            try:
                if not os.path.exists(path):
                    self.folder_log.info('{0} [path is not exists] [path={1}]'.format(self.log_base, path))
                    continue
                elif os.path.isfile(path):
                    self.folder_log.info('{0} [file start] [path={1}]'.format(self.log_base, path))
                    self.pdf_process(name, path, img_output_dir, output_dir)
                    self.folder_log.info('{0} [file end] [path={1}]'.format(self.log_base, path))
                else:
                    self.folder_log.info('{0} [path is dir] [path={1}]'.format(self.log_base, path))
                    failed_path = os.path.join(failed_output_dir, name)
                    shutil.copyfile(path, failed_path)
            except OSError:
                os_error_filename_set.add(name)
                self.folder_log.error('{0} [os error] [path={1}] [error={2}]'.format(
                    self.log_base, path, traceback.format_exc()))
            except Exception as e:
                try:
                    self.folder_log.error('{0} [file error] [path={1}] [error={2}]'.format(self.log_base, path,
                                                                                           traceback.format_exc()))
                    failed_path = os.path.join(failed_output_dir, name)
                    shutil.copyfile(path, failed_path)
                except Exception as e:
                    os_error_filename_set.add(name)
                    self.folder_log.error('{0} [file move error] [path={1}] [error={2}]'.format(
                        self.log_base, path, traceback.format_exc()))

    def main_folder_process(self, input_dir):
        while not os.path.isdir(input_dir):
            self.folder_log.info('{0} [input dir is not dir] [input_dir={1}]'.format(self.log_base, input_dir))
            if self.switch:
                time.sleep(self.sleep_time)
                continue
            else:
                return

        output_dir = os.path.join(os.path.dirname(input_dir), 'Output')
        completed_output_dir = os.path.join(output_dir, 'Completed')
        failed_output_dir = os.path.join(output_dir, 'Failed')
        os.makedirs(output_dir, exist_ok=True)
        os.makedirs(completed_output_dir, exist_ok=True)
        os.makedirs(failed_output_dir, exist_ok=True)

        os_error_filename_set = set()
        while self.switch:
            # 1. 从input dir获取pdf or image
            list_dir = os.listdir(input_dir)
            if not list_dir and len(os_error_filename_set) == 0:
                self.folder_log.info('{0} [input dir empty] [input_dir={1}]'.format(self.log_base, input_dir))
                time.sleep(self.sleep_time)
                continue
            all_file_set = set(list_dir)
            true_file_set = all_file_set - os_error_filename_set
            if len(true_file_set) == 0 and len(os_error_filename_set) > 0:
                true_file_set.add(os_error_filename_set.pop())
            for name in true_file_set:
                time.sleep(10)
                unique_folder_name = '{0}_{1}'.format(datetime.now().strftime('%Y-%m-%d_%H:%M:%S'), name)
                path = os.path.join(input_dir, name)

                try:
                    if not os.path.exists(path):
                        self.folder_log.info('{0} [path is not exists] [path={1}]'.format(self.log_base, path))
                        continue
                    elif os.path.isdir(path):
                        self.folder_log.info('{0} [dir start] [path={1}]'.format(self.log_base, name))
                        self.folder_process(unique_folder_name, path, output_dir)
                        completed_path = os.path.join(completed_output_dir, unique_folder_name)
                        shutil.move(path, completed_path)
                        self.folder_log.info('{0} [dir end] [path={1}]'.format(self.log_base, name))
                    else:
                        self.folder_log.info('{0} [path is not dir] [path={1}]'.format(self.log_base, path))
                        failed_path = os.path.join(failed_output_dir, unique_folder_name)
                        shutil.move(path, failed_path)
                except OSError:
                    os_error_filename_set.add(name)
                    self.folder_log.error('{0} [os error] [path={1}] [error={2}]'.format(
                        self.log_base, path, traceback.format_exc()))
                except Exception as e:
                    try:
                        self.folder_log.error('{0} [file error] [path={1}] [error={2}]'.format(self.log_base, path,
                                                                                               traceback.format_exc()))
                        failed_path = os.path.join(failed_output_dir, unique_folder_name)
                        shutil.move(path, failed_path)
                    except Exception as e:
                        os_error_filename_set.add(name)
                        self.folder_log.error('{0} [file move error] [path={1}] [error={2}]'.format(
                            self.log_base, path, traceback.format_exc()))

    def handle(self, *args, **kwargs):
        process_list = []

        process = Process(target=self.main_folder_process, args=(self.input_dir, ))
        process_list.append(process)

        for p in process_list:
            p.start()
        for p in process_list:
            p.join()

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