dataflow_agent/workflow/wf_paper2ppt_parallel.py

from __future__ import annotations

import asyncio
import json
import os
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

from dataflow_agent.graphbuilder.graph_builder import GenericGraphBuilder
from dataflow_agent.logger import get_logger
from dataflow_agent.state import Paper2FigureState
from dataflow_agent.utils import get_project_root
from dataflow_agent.workflow.registry import register
from dataflow_agent.agentroles import create_react_agent
from dataflow_agent.toolkits.imtool.req_img import generate_or_edit_and_save_image_async
from dataflow_agent.toolkits.imtool.ppt_tool import convert_images_dir_to_pdf_and_ppt, convert_images_dir_to_pdf_and_ppt_api

log = get_logger(__name__)


def _ensure_result_path(state: Paper2FigureState) -> str:
    """
    统一 paper2ppt workflow 的根输出目录：
    - 若 state.result_path 已存在（通常由调用方传入），直接使用；
    - 否则：使用 get_project_root()/outputs/paper2ppt/<timestamp> 初始化，并写回 state.result_path。
    """
    raw = getattr(state, "result_path", None)
    if raw:
        return raw

    root = get_project_root()
    ts = int(time.time())
    base_dir = (root / "outputs" / "paper2ppt" / str(ts)).resolve()
    base_dir.mkdir(parents=True, exist_ok=True)
    state.result_path = str(base_dir)
    return state.result_path


def _abs_path(p: str) -> str:
    if not p:
        return ""
    try:
        return str(Path(p).expanduser().resolve())
    except Exception:
        return p


def _is_table_asset(asset_ref: Optional[str]) -> bool:
    """
    你给的约定：asset 是 Table 时，通过 asset_ref: "Table 2" 这种字符串标记。
    """
    if not asset_ref:
        return False
    s = str(asset_ref).strip().lower()
    return s.startswith("table")


def _serialize_prompt_dict(d: Dict[str, Any]) -> str:
    """
    把 dict 安全序列化为 prompt 文本（中文不转义）。
    """
    try:
        return json.dumps(d, ensure_ascii=False, indent=2)
    except Exception:
        # 兜底：不要因为序列化失败而中断
        return str(d)


def _normalize_single_asset_ref(asset_ref: str) -> str:
    """
    规范化 asset_ref，仅保留第一张图的路径/文件名。

    当前版本不支持多图编辑：
    - 如果 asset_ref 中包含逗号等分隔符，如 "a.jpg,b.jpg"，
      只取第一段 "a.jpg"。
    - TODO: 后续可以扩展多图 asset_ref 支持。
    """
    if not asset_ref:
        return ""
    s = str(asset_ref).strip()
    if not s:
        return ""

    # 简单按逗号切分，保留第一个
    parts = [p.strip() for p in s.split(",") if p.strip()]
    if not parts:
        return ""

    if len(parts) > 1:
        log.warning(
            "[paper2ppt] asset_ref 包含多张图片，仅使用第一张。"
            f" raw={asset_ref!r}, first={parts[0]!r}  # TODO: 支持多图 asset_ref"
        )

    return parts[0]


async def _make_prompt_for_structured_page(item: Dict[str, Any], style: str, state: Paper2FigureState) -> Tuple[str, Optional[str], bool]:
    """
    根据结构化 page item 生成:
    - prompt
    - image_path (如果是编辑模式)
    - use_edit

    规则：
    1) asset 为空：text2img，用 “json(去asset)” + “根据上述内容生成{style}风格的PPT”
    2) asset 是图片路径：img2img/edit，用 “json(去asset)” + “把这个图作为PPT的一部分...”
    3) asset 是 Table（asset_ref="Table 2"）：先提取 table png（这里先占位），再走 edit
    """
    asset_ref = item.get("asset_ref") or item.get("asset") or item.get("assetRef") or ""
    asset_ref = str(asset_ref).strip() if asset_ref is not None else ""
    # TODO 当前版本仅支持单图 asset_ref；若包含多图，仅保留第一张。
    asset_ref = _normalize_single_asset_ref(asset_ref)

    prompt_dict = dict(item)
    for k in ["asset_ref", "asset", "assetRef", "asset_type", "type"]:
        if k in prompt_dict:
            prompt_dict.pop(k, None)

    base = _serialize_prompt_dict(prompt_dict)

    if not asset_ref:
        prompt = f"{base}\n\n根据上述内容。生成{style}风格的 PPT 图像, \n 使用语言：{state.request.language}"
        return prompt, None, False

    # table 走占位提取
    if _is_table_asset(asset_ref):
        # 优先使用 item 自己带的表格图（如果调用方已经生成过）
        table_img_path = item.get("table_img_path") or item.get("table_png_path") or ""
        table_img_path = str(table_img_path).strip()

        # 若没有，则调用 table_extractor agent：生成 html->png，并写入 state.table_img_path
        if not table_img_path:
            state.asset_ref = asset_ref
            agent = create_react_agent(
                name="table_extractor",
                temperature=0.1,
                max_retries=6,
                parser_type="json",
            )
            state = await agent.execute(state=state)

            table_img_path = str(getattr(state, "table_img_path", "") or "").strip()
            log.critical(f'[table_img_path 表格图像路径]:   {table_img_path}')

        if not table_img_path:
            raise ValueError(f"[paper2ppt] 表格提取失败，未得到 table_img_path。asset_ref={asset_ref}")

        image_path = _resolve_asset_path(table_img_path, state)
        # 如果表格图像不存在，则退化为 text2img：不走编辑，返回 use_edit=False
        if not image_path or not os.path.exists(image_path):
            log.error(f"[paper2ppt] 表格图像文件不存在: {image_path!r} (asset_ref={asset_ref})")
            prompt = f"{base}\n\n根据上述内容生成{style}风格的 PPT 图像, \n 使用语言：{state.request.language}"
            return prompt, None, False

        prompt = f"{base}\n\n根据上述内容绘制ppt，把这个图作为PPT的一部分。生成{style}风格的PPT. \n 使用语言：{state.request.language} !!!"
        return prompt, image_path, True

    # 默认：当作图片路径，走编辑
    image_path = _resolve_asset_path(asset_ref, state)
    # 如果图片不存在，则退化为 text2img：不走编辑，返回 use_edit=False
    if not image_path or not os.path.exists(image_path):
        log.error(f"[paper2ppt] 图片文件不存在: {image_path!r} (asset_ref={asset_ref})")
        prompt = f"{base}\n\n根据上述内容生成{style}风格的 PPT 图像, \n 使用语言：{state.request.language}"
        return prompt, None, False

    prompt = f"{base}\n\n根据上述内容绘制ppt，把这个图作为PPT的一部分。生成{style}风格的PPT. \n 使用语言：{state.request.language} !!!"
    return prompt, image_path, True


def _resolve_asset_path(asset_ref: str, state: Paper2FigureState) -> str:
    """
    根据 state 解析 asset 引用为绝对路径。

    规则：
    - 为空直接返回 ""；
    - 绝对路径或以 ~ 开头：直接通过 _abs_path 规范化；
    - 相对路径：
        * 优先挂在 state.mineru_root（MinerU 输出根目录）下；
        * 否则挂在 state.result_path 下；
        * 再否则退化为当前工作目录下的相对路径解析（_abs_path）。
    """
    if not asset_ref:
        return ""
    s = str(asset_ref).strip()
    if not s:
        return ""

    p = Path(s)

    # 已经是绝对路径，或者显式使用家目录
    if p.is_absolute() or s.startswith("~"):
        return _abs_path(s)

    base_dir = getattr(state, "mineru_root", None) or getattr(state, "result_path", None)
    log.critical(f'[base_dir _resolve_asset_path]:   {base_dir}')
    
    if base_dir:
        try:
            return str((Path(base_dir) / p).resolve())
        except Exception:
            return _abs_path(s)

    return _abs_path(s)


def _extract_image_path_from_pagecontent_item(item: Any) -> Optional[str]:
    """
    支持 pagecontent 直接是图片路径的几种形态：
    - "/abs/xxx.png"
    - {"ppt_img_path": "/abs/xxx.png"}
    - {"img_path": "/abs/xxx.png"}
    - {"path": "/abs/xxx.png"}
    """
    if not item:
        return None
    if isinstance(item, str):
        return item.strip()
    if isinstance(item, dict):
        for k in ["ppt_img_path", "img_path", "path", "image_path"]:
            v = item.get(k)
            if v:
                return str(v).strip()
    return None


@register("paper2ppt_parallel")
def create_paper2ppt_parallel_graph() -> GenericGraphBuilder:  # noqa: N802
    """
    Workflow factory: dfa run --wf paper2ppt_parallel

    功能：
    - 并行版本：并发调用 AI 接口生成/编辑所有 PPT 页面
    - 若 state.gen_down == False：批量生成/编辑每页 PPT 图，保存到统一目录
    - 若 state.gen_down == True：按 0-based edit_page_num 对已有页面图做二次编辑（edit_page_prompt）
    """
    builder = GenericGraphBuilder(state_model=Paper2FigureState, entry_point="_start_")

    def _start_(state: Paper2FigureState) -> Paper2FigureState:
        _ensure_result_path(state)
        state.pagecontent = state.pagecontent or []
        state.generated_pages = state.generated_pages or []
        # 兼容：有些调用方把 style 放 state.style，而不是 request.style
        if not getattr(state.request, "style", None) and getattr(state, "style", None):
            state.request.style = getattr(state, "style")
        return state

    def _route(state: Paper2FigureState) -> str:
        # 如果是 all_edited_down，说明用户只想打包下载，不需要生成或编辑，直接去导出
        if getattr(state.request, "all_edited_down", False):
            return "export_ppt_assets"

        # gen_down == False: 第一次批量生成
        if not getattr(state, "gen_down", False):
            return "generate_pages"
        # gen_down == True: 进入按页编辑
        return "edit_single_page"

    async def generate_pages(state: Paper2FigureState) -> Paper2FigureState:
        """
        批量生成/编辑页面图（并行版本）：
        - pagecontent 是结构化 list[dict]：按 asset 规则决定 text2img / img2img
        - pagecontent 直接是图片路径列表：逐页用“修改成xxx风格”编辑
        
        并发调用：使用 asyncio.gather 同时处理所有页面。
        """
        import asyncio

        async def _call_image_api_with_retry(coro_factory, retries: int = 3, delay: float = 1.0) -> bool:
            """
            对图像生成/编辑进行最多 retries 次重试。
            - 成功：返回 True
            - 多次失败：返回 False（由上层决定如何处理当前页）
            """
            last_err: Optional[Exception] = None
            for attempt in range(1, retries + 1):
                try:
                    await coro_factory()
                    return True
                except Exception as e:  # noqa: BLE001
                    last_err = e
                    log.error(f"[paper2ppt] image gen failed attempt {attempt}/{retries}: {e}")
                    if attempt < retries:
                        try:
                            await asyncio.sleep(delay)
                        except Exception:
                            # sleep 失败不影响后续重试
                            pass
            log.error(f"[paper2ppt] image gen failed after {retries} attempts, skip this page. last_err={last_err}")
            return False

        result_root = Path(_ensure_result_path(state))
        img_dir = result_root / "ppt_pages"
        img_dir.mkdir(parents=True, exist_ok=True)

        style = getattr(state.request, "style", None) or "kartoon"
        aspect_ratio = getattr(state, "aspect_ratio", None) or "16:9"

        # 清空旧数据（避免重复执行堆积）
        state.generated_pages = []
        
        # 定义单个页面处理任务
        async def _process_single_page(idx: int, item: Any) -> Dict[str, Any]:
            """
            处理单个页面：返回生成的 result item (dict)。
            如果失败，result item 中的 generated_img_path 为 None。
            """
            save_path = str((img_dir / f"page_{idx:03d}.png").resolve())
            
            # Case B: pagecontent 本身就是图片路径
            direct_img_path = _extract_image_path_from_pagecontent_item(item)
            is_direct_image_list = bool(direct_img_path) and (
                isinstance(item, str)
                or (isinstance(item, dict) and set(item.keys()).intersection({"ppt_img_path", "img_path", "path", "image_path"}))
            )

            if is_direct_image_list and (not isinstance(item, dict) or ("title" not in item and "layout_description" not in item)):
                # 规则 2：只做风格化编辑
                image_path = _abs_path(direct_img_path)
                # 强化提示词，确保模型进行重绘而不是原图输出
                prompt = (
                    f"Please beautify and re-design this PowerPoint slide image. "
                    f"Keep all the original text and structure, but completely transform it into a professional, "
                    f"visually stunning presentation slide in {style} style. "
                    f"Make sure the colors, layout, and background are improved significantly."
                )
                log.info(f"[paper2ppt] page={idx} direct image edit: image={image_path}, save={save_path}")

                log.critical(f'[强化提示词，确保模型进行重绘而不是原图输出]: {prompt}')

                ok = await _call_image_api_with_retry(
                    lambda: generate_or_edit_and_save_image_async(
                        prompt=prompt,
                        save_path=save_path,
                        aspect_ratio=aspect_ratio,
                        api_url=state.request.chat_api_url,
                        api_key=state.request.chat_api_key or os.getenv("DF_API_KEY") ,
                        model=state.request.gen_fig_model,
                        image_path=image_path,
                        use_edit=True,
                    )
                )
                if not ok:
                    # 记录失败信息
                    return {
                        "source_img_path": image_path,
                        "generated_img_path": None,
                        "page_idx": idx,
                        "mode": "edit_direct_image_failed",
                        "style": style,
                    }

                return {
                    "source_img_path": image_path,
                    "generated_img_path": save_path,
                    "page_idx": idx,
                    "mode": "edit_direct_image",
                    "style": style,
                }

            # Case A: 结构化页面
            if not isinstance(item, dict):
                log.warning(f"[paper2ppt] page={idx} 非 dict 且非 image path，跳过。item={item}")
                return {
                    "page_idx": idx,
                    "mode": "invalid_item_skipped",
                    "generated_img_path": None,
                }

            try:
                # 注意：_make_prompt_for_structured_page 可能是 async 的，因为它可能调用 table_extractor agent
                prompt, image_path, use_edit = await _make_prompt_for_structured_page(item, style=style, state=state)
            except Exception as e:  # noqa: BLE001
                log.error(f"[paper2ppt] page={idx} prompt 构造失败: {e}")
                failed_item = dict(item)
                failed_item.update({
                    "generated_img_path": None,
                    "page_idx": idx,
                    "mode": "prompt_build_failed",
                    "style": style,
                    "error": str(e),
                })
                return failed_item

            log.info(
                f"[paper2ppt] page={idx} structured: use_edit={use_edit}, "
                f"image_path={image_path}, save={save_path}, \n 本次生成的 prompt 为:\n{prompt}"
            )

            ok = await _call_image_api_with_retry(
                lambda: generate_or_edit_and_save_image_async(
                    prompt=prompt,
                    save_path=save_path,
                    aspect_ratio=aspect_ratio,
                    api_url=state.request.chat_api_url,
                    api_key=state.request.chat_api_key or os.getenv("DF_API_KEY") ,
                    model=state.request.gen_fig_model,
                    image_path=image_path,
                    use_edit=use_edit,
                )
            )
            if not ok:
                failed_item = dict(item)
                failed_item.update({
                    "generated_img_path": None,
                    "page_idx": idx,
                    "mode": "generate_failed" if not use_edit else "edit_failed",
                    "style": style,
                })
                return failed_item

            # 成功
            out_item = dict(item)
            out_item.update({
                "generated_img_path": save_path,
                "page_idx": idx,
                "mode": "edit" if use_edit else "generate",
                "style": style,
            })
            return out_item

        # -----------------------------------------------------------
        # 并发执行逻辑
        # -----------------------------------------------------------
        page_items = state.pagecontent or []
        tasks = []
        for idx, item in enumerate(page_items):
            tasks.append(_process_single_page(idx, item))
        
        log.info(f"[paper2ppt_parallel] start generating {len(tasks)} pages concurrently...")
        start_time = time.time()
        
        # 使用 gather 并发执行所有任务
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        cost_time = time.time() - start_time
        log.info(f"[paper2ppt_parallel] all pages finished. cost={cost_time:.2f}s")

        # 整理结果
        new_pagecontent: List[Dict[str, Any]] = []
        state.generated_pages = []

        for i, res in enumerate(results):
            if isinstance(res, Exception):
                # 理论上 _process_single_page 内部捕获了大部分异常，这里是防漏
                log.error(f"[paper2ppt_parallel] page {i} unhandled exception: {res}")
                # 构造一个失败项
                failed_item = dict(page_items[i]) if isinstance(page_items[i], dict) else {"raw_item": str(page_items[i])}
                failed_item.update({
                    "generated_img_path": None,
                    "page_idx": i,
                    "mode": "unhandled_exception",
                    "error": str(res)
                })
                new_pagecontent.append(failed_item)
            else:
                # res 是 dict
                res_dict = res  # type: ignore
                new_pagecontent.append(res_dict)
                gen_path = res_dict.get("generated_img_path")

                if gen_path:
                    state.generated_pages.append(gen_path)
                else:
                    # 占位，防止索引错位
                    state.generated_pages.append("")

        state.pagecontent = new_pagecontent
        return state

    async def edit_single_page(state: Paper2FigureState) -> Paper2FigureState:
        """
        gen_down == True 时的路径：
        通过 edit_page_num(0-based) + edit_page_prompt 对已经生成好的某一页做二次编辑。

        当前策略（B1）：
        - 不再生成 *_edit_*.png 新文件；
        - 直接覆盖原来的 page_{idx:03d}.png，保证导出时每页只有一张图。
        """
        idx = int(getattr(state, "edit_page_num", -1))
        prompt = (getattr(state, "edit_page_prompt", "") or "").strip()
        if idx < 0:
            raise ValueError("[paper2ppt] edit_page_num 必须是 0-based 且 >=0")
        
        # 取出原图路径：优先 generated_pages，其次 pagecontent[i].ppt_img_path
        old_path: Optional[str] = None
        
        # Debug log
        log.info(f"[paper2ppt] edit_single_page: idx={idx}")
        # log.info(f"[paper2ppt] generated_pages={getattr(state, 'generated_pages', None)}")
        
        if getattr(state, "generated_pages", None) and idx < len(state.generated_pages):
            old_path = state.generated_pages[idx]
            log.info(f"[paper2ppt] got old_path from generated_pages: {old_path}")
        
        if not old_path and idx < len(state.pagecontent or []):
            it = state.pagecontent[idx]
            if isinstance(it, dict):
                old_path = it.get("generated_img_path") or it.get("ppt_img_path") or it.get("img_path")
                log.info(f"[paper2ppt] got old_path from pagecontent: {old_path}")
        
        if not old_path:
            raise ValueError(f"[paper2ppt] 找不到要编辑的页图路径: idx={idx}")

        old_path = _abs_path(old_path)

        result_root = Path(_ensure_result_path(state))
        img_dir = result_root / "ppt_pages"
        img_dir.mkdir(parents=True, exist_ok=True)

        # B1 策略：编辑时直接覆盖原始 page_{idx:03d}.png，避免 *_edit_*.png 累积
        save_path = str((img_dir / f"page_{idx:03d}.png").resolve())
        aspect_ratio = getattr(state, "aspect_ratio", None) or "16:9"
        style = getattr(state.request, "style", None) or "kartoon"

        # 强化提示词，确保模型进行重绘
        if prompt:
            # 用户提供了具体修改意见
            full_prompt = (
                f"Beautify this PowerPoint slide based on this instruction: '{prompt}'. "
                f"Transform the existing design into a high-end, professional {style} style presentation. "
                f"Enhance the visual aesthetics, layout, and background while preserving the core message."
            )
        else:
            # 用户未提供具体修改意见，仅仅请求重新生成/美化
            full_prompt = (
                f"Beautify and re-design this PowerPoint slide. "
                f"Transform the existing design into a high-end, professional {style} style presentation. "
                f"Enhance the visual aesthetics, layout, and background while preserving the core message."
            )

        log.info(f"[paper2ppt] edit_single_page idx={idx} old={old_path} save={save_path}")
        log.critical(f'[full_prompt] {full_prompt}')

        await generate_or_edit_and_save_image_async(
            prompt=full_prompt,
            save_path=save_path,
            aspect_ratio=aspect_ratio,
            api_url=state.request.chat_api_url,
            api_key=state.request.chat_api_key or os.getenv("DF_API_KEY") ,
            model=state.request.gen_fig_model,
            image_path=old_path,
            use_edit=True,
        )

        # 回写路径
        if getattr(state, "generated_pages", None) and idx < len(state.generated_pages):
            state.generated_pages[idx] = save_path
        if idx < len(state.pagecontent or []):
            it = state.pagecontent[idx]
            if isinstance(it, dict):
                it["generated_img_path"] = save_path
                it["edit_prompt"] = prompt
                it["mode"] = "edit_again"

        # 清理编辑请求（可选）
        state.edit_page_prompt = ""
        state.edit_page_num = -1
        return state

    async def export_ppt_assets(state: Paper2FigureState) -> Paper2FigureState:
        """
        最终导出节点：
        - 使用 ppt_tool.convert_images_dir_to_pdf_and_ppt_api（带 API inpainting 支持）
          将 result_path/ppt_pages 中的页面图导出为 PDF 和可编辑 PPTX。

        注意：
        - gen_down == False（首次生成）：始终导出；
        - gen_down == True（编辑模式）：只有在 request.all_edited_down == True 时才导出，
          否则直接跳过该节点。
        """
        # 若处于编辑模式且未标记全部编辑完成，则跳过导出
        if getattr(state, "gen_down", False):
            all_done = getattr(getattr(state, "request", None), "all_edited_down", False)
            if not all_done:
                log.info("[paper2ppt] export_ppt_assets skipped: gen_down=True & all_edited_down is False")
                return state

        result_root = Path(_ensure_result_path(state))
        img_dir = result_root / "ppt_pages"

        if not img_dir.exists():
            raise ValueError(f"[paper2ppt] export_ppt_assets: image dir not found: {img_dir}")

        pdf_path = result_root / "paper2ppt.pdf"
        pptx_path = result_root / "paper2ppt_editable.pptx"

        log.info(
            f"[paper2ppt] export_ppt_assets: images_dir={img_dir}, "
            f"pdf={pdf_path}, pptx={pptx_path}"
        )

        # 使用新的 API 版本函数（带 inpainting 支持）
        out = await convert_images_dir_to_pdf_and_ppt_api(
            input_dir=str(img_dir),
            output_pdf_path=str(pdf_path),
            output_pptx_path=None,
            api_url=state.request.chat_api_url,
            api_key=state.request.chat_api_key or os.getenv("DF_API_KEY") ,
            model=state.request.gen_fig_model,
            use_api_inpaint=False,  # 启用 API inpainting
        )

        # 可选：把导出结果路径挂到 state 上，方便后续使用
        setattr(state, "ppt_pdf_path", out.get("pdf") or str(pdf_path))
        setattr(state, "ppt_pptx_path", None)

        return state

    nodes = {
        "_start_": _start_,
        "generate_pages": generate_pages,
        "edit_single_page": edit_single_page,
        "export_ppt_assets": export_ppt_assets,
        "_end_": lambda state: state,
    }

    edges = [
        ("generate_pages", "export_ppt_assets"),
        ("edit_single_page", "export_ppt_assets"),
        ("export_ppt_assets", "_end_"),
    ]

    builder.add_nodes(nodes).add_edges(edges).add_conditional_edge("_start_", _route)
    return builder