Blame for processor/app/routes/export.py

"""Export APIs — 本地文件打包(无数据库)。

导出任务存 data/state/export_tasks/<job_id>.json;打包 = 批次目录 zip。
"""

from __future__ import annotations

import hashlib
import shutil
import tempfile
import zipfile
import re
import json
from datetime import datetime, timezone
from pathlib import Path
from uuid import uuid4

from fastapi import (APIRouter, BackgroundTasks, Depends, HTTPException, Query,
                     Request)
from fastapi.responses import FileResponse
from starlette.background import BackgroundTask

from app.config import settings
from app.localstore import (
    STATE_ROOT, scan_sessions, get_episode, list_deleted_episodes, list_runs,
    get_workflow, list_projects, _read_json, _write_json, _remove_json,
)
from app.security import verify_api_key
from app.workflow_types import HAND_PROCESS_TYPES, RGBD_3D_TYPES, canonical_node_type
from app.project_dataset import (
    episode_chunk_for_index,
    episode_files,
    episode_row,
    is_project_dataset,
)

router = APIRouter(prefix="/api/v1/export", tags=["export"])

EXPORTS_DIR = STATE_ROOT / "exports"
EXPORT_TASKS_DIR = STATE_ROOT / "export_tasks"
# 可训练的数据集 —— 与 EXPORTS_DIR 分开:那边是按 run_id 命名的**处理产物**
# (每条一集),这边是**能拿去训练的成品**,训练页只列这个目录。
DATASETS_DIR = STATE_ROOT / "datasets"
# 导出页/批量导出的构建目录与 zip 放本地盘:同一份视频数据不再在
# NAS(SSHFS ~28MB/s)上读写三遍,只保留一次源文件读取。
LOCAL_EXPORTS_DIR = settings.upload_staging_root / "exports"


def _utcnow() -> str:
    return datetime.now(timezone.utc).isoformat()


def _tmp_root() -> Path:
    path = STATE_ROOT / "tmp"
    path.mkdir(parents=True, exist_ok=True)
    return path


def _safe_download_stem(value: object, fallback: str = "egodata") -> str:
    """Keep project names readable in downloaded filenames.

    The old exporter used an ASCII-only whitelist, which silently removed
    Chinese project names and produced names such as ``D435-__3D_AI.zip``.
    Preserve Unicode while still preventing path separators and control
    characters from escaping the filename boundary.
    """
    name = str(value or "").strip()
    name = re.sub(r"[\x00-\x1f\x7f]+", "_", name)
    name = name.replace("/", "_").replace("\\", "_")
    name = re.sub(r"\s+", " ", name).strip(" .")
    return name or fallback


def _zip_batches(batch_dirs: list[Path], dest: Path) -> None:
    with zipfile.ZipFile(dest, "w", allowZip64=True) as archive:
        for batch in batch_dirs:
            for path in batch.rglob("*"):
                if not path.is_file():
                    continue
                rel = path.relative_to(batch)
                if not rel.parts or rel.parts[0] not in {"data", "meta", "videos"}:
                    continue
                arcname = Path(batch.name) / rel
                compression = zipfile.ZIP_STORED if path.suffix.lower() in {".mp4", ".parquet", ".npz"} else zipfile.ZIP_DEFLATED
                archive.write(path, str(arcname), compress_type=compression)


def _zip_directory(root: Path, dest: Path) -> None:
    with zipfile.ZipFile(dest, "w", allowZip64=True) as archive:
        for path in root.rglob("*"):
            if path.is_file():
                archive.write(path, str(path.relative_to(root)),
                              compress_type=zipfile.ZIP_STORED if path.suffix.lower() in {".mp4", ".parquet"} else zipfile.ZIP_DEFLATED)


# 旧版本导出代码遗留的非标准边车文件:标注已在 parquet 列里
# (annotation/annotation_index/annotation_scope),设备信息在 info.json
# (robot_type)。打包时剥离,不随 zip 分发。
_SIDECAR_EXCLUDE = {"meta/annotations.jsonl", "meta/devices.json"}


def _zip_lerobot(root: Path, dest: Path) -> None:
    """打包 LeRobot 标准结构(data/ + meta/ + videos/),不携带处理副产品
    (hand_3d/skeleton/calibration 等)。连接驱动导出带视频 feature 的
    数据集必须带上 videos/,否则查看器/官方加载器在归档里找不到视频
    文件(File not found in archive)。"""
    with zipfile.ZipFile(dest, "w", allowZip64=True) as archive:
        for sub in ("data", "meta", "videos"):
            sub_dir = root / sub
            if not sub_dir.is_dir():
                continue
            for path in sub_dir.rglob("*"):
                if not path.is_file():
                    continue
                if str(path.relative_to(root)) in _SIDECAR_EXCLUDE:
                    continue
                archive.write(path, str(path.relative_to(root)),
                              compress_type=zipfile.ZIP_STORED if path.suffix.lower() in {".mp4", ".parquet"} else zipfile.ZIP_DEFLATED)


def _cleanup_later(path: str) -> None:
    try:
        Path(path).unlink(missing_ok=True)
    except OSError:
        pass


def _episode_dir(episode_id: str) -> Path:
    ep = get_episode(episode_id)
    if ep is None:
        raise HTTPException(status_code=404, detail=f"Episode not found: {episode_id}")
    return Path(ep["path"])


def _require_exportable_episode(episode_id: str) -> dict:
    """Keep Reviewing batches out of every download path, not only the UI."""
    ep = get_episode(episode_id)
    if ep is None:
        raise HTTPException(status_code=404, detail=f"Episode not found: {episode_id}")
    if ep.get("status") not in ("reviewed", "approved"):
        raise HTTPException(status_code=409, detail="Episode is still in Reviewing")
    return ep


def _build_zip(batches: list[Path], prefix: str, filename: str) -> FileResponse:
    tmp = tempfile.NamedTemporaryFile(prefix=prefix, suffix=".zip",
                                      dir=str(_tmp_root()), delete=False)
    tmp.close()
    try:
        _zip_batches(batches, Path(tmp.name))
    except Exception:
        _cleanup_later(tmp.name)
        raise HTTPException(status_code=500, detail="Failed to package episodes")
    return FileResponse(tmp.name, media_type="application/zip", filename=filename,
                        background=BackgroundTask(_cleanup_later, tmp.name))


def _zip_episode_selection(episode_ids: list[str], dest: Path) -> None:
    """Zip exactly the requested episodes from canonical project datasets.

    A project root can now contain hundreds of episodes.  Packaging the root
    directly would silently export every episode when the user requested one,
    and would also include unrelated episodes in a batch download.
    """
    groups: dict[Path, list[tuple[str, dict, dict]]] = {}
    for episode_id in episode_ids:
        ep = get_episode(str(episode_id))
        if ep is None:
            continue
        root = Path(ep["path"])
        if not is_project_dataset(root):
            continue
        row = episode_row(root, str(episode_id))
        if row is None:
            continue
        index = int(row.get("episode_index"))
        groups.setdefault(root, []).append(
            (str(episode_id), row, episode_files(root, index)))

    with zipfile.ZipFile(dest, "w", allowZip64=True) as archive:
        written: set[str] = set()

        def add_file(path: Path, arcname: str) -> None:
            if not path.is_file() or arcname in written:
                return
            compression = (zipfile.ZIP_STORED
                           if path.suffix.lower() in {".mp4", ".parquet", ".npz"}
                           else zipfile.ZIP_DEFLATED)
            archive.write(path, arcname, compress_type=compression)
            written.add(arcname)

        for root, entries in groups.items():
            selected_rows = [row for _episode_id, row, _files in entries]
            selected_indexes = {int(row.get("episode_index")) for row in selected_rows}
            for _episode_id, _row, files in entries:
                for path in files.get("data", []):
                    add_file(path, str(path.relative_to(root)))
                for _source, path in files.get("videos", []):
                    add_file(path, str(path.relative_to(root)))
                for path in files.get("meta", []):
                    meta_path = path if path.is_absolute() else root / path
                    rel = str(meta_path.relative_to(root))
                    if (rel not in {"meta/info.json", "meta/stats.json",
                                    "meta/tasks.json"}
                            and not rel.startswith("meta/episodes/")):
                        add_file(meta_path, rel)

            info = root / "meta" / "info.json"
            tasks = root / "meta" / "tasks.json"
            stats = root / "meta" / "stats.json"
            add_file(info, "meta/info.json")
            add_file(tasks, "meta/tasks.json")
            if selected_rows:
                import pyarrow as pa
                import pyarrow.parquet as pq
                rows_by_chunk: dict[int, list[dict]] = {}
                for row in selected_rows:
                    index = int(row.get("episode_index", 0))
                    rows_by_chunk.setdefault(episode_chunk_for_index(index), []).append(row)
                for chunk_index, chunk_rows in sorted(rows_by_chunk.items()):
                    for row in chunk_rows:
                        index = int(row.get("episode_index", 0))
                        sink = pa.BufferOutputStream()
                        pq.write_table(pa.Table.from_pylist([row]), sink)
                        archive.writestr(
                            f"meta/episodes/chunk-{chunk_index:03d}/"
                            f"episode_{index:06d}.parquet",
                            sink.getvalue().to_pybytes(),
                        )
            add_file(stats, "meta/stats.json")

        # Keep compatibility for pre-migration sessions.


def _build_episode_zip(episode_ids: list[str], prefix: str,
                       filename: str) -> FileResponse:
    tmp = tempfile.NamedTemporaryFile(prefix=prefix, suffix=".zip",
                                      dir=str(_tmp_root()), delete=False)
    tmp.close()
    try:
        _zip_episode_selection([str(value) for value in episode_ids], Path(tmp.name))
    except Exception:
        _cleanup_later(tmp.name)
        raise HTTPException(status_code=500, detail="Failed to package episodes")
    return FileResponse(tmp.name, media_type="application/zip", filename=filename,
                        background=BackgroundTask(_cleanup_later, tmp.name))


# ── 按工作流自动匹配导出格式 ─────────────────────────

def _episode_export_target(episode_id: str) -> dict | None:
    """Return the latest published export product for one episode.

    Export products live in the system-level export cache.  The project
    dataset itself remains only ``data/``, ``meta/`` and ``videos/``.
    """
    ep = get_episode(episode_id)
    if ep is None:
        return None
    runs = [r for r in list_runs() if r.get("episode_id") == episode_id]
    runs.sort(key=lambda r: r.get("created_at") or "", reverse=True)
    for run in runs:
        if run.get("status") != "completed":
            continue
        export_node = next(
            (n for n in (run.get("graph") or {}).get("nodes", [])
             if str((n.get("data") or {}).get("nodeType") or "").lower()
                in ("lerobot_export", "hdf5_export")),
            None,
        )
        if export_node is None:
            continue
        node_type = str((export_node.get("data") or {}).get("nodeType") or "").lower()
        cfg = (export_node.get("data") or {}).get("config") or {}

        # New runs publish export products to the system-level export cache
        # and store the published path in the run snapshot.  The source
        # project still contains only data/meta/videos.
        run_outputs = run.get("outputs") or {}
        published = run_outputs.get("artifacts", run_outputs)
        if isinstance(published, dict):
            for handles in published.values():
                if not isinstance(handles, dict):
                    continue
                for ref in handles.values():
                    if not isinstance(ref, dict):
                        continue
                    candidate = Path(str(ref.get("path") or ""))
                    if node_type == "hdf5_export":
                        if candidate.is_file() and candidate.suffix.lower() in {".h5", ".hdf5"}:
                            return {"kind": "hdf5", "version": None,
                                    "product": candidate, "root": candidate.parent}
                    elif (candidate / "meta" / "info.json").is_file():
                        # 版本以产物自身的 codebase_version 为准(实际
                        # 构建版本),不能信任工作流配置 —— 旧产物可能是
                        # v2.1 代码按配置 v3.0 构建前的遗留,文件名/徽标
                        # 会与实际布局不符。
                        product_version = str(cfg.get("version") or "v3.0")
                        try:
                            import json as _json
                            info = _json.loads(
                                (candidate / "meta" / "info.json")
                                .read_text(encoding="utf-8"))
                            built = str(info.get("codebase_version") or "")
                            if built.lower().startswith("v2"):
                                product_version = "v2.1"
                            elif built.lower().startswith("v3"):
                                product_version = "v3.0"
                        except (OSError, ValueError):
                            pass
                        return {"kind": "lerobot",
                                "version": product_version,
                                "product": candidate, "root": candidate}

    return None


def _workflow_export_version(episode_id: str) -> str | None:
    """Resolve the LeRobot version from the episode's bound workflow.

    Review-page exports intentionally do not have a separate format selector:
    the workflow's ``lerobot_export.version`` is the source of truth.  A
    completed run is used only as a fallback for older projects whose current
    binding no longer contains the export node.
    """
    episode = get_episode(str(episode_id))
    if episode is None:
        return None
    project = next((item for item in list_projects()
                    if item.get("name") == episode.get("project")), None)
    workflow_ids: list[str] = []
    if project:
        raw_ids = project.get("workflow_ids") or []
        if not isinstance(raw_ids, list):
            raw_ids = [raw_ids]
        if project.get("workflow_id"):
            raw_ids = [project.get("workflow_id"), *raw_ids]
        workflow_ids = [str(value) for value in raw_ids if str(value).strip()]
    for workflow_id in dict.fromkeys(workflow_ids):
        workflow = get_workflow(workflow_id)
        if not workflow:
            continue
        for node in (workflow.get("graph") or {}).get("nodes") or []:
            data = node.get("data") or {}
            if str(data.get("nodeType") or "").lower() != "lerobot_export":
                continue
            config = dict(data.get("config") or {})
            config.update((workflow.get("node_configs") or {}).get(
                node.get("id"), {}))
            version = str(config.get("version") or "v3.0").strip().lower()
            if version.startswith("v2"):
                return "v2.1"
            if version.startswith("v3"):
                return "v3.0"

    # Compatibility fallback: an older completed run still records the exact
    # export-node configuration used for that episode.
    target = _episode_export_target(str(episode_id))
    if target and target.get("kind") == "lerobot":
        version = str(target.get("version") or "").strip().lower()
        if version.startswith("v2"):
            return "v2.1"
        if version.startswith("v3"):
            return "v3.0"
    return None


@router.get("/episode-format/{episode_id}")
def episode_export_format(episode_id: str):
    """前端按钮徽标:该批次下载时实际返回的格式。

    同步端点(FastAPI 丢线程池执行)。以前写成 ``async def`` 却调
    :func:`_episode_export_target` —— 那里读历史 run 快照 + 逐个 ``is_file()``
    全是远程阻塞 I/O,在事件循环上跑会把同一次开页的其它请求一起冻住:
    实测该接口冷调用 4.4s,期间 hand-3d / slam / tactile 全部排队,这正是
    "打开一条数据时而 2 秒时而 7 秒"的来源。
    """
    if get_episode(episode_id) is None:
        raise HTTPException(status_code=404, detail="Episode not found")
    target = _episode_export_target(episode_id)
    if target is None:
        return {"episode_id": episode_id, "format": "raw", "version": None,
                "label": "Raw", "available": False}
    if target["kind"] == "hdf5":
        fmt, label = "hdf5", "HDF5"
    else:
        version = target["version"] or "v3.0"
        fmt = "lerobot_v2" if version.lower().startswith("v2") else "lerobot_v3"
        label = f"LeRobot {version}"
    return {"episode_id": episode_id, "format": fmt,
            "version": target["version"], "label": label, "available": True}


@router.get("/download-episode/{episode_id}")
async def download_single_episode(episode_id: str):
    """下载 = 按工作流导出格式返回最终数据集;无产物回退原始批次 zip。"""
    episode = _require_exportable_episode(episode_id)
    project_name = _safe_download_stem(episode.get("project"), episode_id)
    target = _episode_export_target(episode_id)
    if target is not None:
        if target["kind"] == "hdf5":
            return FileResponse(str(target["product"]),
                                media_type="application/x-hdf5",
                                filename=f"{project_name}.h5")
        tmp = tempfile.NamedTemporaryFile(prefix=f"egodata-{episode_id}-",
                                          suffix=".zip", dir=_tmp_root(),
                                          delete=False)
        tmp.close()
        _zip_lerobot(target["root"], Path(tmp.name))
        version = target["version"] or "v3.0"
        return FileResponse(
            tmp.name, media_type="application/zip",
            filename=f"{project_name}-{version}.zip",
            background=BackgroundTask(_cleanup_later, tmp.name))
    return _build_episode_zip([episode_id], f"egodata-{episode_id}-", f"{project_name}.zip")


@router.get("/download-reviewed")
async def download_reviewed_sessions():
    eps = [e for e in scan_sessions() if e.get("status") in ("reviewed", "approved")]
    deleted = {d["id"] for d in list_deleted_episodes()}
    eps = [e for e in eps if e["id"] not in deleted]
    if not eps:
        raise HTTPException(status_code=404, detail="No reviewed episodes")
    return _build_episode_zip([str(e["id"]) for e in eps], "egodata-reviewed-", "reviewed.zip")


async def _episode_ids_from_request(request: Request) -> list[str]:
    """从 JSON body 或表单字段取 episode_ids(两种入参都支持)。

    表单路径的值是逗号分隔的单字段 —— 浏览器原生 form POST 只能提交
    表单编码,而它带来的流式下载正是我们想要的。
    """
    values: list[str] = []
    content_type = request.headers.get("content-type", "")
    try:
        if "application/json" in content_type:
            body = await request.json()
            raw = (body or {}).get("episode_ids") or []
            values = raw if isinstance(raw, list) else [raw]
        else:
            form = await request.form()
            raw = form.get("episode_ids") or ""
            values = str(raw).split(",")
    except Exception:
        return []
    return [str(value).strip() for value in values if str(value).strip()]


@router.post("/batch-download")
async def batch_download_sessions(request: Request):
    """打包下载所选批次 —— 从 canon 数据集里裁出选中的集,不重建。

    产出仍是**单个** LeRobot 标准结构(data/ + meta/ + videos/),与单集
    ``download-episode`` 走的是同一个 ``_zip_episode_selection``,区别只在
    传入的 id 数量。与「导出数据集」(走 /start 的整库重建,约 30 秒/集)
    的差异不在目录结构,而在内容:这里文件原样拷贝、episode 编号保留原值
    (故可能不连续)、meta 整份拷贝自源数据集;重建则重新编号 0..N-1 并
    按选中集重算 stats。

    秒级返回,前端用两个按钮区分快慢。

    入参两种都收:JSON ``{"episode_ids": [...]}`` 供程序调用;表单字段
    ``episode_ids``(逗号分隔)供浏览器原生表单提交 —— 后者让浏览器自己
    流式下载,不必把整个 zip 缓冲进内存再由 JS 触发。

    鉴权:与 download-episode / download/{job_id} / download-reviewed 一致,
    只靠 AuthMiddleware 的登录 cookie。此前额外的 Depends(verify_api_key)
    会让浏览器 401(它拿不到 X-API-Key 头),而这三个兄弟端点都没有该依赖。
    """
    ids = await _episode_ids_from_request(request)
    if not ids:
        raise HTTPException(status_code=400, detail="No episode_ids provided")
    episode_ids: list[str] = []
    project_names: set[str] = set()
    for ep_id in ids:
        episode = _require_exportable_episode(str(ep_id))
        episode_ids.append(str(ep_id))
        project = _safe_download_stem(episode.get("project"), "")
        if project:
            project_names.add(project)
    if not episode_ids:
        raise HTTPException(status_code=404, detail="No valid episodes")
    filename = f"{next(iter(project_names))}.zip" if len(project_names) == 1 else "egodata-batch.zip"
    return _build_episode_zip(episode_ids, "egodata-batch-", filename)


# ── 导出任务(JSON,后台打包)─────────────────────────

def _persist_trainable_dataset(export_root: Path, dataset_name: str) -> Path | None:
    """把导出好的数据集落到 ``state/datasets/<名字>/``,让它能当训练输入。

    原来是打包完直接 ``rmtree``(省本地空间)—— 结果导出的成果只剩一个 zip,
    磁盘上什么都没有,训练页根本无从列起。

    **交换式替换**,不是"先删旧的再搬新的":先把新数据集搬进 ``.staging``,
    再把旧的挪到 ``.old``,然后 rename 上位,最后删 ``.old``。任何一步失败都不会
    让目标位置处于"半个数据集"的状态 —— 直接 rmtree 再 move 的话,搬到一半
    失败就两头落空。

    返回落位后的路径;``export_root`` 里没有 ``meta/info.json``(不是数据集)时
    返回 None,调用方照旧删掉临时目录。
    """
    if not (export_root / "meta" / "info.json").is_file():
        return None

    DATASETS_DIR.mkdir(parents=True, exist_ok=True)
    target = DATASETS_DIR / dataset_name
    staging = DATASETS_DIR / f".{dataset_name}.staging"
    previous = DATASETS_DIR / f".{dataset_name}.old"

    shutil.rmtree(staging, ignore_errors=True)
    shutil.move(str(export_root), str(staging))
    shutil.rmtree(previous, ignore_errors=True)
    if target.exists():
        target.rename(previous)
    staging.rename(target)          # 同盘 rename,原子
    shutil.rmtree(previous, ignore_errors=True)
    return target


def _load_task(job_id: str) -> dict | None:
    return _read_json(EXPORT_TASKS_DIR / f"{job_id}.json", None)


def _save_task(task: dict) -> None:
    _write_json(EXPORT_TASKS_DIR / f"{task['id']}.json", task)


def _list_tasks() -> list[dict]:
    if not EXPORT_TASKS_DIR.is_dir():
        return []
    return [t for t in (_read_json(f, {}) for f in sorted(EXPORT_TASKS_DIR.glob("*.json"))) if t]


# running 任务超过该时长仍未结束 → 线程大概率已异常死亡,标记 failed
# (正常 3 批次导出实测 ~3 分钟;30 分钟是宽裕兜底)
_EXPORT_RUNNING_TIMEOUT_SECONDS = 30 * 60


def _mark_stale_running_failed(task: dict) -> bool:
    """把卡死在 running 的超时任务标记为 failed(返回是否发生了修改)。"""
    if task.get("status") != "running":
        return False
    try:
        created = datetime.fromisoformat(
            str(task.get("created_at") or "").replace("Z", "+00:00"))
        elapsed = (datetime.now(timezone.utc) - created).total_seconds()
        if elapsed > _EXPORT_RUNNING_TIMEOUT_SECONDS:
            task["status"] = "failed"
            task["error"] = ("Export timed out (running for over "
                             "30 minutes without finishing)")
            task["finished_at"] = _utcnow()
            _save_task(task)
            return True
    except (ValueError, TypeError):
        pass
    return False


def _episode_set_fingerprint(episode_ids: list[str]) -> str:
    """所选批次源文件的指纹(size+mtime):数据变过才需要重新构建。"""
    digest = hashlib.sha256()
    for episode_id in episode_ids:
        ep = get_episode(str(episode_id))
        if ep is None:
            continue
        root = Path(ep["path"])
        if not is_project_dataset(root):
            continue
        try:
            row = episode_row(root, str(episode_id))
        except Exception:
            continue
        if row is None:
            continue
        index = int(row.get("episode_index"))
        files = episode_files(root, index)
        candidates = [*files.get("data", []),
                      *(path for _, path in files.get("videos", []))]
        for path in candidates:
            try:
                st = Path(path).stat()
                digest.update(
                    f"{Path(path).name}|{st.st_size}|{st.st_mtime_ns}"
                    .encode("utf-8"))
            except OSError:
                continue
    return digest.hexdigest()


def _run_export_task(job_id: str, episode_ids: list[str], export_format: str,
                     dataset_name: str, split_ratio: float) -> None:
    def _set_progress(progress: float, **extra) -> None:
        """后台任务阶段进度(前端按钮实时显示,不再静默等待)。"""
        task = _load_task(job_id)
        if not task:
            return
        task["progress"] = max(0.0, min(1.0, float(progress)))
        task.update(extra)
        _save_task(task)

    try:
        blocked = [
            str(episode_id) for episode_id in episode_ids
            if (get_episode(str(episode_id)) or {}).get("status")
            not in ("reviewed", "approved")
        ]
        if blocked:
            raise RuntimeError("Some episodes are still in Reviewing")
        # 构建与 zip 都在本地盘进行(NAS 只读一次源文件)。产物两样都留:
        # zip 供下载,数据集目录(LeRobot 分支)供训练页选用 —— 后者走
        # _persist_trainable_dataset 落到 state/datasets/,不再是"打完包就删"。
        dest = LOCAL_EXPORTS_DIR / f"{job_id}.zip"
        dest.parent.mkdir(parents=True, exist_ok=True)
        export_root = LOCAL_EXPORTS_DIR / f"{dataset_name}-{job_id[:8]}"
        if export_format in {"lerobot_v2", "lerobot_v3"}:
            from app.lerobot_export import build_lerobot_dataset
            version = "v2.1" if export_format == "lerobot_v2" else "v3.0"

            def _on_episode(done: int, total_episodes: int) -> None:
                _set_progress(0.05 + 0.65 * (done / max(1, total_episodes)))

            _set_progress(0.05)
            build_lerobot_dataset(dataset_name, episode_ids, export_root,
                                  split_ratio, version=version,
                                  progress_callback=_on_episode)
            _set_progress(0.75)
            _zip_directory(export_root, dest)
            # 打包后**留下**数据集,不要删 —— 训练页的数据集列表就是从这里读的
            # (state/datasets/)。只有它压根不是数据集时才清掉临时目录。
            if _persist_trainable_dataset(export_root, dataset_name) is None:
                shutil.rmtree(export_root, ignore_errors=True)
            _set_progress(0.95)
        else:
            _set_progress(0.2)
            _zip_episode_selection([str(e) for e in episode_ids], dest)
            _set_progress(0.95)
            # 非 LeRobot 分支产出的只是一个选集的 zip,没有数据集目录可留。
            shutil.rmtree(export_root, ignore_errors=True)
        task = _load_task(job_id)
        if task:
            task["status"] = "completed"
            task["progress"] = 1.0
            task["output_dir"] = str(dest)
            task["finished_at"] = _utcnow()
            _save_task(task)
    except Exception as e:
        task = _load_task(job_id)
        if task:
            task["status"] = "failed"
            task["error"] = str(e)
            task["finished_at"] = _utcnow()
            _save_task(task)


@router.post("/start", status_code=202)
async def export_start(body: dict, background_tasks: BackgroundTasks,
                       _: str = Depends(verify_api_key)):
    requested_ids = body.get("episode_ids")
    if requested_ids:
        if not isinstance(requested_ids, list):
            raise HTTPException(status_code=400, detail="episode_ids must be a list")
        episode_ids = [str(episode_id) for episode_id in requested_ids]
        blocked = [
            episode_id for episode_id in episode_ids
            if (get_episode(episode_id) or {}).get("status")
            not in ("reviewed", "approved")
        ]
        if blocked:
            raise HTTPException(status_code=409,
                                detail="Some episodes are still in Reviewing")
    else:
        episode_ids = [
            e["id"] for e in scan_sessions()
            if e.get("status") in ("reviewed", "approved")
        ]
    if not episode_ids:
        raise HTTPException(status_code=404, detail="No reviewed episodes")
    requested_name = str(body.get("dataset_name") or "").strip()
    selected_projects = {
        _safe_download_stem((get_episode(episode_id) or {}).get("project"), "")
        for episode_id in episode_ids
    }
    selected_projects.discard("")
    # The review page does not choose a separate dataset name.  A single
    # project batch should download with that project's full Unicode name;
    # mixed-project selections keep a neutral batch name.
    if not requested_name or requested_name in {"egodata-batch", "egodata-export", "dataset"}:
        requested_name = (next(iter(selected_projects))
                          if len(selected_projects) == 1 else "egodata-batch")
    dataset_name = _safe_download_stem(requested_name, "egodata-export")
    requested_format = str(body.get("export_format") or "").strip().lower()
    if requested_format:
        export_format = requested_format
    else:
        # Review-page exports have no independent format selector. Resolve
        # each selected episode from its project's workflow export node.
        workflow_versions = {
            _workflow_export_version(episode_id)
            for episode_id in episode_ids
        }
        workflow_versions.discard(None)
        if len(workflow_versions) > 1:
            raise HTTPException(
                status_code=409,
                detail=("Selected episodes use different workflow export "
                        "versions; select episodes from one workflow/project."),
            )
        resolved_version = next(iter(workflow_versions), "v3.0")
        export_format = ("lerobot_v2" if resolved_version.startswith("v2")
                         else "lerobot_v3")
    if export_format not in {"lerobot_v2", "lerobot_v3"}:
        raise HTTPException(status_code=400,
                            detail="Unsupported export format")
    # 防重复提交:同一批次集合已有 running 任务时直接拒绝,避免并行
    # 任务互相抢 NAS 带宽(此前实测两次点击各跑一遍,时间翻倍)。
    requested = set(episode_ids)
    for existing in _list_tasks():
        if existing.get("status") != "running":
            continue
        if set(existing.get("episode_ids") or []) == requested:
            raise HTTPException(
                status_code=409,
                detail="An export for the same episodes is already running",
            )
    # 复用已完成结果:同批次(顺序一致)+ 同格式 + 源文件未变 → 直接
    # 返回既有任务,前端轮询立即 completed 并触发下载(秒级)。
    fingerprint = _episode_set_fingerprint(episode_ids)
    for existing in _list_tasks():
        if existing.get("status") != "completed":
            continue
        if (existing.get("episode_ids") or []) != episode_ids:
            continue
        if existing.get("export_format") != export_format:
            continue
        if existing.get("source_fingerprint") != fingerprint:
            continue
        output = Path(str(existing.get("output_dir") or ""))
        if output.is_file():
            return existing
    job = {
        "id": str(uuid4()),
        "dataset_name": dataset_name,
        "status": "running",
        "episode_ids": episode_ids,
        "split_ratio": body.get("split_ratio", 0.9),
        "progress": 0.0,
        "export_format": export_format,
        "error": None,
        "output_dir": None,
        "source_fingerprint": fingerprint,
        "created_at": _utcnow(),
        "finished_at": None,
    }
    _save_task(job)
    background_tasks.add_task(
        _run_export_task,
        job["id"],
        episode_ids,
        job["export_format"],
        job["dataset_name"],
        float(job["split_ratio"] or 0.9),
    )
    return job


@router.get("/list")
async def export_list(limit: int = Query(20, ge=1, le=100), offset: int = 0):
    tasks = sorted(_list_tasks(), key=lambda t: t.get("created_at") or "", reverse=True)
    for task in tasks:
        _mark_stale_running_failed(task)
    return tasks[offset:offset + limit]


@router.get("/{job_id}")
async def export_status(job_id: str):
    task = _load_task(job_id)
    if task is None:
        raise HTTPException(status_code=404, detail="Export job not found")
    _mark_stale_running_failed(task)
    return task


@router.get("/download/{job_id}")
async def download_export_job(job_id: str):
    task = _load_task(job_id)
    if task is None:
        raise HTTPException(status_code=404, detail="Export job not found")
    # 新任务 zip 在本地盘(output_dir);历史任务回退 NAS 路径。
    dest = Path(str(task.get("output_dir") or ""))
    if not dest.is_file():
        dest = EXPORTS_DIR / f"{job_id}.zip"
    if not dest.is_file():
        raise HTTPException(status_code=404, detail="Export output not ready")
    # 文件名带格式版本后缀,下载后一眼可辨 v2.1 / v3.0
    export_format = str(task.get("export_format") or "")
    suffix = ("-v2.1" if export_format == "lerobot_v2"
              else "-v3.0" if export_format == "lerobot_v3" else "")
    return FileResponse(
        dest, media_type="application/zip",
        filename=f"{task.get('dataset_name', 'dataset')}{suffix}.zip")


@router.delete("/{job_id}")
async def delete_export_job(job_id: str, _: str = Depends(verify_api_key)):
    task = _load_task(job_id)
    _remove_json(EXPORT_TASKS_DIR / f"{job_id}.json")
    for path in (Path(str((task or {}).get("output_dir") or "")),
                 EXPORTS_DIR / f"{job_id}.zip"):
        try:
            path.unlink(missing_ok=True)
        except Exception:
            pass
    return {"message": "Export job deleted"}


# ── Re-export:只重建导出,不重跑检测 ─────────────────────────
# 工作流改导出节点(配置/连线)后,本接口按"最新 completed run 的 canonical
# episode 数据 + 当前工作流的导出配置/连线"重建系统级导出缓存,并回写 run 快照
# 的导出节点配置
# (下载徽标/格式检测按快照驱动,回写后自动显示新格式)。

import asyncio

from app.localstore import save_run

_RE_TASKS: dict[str, dict] = {}
_RE_TASKS_DIR = STATE_ROOT / "re_export_tasks"

_PASSTHROUGH_TYPES = {
    # data_quality 由 ai_quality_review + data_cleaning 合并而来；此处的 t 已过
    # canonical_node_type()，旧 slug 不会出现。
    "annotation", "human_review", "ai_annotation", "data_quality",
}
_CAMERA_TYPES = {"rgb_camera", "fisheye_camera", "rgbd_camera", "stereo_camera",
                 "stereo_rgbd_camera", "mono_camera"}


def _re_set_task(task_id: str, **kw) -> None:
    t = _RE_TASKS.get(task_id)
    if t is None:
        return
    t.update(kw)
    try:
        ep = t.get("episode_id")
        if ep:
            _RE_TASKS_DIR.mkdir(parents=True, exist_ok=True)
            _write_json(_RE_TASKS_DIR / f"{ep}.json", t)
    except OSError:
        pass


def purge_episode_re_export_tasks(episode_id: str) -> None:
    """Cancel re-export jobs for a permanently deleted episode.

    Pop first so a still-running job's final _re_set_task call no-ops and
    cannot recreate the on-disk mirror; then remove the mirror file.
    """
    episode_id = str(episode_id)
    for task_id, task in list(_RE_TASKS.items()):
        if str(task.get("episode_id")) == episode_id:
            _RE_TASKS.pop(task_id, None)
    try:
        (_RE_TASKS_DIR / f"{episode_id}.json").unlink(missing_ok=True)
    except OSError:
        pass


def _export_spec_from_graph(graph: dict, node_configs: dict) -> dict | None:
    """从工作流图静态推导导出节点的输入范围(与 DAG 执行同语义):
    - video_keys:相机输入节点沿纯透传链(annotation/human_review/
      ai_annotation)连到导出节点 → 其 source_keys。处理节点
      (mediapipe 等)不透传视频,链式场景视频不进数据集(需求一致)
    - kinds:mediapipe_hand + RGB-only processors → hand_keypoints;
      RGB-D processors → hand_3d (+hand_3d#right for bare-hand multi-view);
      透传节点继续回溯
    图里没有 lerobot_export 节点返回 None。
    """
    nodes = {n.get("id"): n for n in graph.get("nodes") or []}
    edges = graph.get("edges") or []
    export_id = next((nid for nid, n in nodes.items()
                      if (n.get("data") or {}).get("nodeType") == "lerobot_export"),
                     None)
    if export_id is None:
        return None
    video_keys: list[str] = []
    kinds: set[str] = set()
    stack = [e.get("source") for e in edges if e.get("target") == export_id]
    visited: set[str] = set()
    while stack:
        nid = stack.pop()
        if not nid or nid in visited:
            continue
        visited.add(nid)
        node = nodes.get(nid)
        if node is None:
            continue
        t = canonical_node_type((node.get("data") or {}).get("nodeType") or "")
        cfg = dict((node.get("data") or {}).get("config") or {})
        cfg.update((node_configs or {}).get(nid, {}))
        if t in _CAMERA_TYPES:
            for k in (cfg.get("source_keys") or cfg.get("source_key") or "").split(","):
                k = k.strip()
                if k and k not in video_keys:
                    video_keys.append(k)
        elif t == "mediapipe_hand" or t in HAND_PROCESS_TYPES - RGBD_3D_TYPES:
            kinds.add("hand_keypoints")
        if t in RGBD_3D_TYPES:
            kinds.add("hand_3d")
            if t == "rgbd_to_3d_bare_hand":
                kinds.add("hand_3d#right")
        elif t in _PASSTHROUGH_TYPES:
            stack.extend(e.get("source") for e in edges
                         if e.get("target") == nid)
        # 其他节点类型不贡献导出输入,不回溯
    return {"video_keys": video_keys, "kinds": kinds}


async def _run_re_export(task_id: str, episode_id: str) -> None:
    from app.lerobot_export import build_lerobot_dataset
    try:
        _re_set_task(task_id, status="loading", detail="定位 run 与工作流")
        await asyncio.sleep(0)
        ep = get_episode(episode_id)
        if ep is None:
            raise RuntimeError("Episode not found")
        runs = [r for r in list_runs()
                if r.get("episode_id") == episode_id
                and r.get("status") == "completed" and r.get("created_at")]
        runs.sort(key=lambda r: r.get("created_at") or "", reverse=True)
        run = runs[0] if runs else None
        if run is None:
            raise RuntimeError("无 completed run 记录(先跑一次工作流再重导出)")
        project_root = Path(ep["path"])
        source_row = episode_row(project_root, episode_id)
        if source_row is None:
            raise RuntimeError("canonical episode metadata 不存在")
        source_index = int(source_row.get("episode_index", 0))
        source_data = (project_root / "data"
                       / f"chunk-{episode_chunk_for_index(source_index):03d}"
                       / f"episode_{source_index:06d}.parquet")
        if not source_data.is_file():
            raise RuntimeError("canonical episode data parquet 不存在")
        # 当前工作流:优先项目**当前绑定**的工作流(用户改的连线/配置
        # 就生效在那里);run 的来源工作流仅作兜底(run 可能出自已被
        # 解绑的旧工作流,若优先它会拿到过期的连线/源键)。
        wf = None
        project = next((p for p in list_projects()
                        if p.get("name") == ep.get("project")), None)
        for wf_id in (project.get("workflow_ids") or []) if project else []:
            w = get_workflow(wf_id)
            if w is not None and _export_spec_from_graph(
                    w.get("graph") or {}, w.get("node_configs") or {}) is not None:
                wf = w
                break
        if wf is None:
            old_wf = get_workflow(run.get("workflow_id") or "")
            if old_wf is not None and _export_spec_from_graph(
                    old_wf.get("graph") or {},
                    old_wf.get("node_configs") or {}) is not None:
                wf = old_wf
        if wf is None:
            raise RuntimeError("项目绑定工作流里没有 LeRobot 导出节点")
        spec = _export_spec_from_graph(wf.get("graph") or {},
                                       wf.get("node_configs") or {})
        if spec is None:
            raise RuntimeError("当前工作流没有 LeRobot 导出节点")
        export_node = next(
            n for n in (wf.get("graph") or {}).get("nodes", [])
            if (n.get("data") or {}).get("nodeType") == "lerobot_export")
        cfg = dict((export_node.get("data") or {}).get("config") or {})
        cfg.update((wf.get("node_configs") or {}).get(export_node.get("id"), {}))
        version = str(cfg.get("version") or "v3.0")
        split_ratio = float(cfg.get("split_ratio", 0.9))
        # 检测结果已经合并进 canonical episode parquet,重导出不重跑检测。
        kinds = spec["kinds"]
        hand_keypoints_paths = ([str(source_data)]
                                if "hand_keypoints" in kinds else None)
        hand_3d_paths = ([str(source_data)] if "hand_3d" in kinds else None)
        hand_3d_right_paths = ([str(source_data)]
                               if "hand_3d#right" in kinds else None)
        include_video_keys = spec["video_keys"] or None  # None=全量(旧语义),[]=无视频
        _re_set_task(task_id, status="building",
                     detail=f"重建 LeRobot {version} 数据集中(复用检测产物,不重跑)")
        await asyncio.sleep(0)
        # 导出副本位于系统级缓存,不污染项目根目录。
        import shutil
        run_root = EXPORTS_DIR / str(run["id"]) / "dataset"
        shutil.rmtree(run_root, ignore_errors=True)
        run_root.parent.mkdir(parents=True, exist_ok=True)
        await asyncio.to_thread(
            build_lerobot_dataset, episode_id, [episode_id], run_root,
            split_ratio, include_video_keys,
            hand_keypoints_paths, hand_3d_paths, hand_3d_right_paths,
            version, None)
        # 回写 run 快照的导出节点配置 → 下载徽标/格式检测显示新配置。
        # 旧 run 快照可能没有导出节点(早期工作流无导出)→ 从当前工作流
        # 把导出节点(带新配置)补进快照图,保证 _episode_export_target
        # 能读到格式。
        patched = False
        run_graph = run.setdefault("graph", {})
        for n in (run_graph.get("nodes") or []):
            if (n.get("data") or {}).get("nodeType") == "lerobot_export":
                n.setdefault("data", {})["config"] = dict(
                    n.get("data", {}).get("config") or {})
                n["data"]["config"].update({
                    "version": version, "split_ratio": split_ratio})
                if (run.get("node_configs") or {}).get(n.get("id")) is not None:
                    run["node_configs"][n.get("id")].update({
                        "version": version, "split_ratio": split_ratio})
                patched = True
                break
        if not patched:
            node_copy = dict(export_node)
            node_copy["id"] = f"re_export_{node_copy.get('id')}"
            node_copy["data"] = dict(node_copy.get("data") or {})
            node_copy["data"]["config"] = dict(cfg)
            run_graph.setdefault("nodes", [])
            run_graph["nodes"].append(node_copy)
            export_output_node_id = node_copy["id"]
        else:
            export_output_node_id = export_node.get("id")
        run_outputs = run.setdefault("outputs", {})
        artifact_map = run_outputs.setdefault("artifacts", {})
        artifact_map.setdefault(export_output_node_id, {})["dataset"] = {
            "kind": "dataset",
            "path": str(run_root),
            "source_key": None,
            "schema_version": "1.0",
            "metadata": {"root": ".", "version": version},
        }
        save_run(run)
        _re_set_task(task_id, status="done",
                     detail=f"重建完成:LeRobot {version}"
                            + (f",视频:{len(spec['video_keys'])} 路(按连线)"
                               if spec["video_keys"]
                               else ",视频:自动识别全量")
                            + (f",骨骼列 {sorted(kinds)}" if kinds else ""))
    except Exception as exc:
        _re_set_task(task_id, status="failed", detail=str(exc))


@router.post("/re-export/{episode_id}")
async def re_export_episode(episode_id: str):
    """只重建导出:复用最新 completed run 的检测产物 + 当前工作流的
    导出节点配置/连线,原地重建 LeRobot 数据集并回写 run 快照。
    不重跑 mediapipe 等检测节点。"""
    if get_episode(episode_id) is None:
        raise HTTPException(status_code=404, detail="Episode not found")
    task_id = str(uuid4())
    _RE_TASKS[task_id] = {"task_id": task_id, "episode_id": episode_id,
                          "status": "queued", "detail": ""}
    asyncio.get_running_loop().create_task(_run_re_export(task_id, episode_id))
    return {"task_id": task_id, "status": "queued"}


@router.get("/re-export/{episode_id}/status")
async def re_export_status(episode_id: str):
    for t in reversed(list(_RE_TASKS.values())):
        if t.get("episode_id") == episode_id:
            return t
    disk = _read_json(_RE_TASKS_DIR / f"{episode_id}.json", None)
    if disk is None:
        return {"episode_id": episode_id, "status": "idle"}
    if disk.get("status") in ("queued", "loading", "building"):
        disk["status"] = "interrupted"
        disk["detail"] = str(disk.get("detail") or "") + "(服务重启,任务中断,请重新触发)"
    return disk