# -*- coding: utf-8 -*-
# @Author:逐梦苍穹
# @Time: 2024/1/18 1:23
import os
import cv2
import fitz
import shutil
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from cnocr import CnOcr
from tqdm import tqdm
from PIL import Image
from docx import Document


def mkdir(folder_path):
    os.makedirs(folder_path, exist_ok=True)


def get_max_folder_number(folder_path):
    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]
    return len(image_files)


def crop_top_and_save(args):
    i, image_path, output_path, crop_percentage = args
    image = cv2.imread(image_path)
    height = image.shape[0]
    crop_pixels = int(height * crop_percentage)
    cropped_image = image[crop_pixels:, :]
    cv2.imwrite(output_path, cropped_image)
    print(f"Image {i + 1} cropped and saved: {output_path}")


def startCutImage(input_folder, output_folder, crop_percentage, max_workers_input):
    mkdir(output_folder)
    max_folder_number = get_max_folder_number(input_folder)
    with ThreadPoolExecutor(max_workers=max_workers_input) as executor:
        args_list = [
            (i, os.path.join(input_folder, f'{i + 1}.jpg'), os.path.join(output_folder, f'{i + 1}.jpg'), crop_percentage)
            for i in range(max_folder_number)
        ]
        executor.map(crop_top_and_save, args_list)


def get_max_folder_number(folder_path):
    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]
    return len(image_files)


def delete_images(folder_path, start_index, end_index):
    for i in range(start_index, end_index + 1):
        imageStr = str(i) + ".jpg"
        file_path = os.path.join(folder_path, imageStr)
        os.remove(file_path)
        print(f"Deleted: {file_path}")


def startDeleteImage(folder_path, end_index, back_start_index_input):
    max_folder_number = get_max_folder_number(folder_path)
    if max_folder_number > 0:
        start_index = 1
        back_start_index = max_folder_number - back_start_index_input + 1
        back_end_index = max_folder_number
        delete_images(folder_path, start_index, end_index)
        delete_images(folder_path, back_start_index, back_end_index)
    else:
        print("No valid image files found.")


def convert_page_to_image(pdf_path, output_folder, page_number):
    pdf_document = fitz.open(pdf_path)
    page = pdf_document[page_number]
    image_list = page.get_images(full=True)
    for img_index, img_info in enumerate(image_list):
        image_index = img_info[0]
        base_image = pdf_document.extract_image(image_index)
        image_bytes = base_image["image"]
        image_filename = os.path.join(output_folder, f"{page_number + 1}.jpg")
        with open(image_filename, "wb") as image_file:
            image_file.write(image_bytes)
    pdf_document.close()


def pdf_to_images(pdf_path, output_folder, num_workers):
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)
    pdf_document = fitz.open(pdf_path)
    total_pages = pdf_document.page_count
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = [executor.submit(convert_page_to_image, pdf_path, output_folder, page_number)
                   for page_number in range(total_pages)]
        for future in tqdm(futures, desc="Converting pages", total=total_pages):
            future.result()


def get_max_folder_number(folder_path):
    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]
    return len(image_files)


def ocr_and_append(image_str, startList, event, lock):
    pagelist = []
    image_number = int(os.path.basename(image_str).split('.')[0])
    image = Image.open(image_str)
    ocr = CnOcr(debug=False)
    res = ocr.ocr(image)
    for j in range(len(res)):
        pagelist.append(res[j]['text'])
    with lock:
        startList[image_number - 11] = pagelist
    print(f"Image {image_number} OCR completed.")
    event.set()


def process_images_with_ocr(start, end, startList, event, max_workers_input, lock):
    with ThreadPoolExecutor(max_workers=max_workers_input) as executor:
        futures = [executor.submit(ocr_and_append, f'output-image-cut/{i}.jpg', startList, event, lock) for i in
                   tqdm(range(start, end + 1), desc="OCR Processing")]
        for future in as_completed(futures):
            future.result()


def startOCR(file_folder, test_end_index, word_file_output_path, max_workers_input):
    max_folder_number = get_max_folder_number(file_folder)
    start = 1 + test_end_index
    end = test_end_index + max_folder_number
    startList = [''] * max_folder_number
    lock = threading.Lock()
    event = threading.Event()
    process_images_with_ocr(start, end, startList, event, max_workers_input, lock)
    event.wait()
    print(len(startList))
    for i in startList:
        print(len(i))
    print(startList)
    startWriteWordFile(word_file_output_path, startList)


def startWriteWordFile(word_file_path, listPages):
    doc = Document()
    for i in listPages:
        write_list_to_word(word_file_path, i, doc)
        doc.add_page_break()
    doc.save(word_file_path)


def write_list_to_word(word_file, my_list, doc):
    endList = []
    for item in my_list:
        item = item.strip()
        if len(item) < 10:
            if endList != None and len(endList) > 0:
                doc.add_paragraph("".join(endList))
                endList = []
            doc.add_paragraph(item)
        else:
            endList.append(item)
            if item == my_list[-1]:
                doc.add_paragraph("".join(endList))
                endList = []


def main():
    try:
        crop_percentage = 0.0
        end_index = 0
        back_start_index_input = 0
        print("==========================")
        print("这是一个OCR识别程序，作者CSDN@逐梦苍穹，该程序不做商用请切记！")
        print("==========================")
        pdf_file_path = str(input("请输入PDF文件的绝对路径："))
        print("该PDF文档是否是书籍?")
        isBookPdf = int(input("(是，请输入1；否，请输入0):"))
        if isBookPdf == 1:
            print("输入要删除非正文的页面(封面、目录、结尾)")
            end_index = int(input("需要删除该文档的前几页: "))
            back_start_index_input = int(input("需要删除该文档的倒数前几页: "))
        print("询问是否需要截取图片区域?")
        isCutImage = int(input("(是，请输入1；否，请输入0):"))
        if isCutImage == 1:
            crop_percentage = float(input("请输入截取图片区域的纵向比例："))
        max_workers = int(input("请输入本次运行过程中想要使用的进程数："))
        if max_workers <= 0:
            max_workers = 1
    except Exception as e:
        print("参数有误,报错信息如下：", e)
    try:
        pdf_to_images(pdf_file_path, 'output-image', max_workers)
        if isCutImage == 1:
            startCutImage('output-image', 'output-image-cut', crop_percentage, max_workers)
        if isBookPdf == 1:
            startDeleteImage('output-image-cut', end_index, back_start_index_input)
        file_name = os.path.splitext(os.path.basename(pdf_file_path))[0]
        directory_path = os.path.dirname(pdf_file_path)
        word_file_path = directory_path + '\\' + file_name + '(OCR_XZL).docx'
        startOCR(r'output-image-cut', end_index, word_file_path, max_workers)
    except Exception as e:
        print("处理过程出错，报错信息如下：", e)
    try:
        delete_path_1 = r"output-image"
        delete_path_2 = r"output-image-cut"
        shutil.rmtree(delete_path_1)
        shutil.rmtree(delete_path_2)
        print(f"Folder '{delete_path_1}' successfully removed.")
        print(f"Folder '{delete_path_2}' successfully removed.")
        print("OCR处理完成！")
    except Exception as e:
        print("删除过程性文件出错，原因如下：", e)


if __name__ == '__main__':
    main()