import os
import json
import base64
import asyncio
import logging
from typing import Any, Optional, Tuple, Dict
import aiofiles  # pip install aiofiles
from config import SEARCH_DIR,IMAGE_DIR
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(os.path.basename(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 find_image_record(content_list: Any, image_filename: str) -> Optional[dict]:
    """Find an image record by the basename stored in ``img_path``."""
    if not isinstance(content_list, list) or not image_filename:
        return None

    target_name = os.path.basename(str(image_filename).strip().replace('\\', '/'))
    if not target_name:
        return None

    for item in content_list:
        if not isinstance(item, dict) or item.get('type') != 'image':
            continue

        img_path = item.get('img_path')
        if isinstance(img_path, str):
            path_name = os.path.basename(img_path.strip().replace('\\', '/'))
            if path_name == target_name:
                return item


    return None


async def find_image_record_by_pdf(
    directory: str,
    pdf_filename: str,
    image_filename: str,
    encoding: str = 'utf-8',
) -> Optional[dict]:
    """Find an image record from the content list matching a PDF filename."""
    content_list, _ = await find_and_read_content_list(
        directory, pdf_filename, encoding=encoding
    )
    return find_image_record(content_list, image_filename)


async def extract_image_caption(
    search_directory: str,
    pdf_filename: str,
    image_filename: str,
) -> list[str]:
    """?? PDF ??????????? image_caption?"""
    image_record = await find_image_record_by_pdf(
        directory=search_directory,
        pdf_filename=pdf_filename,
        image_filename=image_filename,
    )

    if not image_record:
        return []

    image_caption = image_record.get("image_caption", [])
    if isinstance(image_caption, str):
        return [image_caption] if image_caption.strip() else []
    if isinstance(image_caption, list):
        return [str(item).strip() for item in image_caption if str(item).strip()]
    return []


async def find_and_read_md_file(
    directory: str,
    original_filename: str,
    encoding: str = 'utf-8'
) -> Tuple[Optional[str], Optional[str]]:
    """Find and read the Markdown file matching the uploaded document basename."""
    file_name, _ = os.path.splitext(original_filename)
    target_filename = f"{file_name}.md"
    logger.info(f"正在查找 Markdown 文件: {target_filename}")

    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, errors='ignore') as f:
            return await f.read(), file_path
    except Exception as e:
        logger.error(f"读取 Markdown 文件 {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,
    local_image_dir: Optional[str] = None,
    concurrency: int = 32,
) -> Optional[Tuple[list, Dict[str, str], Optional[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, md_content)
        失败: None
    """
    search_dir = SEARCH_DIR
    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))}")

    md_content, md_path = await find_and_read_md_file(search_dir, input_file_name)
    if md_content is None:
        logger.info("未找到对应的 Markdown 文件。")
    else:
        logger.info(f"找到 Markdown 文件: {md_path}")
        logger.info(f"Markdown 内容长度: {len(md_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, md_content


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

    if result is None:
        return 1

    content, images_dict, md_content = 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__":
    result = asyncio.run(
        extract_image_caption(
            r"E:\ZKYNLP\Hjunproject\project0506\kgrag",
            "163-06A0014-B01001_发动机-维修手册.pdf",
            "23fa41fea5e5b5458d6f7df3f8e7e8033b6665bd23203946646651c3e56c2da7.jpg",
        )
    )
    print(result)

