#!/usr/bin/env python3
# coding: utf-8

import argparse
import hashlib
import json
import os
import shutil
import subprocess
import tarfile
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from threading import Lock
from time import perf_counter

from count_visualization_batches import (
    DEFAULT_IMAGE_EXTS,
    DEFAULT_SHOWPATH_ROOT,
    collect_image_stats,
    normalize_image_exts,
)


DEFAULT_BATCH_GLOB = "batch_*"
DEFAULT_ARCHIVE_FORMAT = "tar.gz"
DEFAULT_ARCHIVE_BACKEND = "auto"
DEFAULT_MANIFEST_NAME = "archives_manifest.json"
DEFAULT_FRAMES_PER_PART = 2000
MODALITIES = ("2D", "3D")
LOG_LOCK = Lock()


@dataclass(frozen=True)
class FrameEntry:
    batch_name: str
    relpath: str
    path_2d: Path
    path_3d: Path


def log(message=""):
    with LOG_LOCK:
        print(message, flush=True)


def default_scan_workers():
    return max(1, min(8, os.cpu_count() or 1))


def default_archive_workers():
    return len(MODALITIES)


def default_part_workers():
    return 1


def parse_args():
    parser = argparse.ArgumentParser(
        description="汇总所有 batch 的 2D/3D 可视化结果，并按固定帧数分片打包。"
    )
    parser.add_argument(
        "target_path",
        nargs="?",
        default=DEFAULT_SHOWPATH_ROOT,
        help="可视化 batch 根目录，或单个 batch 目录。",
    )
    parser.add_argument(
        "output_root",
        nargs="?",
        default=None,
        help="归档输出目录。默认输出到 <batch_root> 的并列目录 <name>_archives。",
    )
    parser.add_argument(
        "--batch-glob",
        default=DEFAULT_BATCH_GLOB,
        help="batch 根目录模式下用于查找 batch 的 glob，默认 batch_*。",
    )
    parser.add_argument(
        "--archive-format",
        choices=["tar", "tar.gz"],
        default=DEFAULT_ARCHIVE_FORMAT,
        help="归档格式，默认 tar.gz。",
    )
    parser.add_argument(
        "--frames-per-part",
        type=int,
        default=DEFAULT_FRAMES_PER_PART,
        help="每个 part 包含多少帧，默认 2000。",
    )
    parser.add_argument(
        "--scan-workers",
        type=int,
        default=default_scan_workers(),
        help="扫描 batch 的并发线程数，默认按机器核数自适应，最多 8。",
    )
    parser.add_argument(
        "--archive-workers",
        type=int,
        default=default_archive_workers(),
        help="单个 part 内归档并发数，默认 2，同时处理 2D 和 3D。",
    )
    parser.add_argument(
        "--part-workers",
        type=int,
        default=default_part_workers(),
        help="同时处理多少个 part，默认 1；每个 part 内仍会按 archive-workers 并发处理 2D/3D。",
    )
    parser.add_argument(
        "--archive-backend",
        choices=["auto", "python", "system_tar"],
        default=DEFAULT_ARCHIVE_BACKEND,
        help="归档实现后端，默认 auto；优先使用系统 tar/gzip，以提升大批量打包速度。",
    )
    parser.add_argument(
        "--gzip-compresslevel",
        type=int,
        default=1,
        help="tar.gz 压缩级别，范围 0-9，默认 1 以优先速度。",
    )
    parser.add_argument(
        "--image-exts",
        default=DEFAULT_IMAGE_EXTS,
        help="参与校验统计的图片扩展名，逗号分隔，默认 .jpg,.jpeg,.png。",
    )
    parser.add_argument(
        "--check-pairs",
        action="store_true",
        default=True,
        help="打包前校验 2D/3D 相对路径是否一一对应，默认开启。",
    )
    parser.add_argument(
        "--no-check-pairs",
        dest="check_pairs",
        action="store_false",
        help="关闭 2D/3D 一一对应校验，仅分别收集文件。",
    )
    parser.add_argument(
        "--sample-mismatch-limit",
        type=int,
        default=10,
        help="2D/3D 配对校验失败时，manifest 中最多保留多少条样例。",
    )
    parser.add_argument(
        "--skip-existing",
        action="store_true",
        help="若目标 part 归档已存在且大小大于 0，则跳过该 part。",
    )
    parser.add_argument(
        "--checksum",
        action="store_true",
        help="归档完成后计算 sha256。",
    )
    parser.add_argument(
        "--verify-archive",
        action="store_true",
        default=True,
        help="归档完成后快速校验非空且可读取，默认开启。",
    )
    parser.add_argument(
        "--no-verify-archive",
        dest="verify_archive",
        action="store_false",
        help="关闭归档完成后的快速校验，以进一步提速。",
    )
    parser.add_argument(
        "--remove-source-after-success",
        action="store_true",
        help="所有归档全部成功后删除源 batch 目录。默认关闭，请谨慎使用。",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="只做校验和计划输出，不实际创建归档或删除源目录。",
    )
    return parser.parse_args()


def now_str():
    return datetime.now().isoformat(timespec="seconds")


def is_relative_to(path, other):
    try:
        path.relative_to(other)
        return True
    except ValueError:
        return False


def is_batch_dir(path):
    return path.is_dir() and ((path / "2D").is_dir() or (path / "3D").is_dir())


def resolve_batch_dirs(target_path, batch_glob):
    if is_batch_dir(target_path):
        return target_path.parent, [target_path], "single_batch"

    if not target_path.is_dir():
        raise FileNotFoundError(f"target path does not exist or is not a directory: {target_path}")

    batch_dirs = sorted(path for path in target_path.glob(batch_glob) if is_batch_dir(path))
    if not batch_dirs:
        raise FileNotFoundError(
            f"no batch directories found under {target_path} with batch_glob={batch_glob}"
        )
    return target_path, batch_dirs, "batch_root"


def default_output_root(batch_root):
    return batch_root.parent / f"{batch_root.name}_archives"


def archive_suffix(archive_format):
    return "tar.gz" if archive_format == "tar.gz" else "tar"


def archive_basename(archive_path):
    name = archive_path.name
    if name.endswith(".tar.gz"):
        return name[: -len(".tar.gz")]
    if name.endswith(".tar"):
        return name[: -len(".tar")]
    return archive_path.stem


def build_archive_path(output_root, modality, part_index, archive_format):
    return output_root / modality / f"part_{part_index:04d}.{archive_suffix(archive_format)}"


def write_json(path, data):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)


def ensure_safe_output_root(root_dir, output_root):
    batch_dir_resolved = root_dir.resolve()
    output_root_resolved = output_root.resolve()
    if is_relative_to(output_root_resolved, batch_dir_resolved):
        raise ValueError(
            "output_root 不允许位于源目录内部，"
            f"source_root={batch_dir_resolved}, output_root={output_root_resolved}"
        )


def supports_system_tar_backend(archive_format, gzip_compresslevel):
    if shutil.which("tar") is None:
        return False
    if archive_format != "tar.gz":
        return True
    if gzip_compresslevel == 0:
        return False
    return shutil.which("gzip") is not None


def resolve_archive_backend(requested_backend, archive_format, gzip_compresslevel):
    system_tar_supported = supports_system_tar_backend(
        archive_format,
        gzip_compresslevel,
    )
    if requested_backend == "auto":
        return "system_tar" if system_tar_supported else "python"

    if requested_backend == "system_tar":
        if shutil.which("tar") is None:
            raise RuntimeError("archive-backend=system_tar 需要系统 tar 命令。")
        if archive_format == "tar.gz" and shutil.which("gzip") is None:
            raise RuntimeError("archive-backend=system_tar 在 tar.gz 模式下需要 gzip 命令。")
        if archive_format == "tar.gz" and gzip_compresslevel == 0:
            raise ValueError(
                "archive-backend=system_tar 当前不支持 gzip-compresslevel=0，"
                "请改用 python 后端或将压缩级别设为 1-9。"
            )
        return "system_tar"

    return "python"


def open_tarfile(archive_path, archive_format, gzip_compresslevel):
    mode = "w:gz" if archive_format == "tar.gz" else "w"
    if archive_format == "tar.gz":
        return tarfile.open(archive_path, mode=mode, compresslevel=gzip_compresslevel)
    return tarfile.open(archive_path, mode=mode)


def build_archive_member_dir(archive_path):
    return f"{archive_basename(archive_path)}/images"


def build_archive_entry_name(archive_path, source_path):
    return f"{build_archive_member_dir(archive_path)}/{source_path.name}"


def ensure_unique_archive_names(source_paths, archive_path):
    seen_names = {}
    for source_path in source_paths:
        archive_name = build_archive_entry_name(archive_path, source_path)
        if archive_name in seen_names:
            raise ValueError(
                "archive entry name duplicated: "
                f"{archive_name}, first={seen_names[archive_name]}, duplicate={source_path}"
            )
        seen_names[archive_name] = source_path


def build_relative_source_paths(source_paths, source_root):
    source_root = source_root.resolve()
    relative_paths = []
    for source_path in source_paths:
        try:
            relative_path = source_path.relative_to(source_root)
        except ValueError as exc:
            raise ValueError(
                f"source path is outside source_root: source_root={source_root}, path={source_path}"
            ) from exc
        relative_paths.append(relative_path.as_posix())
    return relative_paths


def decode_stderr(stderr_bytes):
    if not stderr_bytes:
        return ""
    return stderr_bytes.decode("utf-8", errors="replace").strip()


def escape_tar_transform_replacement(value):
    return value.replace("\\", "\\\\").replace("&", r"\&").replace("|", r"\|")


def create_archive_with_python_tar(
    source_paths,
    archive_path,
    archive_format,
    gzip_compresslevel,
):
    archive_path.parent.mkdir(parents=True, exist_ok=True)
    partial_archive = archive_path.with_name(f"{archive_path.name}.partial")
    if partial_archive.exists():
        partial_archive.unlink()

    ensure_unique_archive_names(source_paths, archive_path)
    try:
        with open_tarfile(partial_archive, archive_format, gzip_compresslevel) as tar_obj:
            for source_path in source_paths:
                tar_obj.add(
                    source_path,
                    arcname=build_archive_entry_name(archive_path, source_path),
                    recursive=False,
                )

        partial_archive.replace(archive_path)
    finally:
        if partial_archive.exists():
            partial_archive.unlink()


def create_archive_with_system_tar(
    source_paths,
    archive_path,
    source_root,
    archive_format,
    gzip_compresslevel,
):
    archive_path.parent.mkdir(parents=True, exist_ok=True)
    partial_archive = archive_path.with_name(f"{archive_path.name}.partial")
    if partial_archive.exists():
        partial_archive.unlink()

    ensure_unique_archive_names(source_paths, archive_path)
    relative_paths = build_relative_source_paths(source_paths, source_root)
    relative_paths_input = b"".join(os.fsencode(path) + b"\0" for path in relative_paths)
    archive_member_dir = escape_tar_transform_replacement(build_archive_member_dir(archive_path))
    transform_rule = f"s|.*/||;s|^|{archive_member_dir}/|"

    try:
        if archive_format == "tar":
            tar_cmd = [
                "tar",
                "-cf",
                str(partial_archive),
                "--directory",
                str(source_root),
                "--transform",
                transform_rule,
                "--no-recursion",
                "--null",
                "--files-from",
                "-",
            ]
            completed = subprocess.run(
                tar_cmd,
                input=relative_paths_input,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.PIPE,
                check=False,
            )
            if completed.returncode != 0:
                raise RuntimeError(
                    "system tar failed: "
                    f"{decode_stderr(completed.stderr) or 'unknown error'}"
                )
        else:
            tar_cmd = [
                "tar",
                "-cf",
                "-",
                "--directory",
                str(source_root),
                "--transform",
                transform_rule,
                "--no-recursion",
                "--null",
                "--files-from",
                "-",
            ]
            gzip_cmd = ["gzip", f"-{gzip_compresslevel}", "-c"]
            with open(partial_archive, "wb") as archive_file:
                tar_proc = subprocess.Popen(
                    tar_cmd,
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
                try:
                    gzip_proc = subprocess.Popen(
                        gzip_cmd,
                        stdin=tar_proc.stdout,
                        stdout=archive_file,
                        stderr=subprocess.PIPE,
                    )
                except Exception:
                    tar_proc.kill()
                    tar_proc.wait()
                    raise

                tar_proc.stdout.close()
                _, tar_stderr = tar_proc.communicate(input=relative_paths_input)
                _, gzip_stderr = gzip_proc.communicate()

            if tar_proc.returncode != 0:
                raise RuntimeError(
                    "system tar failed: "
                    f"{decode_stderr(tar_stderr) or 'unknown error'}"
                )
            if gzip_proc.returncode != 0:
                raise RuntimeError(
                    "gzip failed: "
                    f"{decode_stderr(gzip_stderr) or 'unknown error'}"
                )

        partial_archive.replace(archive_path)
    finally:
        if partial_archive.exists():
            partial_archive.unlink()


def create_archive(
    source_paths,
    archive_path,
    source_root,
    archive_backend,
    archive_format,
    gzip_compresslevel,
):
    if archive_backend == "system_tar":
        create_archive_with_system_tar(
            source_paths,
            archive_path,
            source_root,
            archive_format,
            gzip_compresslevel,
        )
        return

    create_archive_with_python_tar(
        source_paths,
        archive_path,
        archive_format,
        gzip_compresslevel,
    )


def verify_archive(archive_path):
    if not archive_path.exists():
        raise FileNotFoundError(f"archive not found: {archive_path}")
    if archive_path.stat().st_size <= 0:
        raise ValueError(f"archive size is 0: {archive_path}")

    with tarfile.open(archive_path, mode="r:*") as tar_obj:
        first_member = tar_obj.next()
        if first_member is None:
            raise ValueError(f"archive is empty: {archive_path}")


def compute_sha256(file_path):
    digest = hashlib.sha256()
    with open(file_path, "rb") as f:
        while True:
            chunk = f.read(1024 * 1024)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()


def part_manifest_path(output_root, part_index):
    return output_root / f"part_{part_index:04d}.manifest.json"


def build_part_manifest(
    args,
    part_index,
    archive_paths,
    entries,
    status,
    error=None,
    archive_sha256=None,
):
    archive_details = {}
    for modality, archive_path in archive_paths.items():
        archive_exists = archive_path.exists()
        archive_details[modality] = {
            "archive_path": str(archive_path.resolve()),
            "archive_exists": archive_exists,
            "archive_size_bytes": archive_path.stat().st_size if archive_exists else 0,
            "archive_sha256": archive_sha256.get(modality) if archive_sha256 else None,
        }

    batch_distribution = Counter(entry.batch_name for entry in entries)
    return {
        "created_at": now_str(),
        "part_name": f"part_{part_index:04d}",
        "part_index": part_index,
        "archive_format": args.archive_format,
        "archive_backend": args.archive_backend_resolved,
        "frames_per_part": args.frames_per_part,
        "frame_count": len(entries),
        "archive_details": archive_details,
        "status": status,
        "dry_run": args.dry_run,
        "check_pairs": args.check_pairs,
        "verify_archive": args.verify_archive,
        "batch_distribution": dict(sorted(batch_distribution.items())),
        "frame_samples": [entry.relpath for entry in entries[: min(5, len(entries))]],
        "error": error,
    }


def summarize_results(results):
    status_counter = Counter(item["status"] for item in results)
    return {
        "total_parts": len(results),
        "status_counts": dict(status_counter),
    }


def collect_batch_relpaths(batch_dir, image_exts):
    modality_relpaths = {}
    modality_counts = {}
    for modality in MODALITIES:
        image_dir = batch_dir / modality
        count, relpaths = collect_image_stats(image_dir, image_exts, collect_relpaths=True)
        modality_counts[modality] = count
        modality_relpaths[modality] = relpaths or set()
    return modality_relpaths, modality_counts


def collect_frame_entries(batch_dir, image_exts, check_pairs, sample_mismatch_limit):
    modality_relpaths, modality_counts = collect_batch_relpaths(batch_dir, image_exts)
    dir_2d = batch_dir / "2D"
    dir_3d = batch_dir / "3D"
    relpaths_2d = modality_relpaths["2D"]
    relpaths_3d = modality_relpaths["3D"]

    only_in_2d = sorted(relpaths_2d - relpaths_3d)
    only_in_3d = sorted(relpaths_3d - relpaths_2d)
    paired_relpaths = sorted(relpaths_2d & relpaths_3d)

    if check_pairs and (only_in_2d or only_in_3d):
        raise ValueError(
            "2D/3D 配对校验失败: "
            f"batch={batch_dir.name}, only_in_2d={len(only_in_2d)}, only_in_3d={len(only_in_3d)}, "
            f"only_in_2d_samples={only_in_2d[:sample_mismatch_limit]}, "
            f"only_in_3d_samples={only_in_3d[:sample_mismatch_limit]}"
        )

    if check_pairs:
        selected_relpaths = paired_relpaths
    else:
        selected_relpaths = sorted(relpaths_2d | relpaths_3d)

    if check_pairs:
        entries = [
            FrameEntry(
                batch_name=batch_dir.name,
                relpath=relpath,
                path_2d=dir_2d / relpath,
                path_3d=dir_3d / relpath,
            )
            for relpath in selected_relpaths
        ]
    else:
        entries = []
        for relpath in selected_relpaths:
            path_2d = dir_2d / relpath
            path_3d = dir_3d / relpath
            if not path_2d.is_file() or not path_3d.is_file():
                continue
            entries.append(
                FrameEntry(
                    batch_name=batch_dir.name,
                    relpath=relpath,
                    path_2d=path_2d,
                    path_3d=path_3d,
                )
            )

    batch_summary = {
        "batch_name": batch_dir.name,
        "count_2d": modality_counts["2D"],
        "count_3d": modality_counts["3D"],
        "paired_count": len(paired_relpaths),
        "selected_count": len(entries),
        "only_in_2d_count": len(only_in_2d),
        "only_in_3d_count": len(only_in_3d),
        "only_in_2d_samples": only_in_2d[:sample_mismatch_limit],
        "only_in_3d_samples": only_in_3d[:sample_mismatch_limit],
    }
    return entries, batch_summary


def scan_single_batch(
    batch_index,
    total_batches,
    batch_dir,
    image_exts,
    check_pairs,
    sample_mismatch_limit,
):
    batch_started = perf_counter()
    log(f"[SCAN] ({batch_index}/{total_batches}) {batch_dir.name}")
    batch_entries, batch_summary = collect_frame_entries(
        batch_dir,
        image_exts,
        check_pairs=check_pairs,
        sample_mismatch_limit=sample_mismatch_limit,
    )
    return batch_index, batch_entries, batch_summary, perf_counter() - batch_started


def build_global_frame_index(batch_dirs, image_exts, check_pairs, sample_mismatch_limit, scan_workers):
    entries = []
    batch_summaries = []
    total_batches = len(batch_dirs)
    scan_started = perf_counter()
    cumulative_entries = 0

    if scan_workers <= 1 or total_batches <= 1:
        ordered_entries = []
        ordered_summaries = []
        for batch_index, batch_dir in enumerate(batch_dirs, start=1):
            batch_index, batch_entries, batch_summary, elapsed = scan_single_batch(
                batch_index,
                total_batches,
                batch_dir,
                image_exts,
                check_pairs,
                sample_mismatch_limit,
            )
            ordered_entries.append(batch_entries)
            ordered_summaries.append(batch_summary)
            cumulative_entries += len(batch_entries)
            log(
                f"[SCAN DONE] ({batch_index}/{total_batches}) {batch_dir.name} "
                f"2D={batch_summary['count_2d']} 3D={batch_summary['count_3d']} "
                f"paired={batch_summary['paired_count']} selected={batch_summary['selected_count']} "
                f"elapsed={elapsed:.1f}s cumulative={cumulative_entries}"
            )
    else:
        ordered_entries = [None] * total_batches
        ordered_summaries = [None] * total_batches
        max_workers = min(scan_workers, total_batches)
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            future_map = {
                executor.submit(
                    scan_single_batch,
                    batch_index,
                    total_batches,
                    batch_dir,
                    image_exts,
                    check_pairs,
                    sample_mismatch_limit,
                ): (batch_index, batch_dir)
                for batch_index, batch_dir in enumerate(batch_dirs, start=1)
            }
            for future in as_completed(future_map):
                batch_index, batch_dir = future_map[future]
                _, batch_entries, batch_summary, elapsed = future.result()
                ordered_entries[batch_index - 1] = batch_entries
                ordered_summaries[batch_index - 1] = batch_summary
                cumulative_entries += len(batch_entries)
                log(
                    f"[SCAN DONE] ({batch_index}/{total_batches}) {batch_dir.name} "
                    f"2D={batch_summary['count_2d']} 3D={batch_summary['count_3d']} "
                    f"paired={batch_summary['paired_count']} selected={batch_summary['selected_count']} "
                    f"elapsed={elapsed:.1f}s cumulative={cumulative_entries}"
                )

    for batch_entries in ordered_entries:
        entries.extend(batch_entries)
    batch_summaries.extend(ordered_summaries)

    log(
        f"[SCAN SUMMARY] batches={total_batches} total_frames={len(entries)} "
        f"elapsed={perf_counter() - scan_started:.1f}s"
    )
    return entries, batch_summaries


def count_parts(total_entries, frames_per_part):
    return (total_entries + frames_per_part - 1) // frames_per_part


def iter_entry_chunks(entries, frames_per_part):
    for index in range(0, len(entries), frames_per_part):
        yield entries[index : index + frames_per_part]


def package_modality_archive(modality, source_paths, archive_path, source_root, args):
    create_archive(
        source_paths,
        archive_path,
        source_root,
        args.archive_backend_resolved,
        args.archive_format,
        args.gzip_compresslevel,
    )
    if args.verify_archive:
        verify_archive(archive_path)
    archive_sha256 = compute_sha256(archive_path) if args.checksum else None
    return modality, archive_sha256


def process_part(part_index, entries, source_root, output_root, args):
    archive_paths = {
        modality: build_archive_path(output_root, modality, part_index, args.archive_format)
        for modality in MODALITIES
    }

    part_started = perf_counter()
    header_lines = [
        "",
        "==========================================",
        f"Part      : part_{part_index:04d}",
        f"Frame cnt : {len(entries)}",
    ]
    for modality in MODALITIES:
        header_lines.append(f"{modality} archive: {archive_paths[modality]}")
    header_lines.append("==========================================")
    log("\n".join(header_lines))

    if args.skip_existing and all(
        archive_path.exists() and archive_path.stat().st_size > 0
        for archive_path in archive_paths.values()
    ):
        archive_sha256 = {
            modality: compute_sha256(archive_path) if args.checksum else None
            for modality, archive_path in archive_paths.items()
        }
        manifest = build_part_manifest(
            args,
            part_index,
            archive_paths,
            entries,
            status="skipped_existing",
            archive_sha256=archive_sha256,
        )
        write_json(part_manifest_path(output_root, part_index), manifest)
        log(
            f"[SKIP] archives already exist for part_{part_index:04d} "
            f"elapsed={perf_counter() - part_started:.1f}s"
        )
        return manifest

    if args.dry_run:
        manifest = build_part_manifest(
            args,
            part_index,
            archive_paths,
            entries,
            status="dry_run",
        )
        write_json(part_manifest_path(output_root, part_index), manifest)
        log(
            f"[DRY RUN] would create archives for part_{part_index:04d} "
            f"elapsed={perf_counter() - part_started:.1f}s"
        )
        return manifest

    archive_sha256 = {}
    try:
        source_paths = {
            "2D": [entry.path_2d for entry in entries],
            "3D": [entry.path_3d for entry in entries],
        }
        max_workers = min(max(1, args.archive_workers), len(MODALITIES))
        if max_workers <= 1:
            for modality in MODALITIES:
                _, archive_sha = package_modality_archive(
                    modality,
                    source_paths[modality],
                    archive_paths[modality],
                    source_root,
                    args,
                )
                archive_sha256[modality] = archive_sha
        else:
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                future_map = {
                    executor.submit(
                        package_modality_archive,
                        modality,
                        source_paths[modality],
                        archive_paths[modality],
                        source_root,
                        args,
                    ): modality
                    for modality in MODALITIES
                }
                for future in as_completed(future_map):
                    modality, archive_sha = future.result()
                    archive_sha256[modality] = archive_sha

        manifest = build_part_manifest(
            args,
            part_index,
            archive_paths,
            entries,
            status="packaged",
            archive_sha256=archive_sha256,
        )
        write_json(part_manifest_path(output_root, part_index), manifest)
        log(f"[OK] packaged part_{part_index:04d} elapsed={perf_counter() - part_started:.1f}s")
        return manifest
    except Exception as exc:
        manifest = build_part_manifest(
            args,
            part_index,
            archive_paths,
            entries,
            status="failed",
            error=str(exc),
        )
        write_json(part_manifest_path(output_root, part_index), manifest)
        log(f"[FAIL] part_{part_index:04d}: {exc}")
        return manifest


def package_parts(entries, source_root, output_root, args):
    part_specs = [
        (
            part_index,
            entries[start_index : start_index + args.frames_per_part],
        )
        for part_index, start_index in enumerate(
            range(0, len(entries), args.frames_per_part),
            start=1,
        )
    ]
    total_parts = len(part_specs)
    if args.part_workers <= 1 or total_parts <= 1:
        return [
            process_part(part_index, part_entries, source_root, output_root, args)
            for part_index, part_entries in part_specs
        ]

    ordered_results = [None] * total_parts
    max_workers = min(args.part_workers, total_parts)
    log(f"[PACKAGE] part_workers={max_workers} archive_workers={args.archive_workers}")
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        future_map = {
            executor.submit(
                process_part,
                part_index,
                part_entries,
                source_root,
                output_root,
                args,
            ): part_index
            for part_index, part_entries in part_specs
        }
        for future in as_completed(future_map):
            part_index = future_map[future]
            ordered_results[part_index - 1] = future.result()

    return ordered_results


def build_overall_manifest(
    args,
    target_path,
    batch_root,
    output_root,
    resolve_mode,
    image_exts,
    batch_summaries,
    total_frames,
    results,
):
    return {
        "generated_at": now_str(),
        "target_path": str(target_path.resolve()),
        "batch_root": str(batch_root.resolve()),
        "output_root": str(output_root.resolve()),
        "resolve_mode": resolve_mode,
        "batch_glob": args.batch_glob,
        "archive_format": args.archive_format,
        "archive_backend_requested": args.archive_backend,
        "archive_backend_resolved": args.archive_backend_resolved,
        "frames_per_part": args.frames_per_part,
        "image_exts": list(image_exts),
        "check_pairs": args.check_pairs,
        "sample_mismatch_limit": args.sample_mismatch_limit,
        "part_workers": args.part_workers,
        "archive_workers": args.archive_workers,
        "skip_existing": args.skip_existing,
        "checksum": args.checksum,
        "verify_archive": args.verify_archive,
        "remove_source_after_success": args.remove_source_after_success,
        "dry_run": args.dry_run,
        "total_frames": total_frames,
        "total_batches": len(batch_summaries),
        "batch_summaries": batch_summaries,
        "summary": summarize_results(results),
        "parts": results,
    }


def main():
    try:
        args = parse_args()
        if args.frames_per_part <= 0:
            raise ValueError("frames-per-part 必须大于 0")
        if args.scan_workers <= 0:
            raise ValueError("scan-workers 必须大于 0")
        if args.archive_workers <= 0:
            raise ValueError("archive-workers 必须大于 0")
        if args.part_workers <= 0:
            raise ValueError("part-workers 必须大于 0")
        if not 0 <= args.gzip_compresslevel <= 9:
            raise ValueError("gzip-compresslevel 必须在 0 到 9 之间")

        image_exts = normalize_image_exts(args.image_exts)
        args.archive_backend_resolved = resolve_archive_backend(
            args.archive_backend,
            args.archive_format,
            args.gzip_compresslevel,
        )

        target_path = Path(args.target_path).resolve()
        batch_root, batch_dirs, resolve_mode = resolve_batch_dirs(target_path, args.batch_glob)
        output_root = Path(args.output_root).resolve() if args.output_root else default_output_root(batch_root)
        ensure_safe_output_root(batch_root, output_root)
        output_root.mkdir(parents=True, exist_ok=True)
        for modality in MODALITIES:
            (output_root / modality).mkdir(parents=True, exist_ok=True)

        log("")
        log("######################################################################")
        log("# Package visualization batches")
        log("######################################################################")
        log(f"Mode           : {resolve_mode}")
        log(f"Target         : {target_path}")
        log(f"Batch root     : {batch_root}")
        log(f"Output root    : {output_root}")
        log(f"Archive format : {args.archive_format}")
        log(
            f"Archive backend: {args.archive_backend_resolved} "
            f"(requested={args.archive_backend})"
        )
        log(f"Frames/part    : {args.frames_per_part}")
        log(f"Scan workers   : {args.scan_workers}")
        log(f"Archive workers: {args.archive_workers}")
        log(f"Part workers   : {args.part_workers}")
        log(f"Gzip level     : {args.gzip_compresslevel}")
        log(f"Check pairs    : {args.check_pairs}")
        log(f"Skip existing  : {args.skip_existing}")
        log(f"Checksum       : {args.checksum}")
        log(f"Verify archive : {args.verify_archive}")
        log(f"Remove source  : {args.remove_source_after_success}")
        log(f"Dry run        : {args.dry_run}")
        log(f"Total batches  : {len(batch_dirs)}")

        entries, batch_summaries = build_global_frame_index(
            batch_dirs,
            image_exts,
            check_pairs=args.check_pairs,
            sample_mismatch_limit=args.sample_mismatch_limit,
            scan_workers=args.scan_workers,
        )
        if not entries:
            raise ValueError("未找到可用于打包的 2D/3D 配对帧")

        total_parts = count_parts(len(entries), args.frames_per_part)
        log(f"Total frames   : {len(entries)}")
        log(f"Total parts    : {total_parts}")

        results = package_parts(entries, batch_root, output_root, args)

        overall_manifest = build_overall_manifest(
            args,
            target_path,
            batch_root,
            output_root,
            resolve_mode,
            image_exts,
            batch_summaries,
            len(entries),
            results,
        )
        overall_manifest_path = output_root / DEFAULT_MANIFEST_NAME
        write_json(overall_manifest_path, overall_manifest)

        summary = overall_manifest["summary"]
        log("")
        log(
            "[DONE] total_parts={total} status_counts={status_counts}".format(
                total=summary["total_parts"],
                status_counts=summary["status_counts"],
            )
        )
        log(f"Manifest: {overall_manifest_path}")

        failed_statuses = {"failed"}
        has_failure = any(item["status"] in failed_statuses for item in results)

        if args.remove_source_after_success and not args.dry_run and not has_failure:
            for batch_dir in batch_dirs:
                if batch_dir.exists():
                    shutil.rmtree(batch_dir)

        return 1 if has_failure else 0
    except Exception as exc:
        log(f"[ERROR] {exc}")
        return 1


if __name__ == "__main__":
    raise SystemExit(main())