import fitz
from urllib import parse
from urllib.parse import urlparse
import os

import asyncio
import aiohttp
from aiohttp import ClientSession
from utils.http_util import url_replace_fastgpt, download_file
from utils.common_util import random_str
from loguru import logger
import json
from core.config import ocr_url


class OCRUtil:
    def __init__(self):
        pass

    # ocr异步接口
    async def ocr_requests_async(self, session, file_path):
        # 文件上传通常使用multipart/form-data编码格式
        # file_path = 'file/1.webp'
        async with session.post(ocr_url, data={'file': open(file_path, 'rb')}) as response:
            # logger.info(f'开始请求:{file_path}')
            rsp = await response.text()
            # logger.info(f'{file_path}结束请求:{json.loads(rsp)["msg"]}')
            return rsp, file_path

    # 异步请求ocr image接口
    async def ocr_image_async(self, path_list):
        timeout = aiohttp.ClientTimeout(total=600)
        connector = aiohttp.TCPConnector(limit=1)
        async with ClientSession(connector=connector, timeout=timeout) as session:
            tasks = [self.ocr_requests_async(session, file_path) for file_path in path_list]
            responses = await asyncio.gather(*tasks)
            res_dict = {}
            for response_text, file_path in responses:
                rsp_json = json.loads(response_text)
                if 'data' not in rsp_json:
                    logger.error(f'ocr_image_async {file_path} error:{rsp_json["msg"]}')
                    continue
                else:
                    content = rsp_json['data']['strRes']
                    # logger.info(f'ocr_image_async append {file_path} success. content:{content[:100]}')
                # add to dict
                page_num = int(self.get_pdf_2_img_page_num(file_path))
                res_dict[page_num] = rsp_json['data']['strRes']
            # 根据页数排序
            logger.info(f'ocr_image_async finish. all pages:{len(res_dict)}')
            sorted_values = [res_dict[key] for key in sorted(res_dict)]
            return sorted_values

    def set_pdf_2_img_page(self, path, page_idx):
        return f'{path}_{page_idx + 1}.png'

    def get_pdf_2_img_page_num(self, path):
        split_path = path.split('_')
        # for example : 'xx_10.png' ==> get 10
        return split_path[-1][:-4]

    def pdf_2_img(self, path, zoom_x=1, zoom_y=1):
        # 打开PDF文件
        pdf = fitz.open(path)
        pdf_list = []
        # 逐页读取PDF
        for pg in range(0, pdf.page_count):
            page = pdf[pg]
            # 设置缩放和旋转系数
            trans = fitz.Matrix(zoom_x, zoom_y)
            pm = page.get_pixmap(matrix=trans, alpha=False)
            # img save
            dest_png = self.set_pdf_2_img_page(path, pg)
            pm.save(dest_png)
            pdf_list.append(dest_png)
        pdf.close()
        return pdf_list

    def ocr_download_path(self, url):
        logger.info(f'ocr url:{url}')
        # url替换
        url = url_replace_fastgpt(url)
        # url分析
        url_parsed = urlparse(url)
        query_dict = parse.parse_qs(url_parsed.query)
        if 'filename' in query_dict:
            filename = query_dict.get('filename')[0]
        else:
            filename = f'{random_str()}.pdf'
        # 获取文件内容
        dest_path = f'ocr/{filename}'
        download_file(url, dest_path)
        return dest_path

    async def ocr_result_pdf(self, dest_path):
        # 解析返回结果
        pdf_list = self.pdf_2_img(dest_path)
        # logger.info(f'pdf 生成完毕:{pdf_list}')
        result = await self.ocr_image_async(pdf_list)
        # 删除文件
        for pdf in pdf_list:
            if os.path.exists(pdf):
                os.remove(pdf)
        return result


if __name__ == '__main__':
    ocr_util = OCRUtil()
    result = asyncio.run(ocr_util.ocr_result_pdf('/home/ccran/lufa-contract/tmp/财03 2023年国科京东方评估合同.pdf'))
    print(f'len(result):{len(result)}')
    print(result)
