from loguru import logger
import configparser
import os
import time

from magic_pdf.libs.commons import join_path
from magic_pdf.libs.ocr_content_type import ContentType
from magic_pdf.libs.pdf_image_tools import cut_image

from multiprocessing import Pool
# vllm：
from magic_pdf.dict2md.ocr_vllm_client import PredictClient,compress_image

# 普通 非vllm
# from magic_pdf.dict2md.ocr_client import PredictClient,compress_image

text = '解析图片内容，直接返回一段带有逻辑性的中文书面语描述，要求表达精准，不脱离图片中的实际内容，不要带换行,文中所有的名词不要用指代词'
client = None
def ocr_image(image_path):
    start = time.time()
    compress_image(image_path)
    txt = os.getpid()
    # global client
    #logger.info(f'image_path:{image_path}')
    generated_text = client.predict(image_path, text)
    end = time.time()
    logger.info(f'qwen解析{image_path}表格的内容为：{generated_text},耗时为：{end - start}')
    return generated_text

def ocr_cut_image_and_table(ocr_status,config_path,local_image_dir,spans, page, page_id, pdf_bytes_md5, imageWriter):
    def return_path(type):
        return join_path(pdf_bytes_md5, type)

    pool = Pool(4)
    global client
    config = configparser.ConfigParser()
    config.read(config_path)
    url = config.get('server', 'ocr_server')
    client = PredictClient(url)

    if not ocr_status:
        logger.warning(f'Health check failed. The server at "{url}" is not responding as expected.')
        logger.info(f'Qwen ocr解析服务无法正常运行,暂不使用qwen解析表格服务')

    for span in spans:
        span_type = span['type']
        if span_type == ContentType.Image:
            if not check_img_bbox(span['bbox']):
                continue
            span['image_path'] = cut_image(span['bbox'], page_id, page, return_path=return_path('images'),
                                           imageWriter=imageWriter)
        elif span_type == ContentType.Table:
            if not check_img_bbox(span['bbox']):
                continue
            image_path = cut_image(span['bbox'], page_id, page, return_path=return_path('tables'),
                                   imageWriter=imageWriter)
            image_path = join_path(local_image_dir, image_path)
            if ocr_status:

                txt = pool.apply_async(ocr_image, args=(image_path,)).get()
                span['image_path'] = str(txt)
            else:
                span['image_path'] = f"----------------图片路径为({image_path})，请检查qwen ocr服务，重新运行文件解析-------------------  \n"
    pool.close()
    pool.join()

    return spans


def txt_save_images_by_bboxes(page_num: int, page, pdf_bytes_md5: str,
                              image_bboxes: list, images_overlap_backup: list, table_bboxes: list,
                              equation_inline_bboxes: list,
                              equation_interline_bboxes: list, imageWriter) -> dict:
    """
    返回一个dict, key为bbox, 值是图片地址
    """
    image_info = []
    image_backup_info = []
    table_info = []
    inline_eq_info = []
    interline_eq_info = []

    # 图片的保存路径组成是这样的： {s3_or_local_path}/{book_name}/{images|tables|equations}/{page_num}_{bbox[0]}_{bbox[1]}_{bbox[2]}_{bbox[3]}.jpg

    def return_path(type):
        return join_path(pdf_bytes_md5, type)

    for bbox in image_bboxes:
        if not check_img_bbox(bbox):
            continue
        image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
        image_info.append({"bbox": bbox, "image_path": image_path})

    for bbox in images_overlap_backup:
        if not check_img_bbox(bbox):
            continue
        image_path = cut_image(bbox, page_num, page, return_path("images"), imageWriter)
        image_backup_info.append({"bbox": bbox, "image_path": image_path})

    for bbox in table_bboxes:
        if not check_img_bbox(bbox):
            continue
        image_path = cut_image(bbox, page_num, page, return_path("tables"), imageWriter)
        table_info.append({"bbox": bbox, "image_path": image_path})

    return image_info, image_backup_info, table_info, inline_eq_info, interline_eq_info


def check_img_bbox(bbox) -> bool:
    if any([bbox[0] >= bbox[2], bbox[1] >= bbox[3]]):
        logger.warning(f"image_bboxes: 错误的box, {bbox}")
        return False
    return True
