import os
import json
import base64
import asyncio
import logging
from typing import Optional, Tuple, Dict
import aiofiles  # pip install aiofiles

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


async def find_and_read_content_list(
    directory: str,
    original_filename: str,
    encoding: str = 'utf-8'
) -> Tuple[Optional[list], Optional[str]]:
    """根据原始文件名查找并读取对应的 _content_list.json 文件。"""
    file_name, _ = os.path.splitext(original_filename)
    target_filename = f"{file_name}_content_list.json"
    logger.info(f"正在查找文件: {target_filename}")

    # os.walk 没有原生 async 版本,放到线程里执行避免阻塞事件循环
    def _walk_for_file():
        for root, _, files in os.walk(directory):
            if target_filename in files:
                return os.path.join(root, target_filename)
        return None

    file_path = await asyncio.to_thread(_walk_for_file)
    if file_path is None:
        return None, None

    try:
        async with aiofiles.open(file_path, 'r', encoding=encoding) as f:
            text = await f.read()
        # JSON 解析是 CPU 操作,大文件可考虑放线程
        return json.loads(text), file_path
    except json.JSONDecodeError:
        logger.error(f"文件 {file_path} 不是有效的 JSON 格式。")
        return None, None
    except Exception as e:
        logger.error(f"读取文件 {file_path} 时发生错误: {e}")
        return None, None


def extract_all_jpg_filenames(doc_data: list) -> set:
    """从文档数据中提取所有 .jpg 图片的纯文件名。纯 CPU 操作,不需要 async。"""
    result_set = set()
    for item in doc_data:
        if 'img_path' in item and item['img_path']:
            img_path = item['img_path']
            if img_path.endswith('.jpg'):
                result_set.add(os.path.basename(img_path))
    return result_set


async def _encode_one_image(
    img_file: str,
    local_image_dir: str,
    mime_map: dict,
) -> Tuple[str, Optional[str]]:
    """编码单张图片,返回 (文件名, data_uri 或 None)。"""
    img_path = os.path.join(local_image_dir, img_file)
    if not os.path.isfile(img_path):
        logger.warning(f"图片不存在: {img_path}")
        return img_file, None

    try:
        async with aiofiles.open(img_path, "rb") as f:
            img_bytes = await f.read()
        # base64 编码是 CPU 操作,小文件直接做即可;大文件可放线程
        img_base64 = base64.b64encode(img_bytes).decode('utf-8')
        ext = img_file.lower().rsplit('.', 1)[-1]
        mime_type = f"image/{mime_map.get(ext, 'png')}"
        return img_file, f"data:{mime_type};base64,{img_base64}"
    except Exception as e:
        logger.error(f"编码图片失败 {img_file}: {str(e)}")
        return img_file, None


async def encode_images_to_base64(
    local_image_dir: str,
    referenced_images: set,
    concurrency: int = 32,
) -> Dict[str, str]:
    """并发将引用的图片编码为 base64。"""
    if not os.path.exists(local_image_dir) or not referenced_images:
        return {}

    mime_map = {'jpg': 'jpeg', 'jpeg': 'jpeg', 'png': 'png', 'gif': 'gif', 'webp': 'webp'}

    # 用 Semaphore 限制并发数,避免一次性打开几百个文件句柄
    sem = asyncio.Semaphore(concurrency)

    async def _bounded(img_file: str):
        async with sem:
            return await _encode_one_image(img_file, local_image_dir, mime_map)

    results = await asyncio.gather(*(_bounded(f) for f in referenced_images))
    return {name: data for name, data in results if data is not None}


async def process_document(
    input_file_name: str,
    search_dir: Optional[str] = "/app/mineru_output",
    local_image_dir: Optional[str] = None,
    concurrency: int = 32,
) -> Optional[Tuple[list, Dict[str, str]]]:
    """
    异步处理文档:查找 content_list.json,提取图片引用,并发编码为 base64。

    Args:
        input_file_name: 原始文件名(如 "xxx.pdf")
        search_dir: 搜索 _content_list.json 的根目录
        local_image_dir: 图片所在目录,None 则自动推断
        concurrency: 图片编码的并发数上限,默认 32

    Returns:
        成功: (content_list, images_dict)
        失败: None
    """
    content, content_path = await find_and_read_content_list(search_dir, input_file_name)
    if content is None:
        logger.error("❌ 未找到指定的 _content_list.json 文件。")
        return None

    logger.info(f"✅ 找到文件: {content_path}")
    logger.info(f"内容类型: {type(content).__name__}, 长度: {len(str(content))}")

    if local_image_dir is None:
        local_image_dir = os.path.join(os.path.dirname(content_path), "images")
    logger.info(f"图片目录: {local_image_dir}")

    referenced_images = extract_all_jpg_filenames(content)
    logger.info(f"引用图片数量: {len(referenced_images)}")

    images_dict = await encode_images_to_base64(
        local_image_dir, referenced_images, concurrency=concurrency
    )
    logger.info(f"成功编码 {len(images_dict)} 张图片")

    return content, images_dict


async def main():
    result = await process_document(
        input_file_name="163-06A0014-B01001_发动机-维修手册.pdf",
        search_dir="/app/mineru_output",
        concurrency=32,
    )

    if result is None:
        return 1

    content, images_dict = result
    print(f"\n=== 处理完成 ===")
    print(f"文档段落数: {len(content)}")
    print(f"图片数量: {len(images_dict)}")
    if images_dict:
        first_key = next(iter(images_dict))
        print(f"示例图片 [{first_key}]: {images_dict[first_key][:80]}...")
    return 0


if __name__ == "__main__":
    exit_code = asyncio.run(main())
    exit(exit_code)