#!/usr/bin/env python3
"""安全发布一个 FastWAM 真机实验到公开 Hugging Face model 仓库。

设计原则
========

这个脚本只发布命令行明确指定的一个实验，不递归扫描整个训练目录或数据集。
每个实验固定放在 ``<task>/<experiment-name>/`` 下。两个或多个 checkpoint
共享一份 config、数据统计、任务说明和文本 embedding，不按 step 重复 metadata。
脚本永远不上传 DeepSpeed state、W&B 目录、评估目录、视频或 parquet 数据。

远端布局严格为：

::

    <task>/<experiment-name>/
    ├── README.md
    ├── artifact_manifest.json
    ├── config.yaml
    ├── dataset/
    │   ├── meta/
    │   │   ├── stats.json
    │   │   └── tasks.jsonl
    │   └── text_embedding/
    │       ├── manifest.json
    │       └── <manifest 引用的 64-hex 文件名>.pt
    └── weights/
        ├── step_010000.pt
        └── step_020000.pt

其中最大 step 是推荐 checkpoint；同时上传 10000 和 20000 时，20000 是推荐权重，
10000 是备用权重。``artifact_manifest.json`` 记录除自身之外每个公开文件的远端
路径、用途、字节数和 SHA-256。它不记录本机源文件绝对路径。

准备与 dry-run
==============

``config.yaml`` 必须是训练 run 中的 resolved Hydra config，并满足：

* ``output_dir`` 等于 ``--run-dir``；
* ``data.train._target_`` 是 FastWAM DreamZero 数据 builder；
* ``data.train.data_path`` 是数据集根目录；
* ``data.train.text_embed_cache_dir`` 等于 ``data_path/text_embedding``。

默认只做本地校验和打印计划，不访问仓库写接口：

.. code-block:: bash

    python scripts/upload_real_machine_experiment.py \
      --run-dir /path/to/real_machine_close_drawer_baseline_bs48 \
      --steps 10000 20000 \
      --repo-id emiliiia/real_machine_exper \
      --task close_drawer \
      --experiment-name real_machine_close_drawer_baseline_bs48 \
      --revision main

正式公开上传
============

正式写入必须同时提供 ``--upload`` 和与 ``--repo-id`` 完全一致的
``--confirm-public-repo``。脚本会确认目标仓库明确为 public，捕获目标 branch HEAD，
并把该 HEAD 作为 ``parent_commit`` 做一次 compare-and-swap commit：

.. code-block:: bash

    python scripts/upload_real_machine_experiment.py \
      --run-dir /path/to/real_machine_close_drawer_baseline_bs48 \
      --steps 10000 20000 \
      --repo-id emiliiia/real_machine_exper \
      --task close_drawer \
      --experiment-name real_machine_close_drawer_baseline_bs48 \
      --revision main \
      --commit-message "Upload close drawer baseline checkpoints" \
      --upload \
      --confirm-public-repo emiliiia/real_machine_exper

目标 prefix 为空时才创建。若它已经恰好包含同一计划，脚本在固定 HEAD 完整核验后
幂等退出；只要路径有缺失、额外文件或内容不同，就拒绝覆盖。上传只使用显式
``CommitOperationAdd``，从不发出 delete。源文件先在 run 同盘父目录生成私有、
不可变的 reflink snapshot；文件系统不支持 reflink 时才从安全打开的 fd 完整复制。
Hugging Face 只能读取已 unlink 的 snapshot 只读 fd，训练进程随后修改原文件不会改变
上传内容。

下载与服务器路径
================

正式上传完成后，脚本打印不可变 commit oid。下载时应把它作为 ``revision``，不要只
依赖会移动的 ``main``：

.. code-block:: python

    from huggingface_hub import snapshot_download

    snapshot_root = snapshot_download(
        repo_id="emiliiia/real_machine_exper",
        revision="<returned-commit-oid>",
        allow_patterns=[
            "close_drawer/real_machine_close_drawer_baseline_bs48/**"
        ],
    )

下载后，相对 ``snapshot_root`` 的主要路径为：

* ``close_drawer/real_machine_close_drawer_baseline_bs48/weights/step_020000.pt``：推荐权重；
* ``.../weights/step_010000.pt``：备用权重；
* ``.../config.yaml``：resolved Hydra config；
* ``.../dataset/meta/{stats.json,tasks.jsonl}``：共享数据 metadata；
* ``.../dataset/text_embedding/``：必须一起保留 manifest 和全部被引用的 ``.pt``；
* ``.../artifact_manifest.json``：公开文件的完整性清单。

``config.yaml`` 中的 ``output_dir``、数据路径、文本 cache 路径及 checkpoint 根路径可能
是训练服务器绝对路径，下载后必须按新服务器覆盖。模型底座仍通过服务器自己的
``DIFFSYNTH_MODEL_BASE_PATH`` 等环境设置提供。

安全边界与常见失败
==================

* 没有上传 DeepSpeed optimizer/training state，因此不能精确恢复 optimizer 续训；
* manifest、tasks 指令、DROID prompt、embedding 文件名/大小/SHA 不一致时拒绝；
* config、metadata、README 或 commit message 出现 token、密码、私钥等秘密时拒绝；
* 发现 symlink 本地源、``state``/``wandb``/``eval`` 远端路径、视频或 parquet 时拒绝；
* 仓库不是明确 public、公开确认值不匹配、branch HEAD 漂移或 prefix 冲突时拒绝；
* snapshot 空间不足、远端 LFS SHA/文件大小/下载后 SHA 不一致时失败且不会静默通过；
* commit 成功后 branch HEAD 必须仍精确等于返回 oid，否则脚本报告失败，要求人工核查。
"""

from __future__ import annotations

import argparse
import ast
import errno
import fcntl
import hashlib
import json
import os
import re
import stat
import tempfile
from contextlib import ExitStack
from dataclasses import dataclass
from io import BytesIO
from pathlib import Path, PurePosixPath
from typing import Callable, Iterable

import yaml


REPO_TYPE = "model"
FORBIDDEN_COMPONENTS = {"state", "wandb", "eval"}
SAFE_SEGMENT = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
SHA256_RE = re.compile(r"^[0-9a-f]{64}$")
SMALL_FILE_LIMIT = 10 * 1024 * 1024
EXPECTED_DATASET_TARGET = (
    "fastwam.datasets.dreamzero.builder.build_droid_track_dataset"
)
SENSITIVE_KEY_SUFFIXES = (
    "token",
    "password",
    "secret",
    "api_key",
    "apikey",
    "credential",
    "credentials",
    "private_key",
)
SECRET_PATTERNS = (
    ("Hugging Face token", re.compile(r"\bhf_[A-Za-z0-9]{20,}\b")),
    ("OpenAI token", re.compile(r"\bsk-(?:proj-)?[A-Za-z0-9_-]{20,}\b")),
    ("GitHub token", re.compile(r"\b(?:gh[pousr]_[A-Za-z0-9]{20,}|github_pat_[A-Za-z0-9_]{20,})\b")),
    ("AWS access key", re.compile(r"\b(?:AKIA|ASIA)[0-9A-Z]{16}\b")),
    ("PEM private key", re.compile(r"-----BEGIN [A-Z0-9 ]*PRIVATE KEY-----")),
    ("credential URL", re.compile(r"[A-Za-z][A-Za-z0-9+.-]*://[^\s/:@]+:[^\s/@]+@")),
    ("authorization header", re.compile(r"(?i)\bauthorization\s*[:=]\s*(?:bearer\s+)?\S+")),
    ("bearer token", re.compile(r"(?i)\bbearer\s+[A-Za-z0-9._~-]{16,}")),
)
_TRAINING_PROMPT_PREFIX = "A multi-view video shows that a robot "
_MULTIVIEW_LAYOUT = (
    " The video is split into two views stacked vertically: The top view shows "
    "the camera view from a fixed third-person camera that observes the robot "
    "arm and workspace, and the bottom view shows the camera view from the "
    "wrist camera mounted on the gripper, moving together with the arm. "
    "The robot "
)
FICLONE = 0x40049409


class UploadValidationError(ValueError):
    """Raised before or after an upload when an integrity invariant fails."""


@dataclass(frozen=True)
class FileIdentity:
    device: int
    inode: int
    size: int
    mtime_ns: int
    ctime_ns: int


@dataclass(frozen=True)
class Artifact:
    remote_path: str
    role: str
    size: int
    sha256: str
    local_path: Path | None = None
    content: bytes | None = None
    identity: FileIdentity | None = None

@dataclass(frozen=True)
class UploadPlan:
    repo_id: str
    repo_prefix: str
    revision: str
    staging_parent: Path
    steps: tuple[int, ...]
    artifacts: tuple[Artifact, ...]

    @property
    def expected_paths(self) -> set[str]:
        return {artifact.remote_path for artifact in self.artifacts}

    @property
    def total_size(self) -> int:
        return sum(artifact.size for artifact in self.artifacts)


def _sha256_bytes(content: bytes) -> str:
    return hashlib.sha256(content).hexdigest()


def _identity(st: os.stat_result) -> FileIdentity:
    return FileIdentity(
        device=st.st_dev,
        inode=st.st_ino,
        size=st.st_size,
        mtime_ns=st.st_mtime_ns,
        ctime_ns=st.st_ctime_ns,
    )


def _validate_segment(value: str, label: str) -> str:
    if (
        not isinstance(value, str)
        or not SAFE_SEGMENT.fullmatch(value)
        or value.casefold() in FORBIDDEN_COMPONENTS
    ):
        raise UploadValidationError(f"unsafe {label}: {value!r}")
    return value


def _validate_remote_path(path: str) -> str:
    posix = PurePosixPath(path)
    if (
        not path
        or path.startswith("/")
        or "\\" in path
        or posix.as_posix() != path
        or any(part in {"", ".", ".."} for part in posix.parts)
        or any(part.casefold() in FORBIDDEN_COMPONENTS for part in posix.parts)
    ):
        raise UploadValidationError(f"unsafe or forbidden remote path: {path!r}")
    return path


def _assert_directory(path: Path, label: str) -> Path:
    expanded = path.expanduser()
    try:
        st = os.lstat(expanded)
    except OSError as exc:
        raise UploadValidationError(f"{label} is unavailable: {expanded}") from exc
    if stat.S_ISLNK(st.st_mode) or not stat.S_ISDIR(st.st_mode):
        raise UploadValidationError(f"{label} must be a real directory: {expanded}")
    return expanded.resolve()


def _reject_symlink_components(path: Path, stop: Path) -> None:
    current = path
    stop = stop.resolve()
    while True:
        if current.is_symlink():
            raise UploadValidationError(f"symlink source component rejected: {current}")
        if current == stop:
            return
        if stop not in current.resolve().parents:
            raise UploadValidationError(f"source escaped expected root: {path}")
        current = current.parent


def _hash_regular_file(path: Path, *, root: Path | None = None) -> tuple[int, str, FileIdentity]:
    path = path.expanduser()
    if root is not None:
        _reject_symlink_components(path, root)
    try:
        before = os.lstat(path)
    except OSError as exc:
        raise UploadValidationError(f"required source file is unavailable: {path}") from exc
    if stat.S_ISLNK(before.st_mode) or not stat.S_ISREG(before.st_mode):
        raise UploadValidationError(f"source must be a non-symlink regular file: {path}")

    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    fd = os.open(path, flags)
    try:
        opened = os.fstat(fd)
        if not stat.S_ISREG(opened.st_mode) or (
            opened.st_dev,
            opened.st_ino,
        ) != (before.st_dev, before.st_ino):
            raise UploadValidationError(f"source changed while opening: {path}")
        digest = hashlib.sha256()
        while True:
            block = os.read(fd, 1024 * 1024)
            if not block:
                break
            digest.update(block)
        after_fd = os.fstat(fd)
    finally:
        os.close(fd)
    after_path = os.lstat(path)
    expected = _identity(before)
    if _identity(after_fd) != expected or _identity(after_path) != expected:
        raise UploadValidationError(f"source changed while hashing: {path}")
    return expected.size, digest.hexdigest(), expected


def _source_artifact(
    local_path: Path,
    remote_path: str,
    role: str,
    *,
    root: Path | None = None,
) -> Artifact:
    remote_path = _validate_remote_path(remote_path)
    size, sha256, identity = _hash_regular_file(local_path, root=root)
    return Artifact(
        remote_path=remote_path,
        role=role,
        size=size,
        sha256=sha256,
        local_path=local_path.resolve(),
        identity=identity,
    )


def _bytes_artifact(content: bytes, remote_path: str, role: str) -> Artifact:
    _scan_public_text_bytes(content, location=remote_path)
    return Artifact(
        remote_path=_validate_remote_path(remote_path),
        role=role,
        size=len(content),
        sha256=_sha256_bytes(content),
        content=content,
    )


def _is_nonempty(value) -> bool:
    if value is None or value is False:
        return False
    if isinstance(value, (str, bytes, list, tuple, dict, set)):
        return bool(value)
    return True


def _is_sensitive_key(key: object) -> bool:
    if not isinstance(key, str):
        return False
    camel_split = re.sub(r"([a-z0-9])([A-Z])", r"\1_\2", key)
    normalized = re.sub(r"[^a-z0-9]+", "_", camel_split.casefold()).strip("_")
    if normalized.endswith(("_token_id", "_token_ids")):
        return False
    tokens = set(normalized.split("_"))
    return bool(
        tokens & {"token", "password", "secret", "credential", "credentials"}
    ) or any(
        normalized == suffix or normalized.endswith("_" + suffix)
        for suffix in SENSITIVE_KEY_SUFFIXES
    )


def _scan_sensitive_values(value, *, location: str) -> None:
    if isinstance(value, dict):
        for key, child in value.items():
            child_location = f"{location}.{key}"
            if _is_sensitive_key(key) and _is_nonempty(child):
                raise UploadValidationError(
                    f"non-empty sensitive value rejected at {child_location}"
                )
            _scan_sensitive_values(child, location=child_location)
    elif isinstance(value, list):
        for index, child in enumerate(value):
            _scan_sensitive_values(child, location=f"{location}[{index}]")


def _scan_public_text(text: str, *, location: str) -> None:
    for label, pattern in SECRET_PATTERNS:
        if pattern.search(text):
            raise UploadValidationError(f"{label} rejected in public text: {location}")


def _scan_public_text_bytes(content: bytes, *, location: str) -> None:
    try:
        text = content.decode("utf-8")
    except UnicodeDecodeError as exc:
        raise UploadValidationError(
            f"public metadata must be UTF-8 text: {location}"
        ) from exc
    _scan_public_text(text, location=location)


def _load_json(path: Path, label: str):
    try:
        raw = path.read_bytes()
        _scan_public_text_bytes(raw, location=label)
        value = json.loads(raw)
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise UploadValidationError(f"invalid {label}: {path}") from exc
    _scan_sensitive_values(value, location=label)
    return value


def _normalize_instruction_text(raw) -> str:
    if not isinstance(raw, str):
        return str(raw).lower()
    try:
        parsed = ast.literal_eval(raw)
        if isinstance(parsed, (list, tuple)):
            if not parsed:
                raise UploadValidationError("instruction list must not be empty")
            return str(parsed[0]).lower()
        return str(parsed).lower()
    except (ValueError, SyntaxError, TypeError):
        return raw.lower()


def _format_droid_prompt(instruction: str) -> str:
    return _TRAINING_PROMPT_PREFIX + instruction + _MULTIVIEW_LAYOUT + instruction


def _validate_tasks_jsonl(path: Path) -> dict[str, set[str]]:
    tasks: dict[str, set[str]] = {}
    try:
        raw_content = path.read_bytes()
        _scan_public_text_bytes(raw_content, location="dataset tasks")
        lines = raw_content.decode("utf-8").splitlines()
        if not lines:
            raise UploadValidationError(f"tasks JSONL is empty: {path}")
        for line_number, line in enumerate(lines, 1):
            if not line.strip():
                raise UploadValidationError(
                    f"blank line in tasks JSONL at line {line_number}: {path}"
                )
            value = json.loads(line)
            if not isinstance(value, dict):
                raise UploadValidationError(
                    f"tasks JSONL line {line_number} is not an object: {path}"
                )
            _scan_sensitive_values(value, location=f"tasks[{line_number}]")
            task = value.get("task")
            if not isinstance(task, str) or not task.strip():
                raise UploadValidationError(
                    f"tasks JSONL line {line_number} has no task instruction: {path}"
                )
            normalized = _normalize_instruction_text(task)
            if not normalized:
                raise UploadValidationError(
                    f"tasks JSONL line {line_number} normalizes to an empty instruction"
                )
            tasks.setdefault(normalized, set()).add(task)
    except (OSError, UnicodeError, json.JSONDecodeError) as exc:
        raise UploadValidationError(f"invalid tasks JSONL: {path}") from exc
    return tasks


def _load_config(config_path: Path, run_root: Path) -> tuple[dict, Path, Path]:
    try:
        raw_config = config_path.read_bytes()
        _scan_public_text_bytes(raw_config, location="resolved config")
        config = yaml.safe_load(raw_config)
        train = config["data"]["train"]
        if train.get("_target_") != EXPECTED_DATASET_TARGET:
            raise UploadValidationError(
                f"data.train._target_ must be {EXPECTED_DATASET_TARGET!r}"
            )
        data_value = train["data_path"]
        text_value = train["text_embed_cache_dir"]
        output_value = config["output_dir"]
        if not all(isinstance(value, str) for value in (data_value, text_value, output_value)):
            raise UploadValidationError(
                "resolved data_path, text_embed_cache_dir, and output_dir must be scalar strings"
            )
        data_path = Path(data_value).expanduser()
        text_path = Path(text_value).expanduser()
        output_path = Path(output_value).expanduser()
    except (OSError, UnicodeError, yaml.YAMLError, KeyError, TypeError) as exc:
        raise UploadValidationError(
            "config.yaml must contain output_dir, data.train.data_path, "
            "data.train.text_embed_cache_dir, and the expected dataset target"
        ) from exc
    if not isinstance(config, dict):
        raise UploadValidationError("resolved config must be a YAML mapping")
    _scan_sensitive_values(config, location="config")
    if not data_path.is_absolute() or not text_path.is_absolute() or not output_path.is_absolute():
        raise UploadValidationError("resolved config paths must be absolute")
    if output_path.resolve() != run_root:
        raise UploadValidationError("config output_dir must equal --run-dir")
    data_root = _assert_directory(data_path, "data.train.data_path")
    text_root = _assert_directory(text_path, "data.train.text_embed_cache_dir")
    if text_root != data_root / "text_embedding":
        raise UploadValidationError(
            "text_embed_cache_dir must equal data_path/text_embedding"
        )
    return config, data_root, text_root


def _manifest_references(
    manifest_path: Path,
    *,
    expected_model_id: str,
    expected_tokenizer_model_id: str,
    task_instructions: dict[str, set[str]],
) -> list[tuple[str, int, str]]:
    manifest = _load_json(manifest_path, "text embedding manifest")
    entries = manifest.get("entries") if isinstance(manifest, dict) else None
    required_top = {
        "schema_version",
        "embodiment_tag",
        "model_id",
        "tokenizer_model_id",
        "entries",
    }
    if (
        not isinstance(manifest, dict)
        or set(manifest) != required_top
        or manifest.get("schema_version") != 1
        or not isinstance(manifest.get("embodiment_tag"), str)
        or not manifest["embodiment_tag"]
        or manifest.get("model_id") != expected_model_id
        or manifest.get("tokenizer_model_id") != expected_tokenizer_model_id
        or not isinstance(entries, list)
        or not entries
    ):
        raise UploadValidationError("invalid text embedding manifest schema/identity")
    references: dict[str, tuple[int, str]] = {}
    manifest_tasks = set()
    required_entry = {
        "prompt",
        "raw_instruction",
        "normalized_instruction",
        "path",
        "file_size",
        "file_sha256",
        "model_id",
        "tokenizer_model_id",
        "embodiment_tag",
        "encoder_identity",
        "tokenizer_identity",
        "pad_token_id",
    }
    for index, entry in enumerate(entries):
        if not isinstance(entry, dict) or set(entry) != required_entry:
            raise UploadValidationError(f"manifest entry {index} is not an object")
        rel = entry.get("path")
        size = entry.get("file_size")
        sha256 = entry.get("file_sha256")
        raw_instruction = entry.get("raw_instruction")
        normalized_instruction = entry.get("normalized_instruction")
        expected_normalized = (
            _normalize_instruction_text(raw_instruction)
            if isinstance(raw_instruction, str)
            else None
        )
        expected_prompt = (
            _format_droid_prompt(expected_normalized)
            if expected_normalized is not None
            else None
        )
        if (
            not isinstance(raw_instruction, str)
            or not raw_instruction.strip()
            or not isinstance(normalized_instruction, str)
            or normalized_instruction != expected_normalized
            or entry.get("prompt") != expected_prompt
            or raw_instruction not in task_instructions.get(normalized_instruction, set())
            or entry.get("model_id") != expected_model_id
            or entry.get("tokenizer_model_id") != expected_tokenizer_model_id
            or entry.get("embodiment_tag") != manifest["embodiment_tag"]
            or not isinstance(entry.get("encoder_identity"), dict)
            or not entry["encoder_identity"]
            or not isinstance(entry.get("tokenizer_identity"), dict)
            or not entry["tokenizer_identity"]
            or not isinstance(entry.get("pad_token_id"), int)
            or isinstance(entry["pad_token_id"], bool)
        ):
            raise UploadValidationError(f"manifest entry {index} identity mismatch")
        manifest_tasks.add(normalized_instruction)
        if not isinstance(rel, str):
            raise UploadValidationError(f"manifest entry {index} has no path")
        _validate_remote_path(rel)
        if (
            PurePosixPath(rel).name != rel
            or rel != hashlib.sha256(expected_prompt.encode("utf-8")).hexdigest() + ".pt"
            or rel in references
        ):
            raise UploadValidationError(f"duplicate/reserved manifest path: {rel!r}")
        if not isinstance(size, int) or isinstance(size, bool) or size < 0:
            raise UploadValidationError(f"manifest entry {index} has invalid file_size")
        if not isinstance(sha256, str) or not SHA256_RE.fullmatch(sha256):
            raise UploadValidationError(f"manifest entry {index} has invalid file_sha256")
        references[rel] = (size, sha256)
    if manifest_tasks != set(task_instructions):
        raise UploadValidationError(
            "tasks.jsonl instructions and text manifest instructions differ"
        )
    return [(path, *references[path]) for path in sorted(references)]


def _readme_bytes(
    repo_id: str,
    repo_prefix: str,
    revision: str,
    steps: tuple[int, ...],
    artifact_paths: Iterable[str],
) -> bytes:
    recommended = max(steps)
    weights = "\n".join(
        f"- `weights/step_{step:06d}.pt`: "
        + ("recommended checkpoint." if step == recommended else "backup checkpoint.")
        for step in steps
    )
    paths = "\n".join(f"- `{path}`" for path in sorted(artifact_paths))
    text = f"""# Real-machine experiment: {repo_prefix}

Public repository: https://huggingface.co/{repo_id}

Experiment directory: https://huggingface.co/{repo_id}/tree/{revision}/{repo_prefix}

## Checkpoints

{weights}

The largest requested step (`step_{recommended:06d}.pt`) is recommended. Earlier steps are backups. DeepSpeed state is intentionally not included, so these files cannot restore an exact optimizer-state training resume.

## Files

{paths}

`config.yaml` is the resolved Hydra configuration. Its machine-specific absolute paths must be overridden for the local machine. `dataset/text_embedding/manifest.json` and every embedding file it references must be downloaded together.

## Download

```python
from huggingface_hub import snapshot_download

snapshot_download(
    repo_id={repo_id!r},
    revision={revision!r},  # Replace with the returned commit oid for an immutable download.
    allow_patterns=[{repo_prefix + '/**'!r}],
)
```

For reproducible use, replace `revision` with the commit oid printed after upload. All listed paths are relative to the repository snapshot root.
"""
    return text.encode("utf-8")


def _artifact_manifest_bytes(
    repo_id: str, repo_prefix: str, artifacts: Iterable[Artifact]
) -> bytes:
    # A manifest cannot contain its own byte hash without a recursive fixed point.
    # It therefore inventories every other artifact and is itself verified by the
    # uploader's fixed-revision verification pass.
    payload = {
        "schema_version": 1,
        "repo_id": repo_id,
        "repo_prefix": repo_prefix,
        "files": [
            {
                "path": artifact.remote_path,
                "role": artifact.role,
                "size": artifact.size,
                "sha256": artifact.sha256,
            }
            for artifact in sorted(artifacts, key=lambda item: item.remote_path)
        ],
    }
    return (json.dumps(payload, indent=2, sort_keys=True) + "\n").encode("utf-8")


def build_upload_plan(
    *,
    run_dir: str | Path,
    steps: Iterable[int],
    repo_id: str,
    task: str,
    experiment_name: str,
    revision: str = "main",
) -> UploadPlan:
    task = _validate_segment(task, "task")
    experiment_name = _validate_segment(experiment_name, "experiment name")
    if (
        not isinstance(revision, str)
        or not revision
        or any(character.isspace() for character in revision)
        or ".." in PurePosixPath(revision).parts
        or "\\" in revision
    ):
        raise UploadValidationError(f"unsafe revision: {revision!r}")
    repo_prefix = f"{task}/{experiment_name}"
    _validate_remote_path(repo_prefix)
    step_values = tuple(sorted(steps))
    if (
        not step_values
        or len(set(step_values)) != len(step_values)
        or any(not isinstance(step, int) or isinstance(step, bool) or step <= 0 for step in step_values)
    ):
        raise UploadValidationError("steps must be unique positive integers")

    run_root = _assert_directory(Path(run_dir), "run directory")
    config_path = run_root / "config.yaml"
    config_artifact = _source_artifact(
        config_path, f"{repo_prefix}/config.yaml", "resolved_hydra_config", root=run_root
    )
    config, data_root, text_root = _load_config(config_path, run_root)
    try:
        model_id = config["model"]["model_id"]
        tokenizer_model_id = config["model"]["tokenizer_model_id"]
    except (KeyError, TypeError) as exc:
        raise UploadValidationError(
            "config.yaml must contain model.model_id and model.tokenizer_model_id"
        ) from exc
    if not all(
        isinstance(value, str) and value.strip()
        for value in (model_id, tokenizer_model_id)
    ):
        raise UploadValidationError("model identities must be non-empty strings")

    stats_path = data_root / "meta" / "stats.json"
    tasks_path = data_root / "meta" / "tasks.jsonl"
    stats_artifact = _source_artifact(
        stats_path,
        f"{repo_prefix}/dataset/meta/stats.json",
        "dataset_stats",
        root=data_root,
    )
    tasks_artifact = _source_artifact(
        tasks_path,
        f"{repo_prefix}/dataset/meta/tasks.jsonl",
        "dataset_tasks",
        root=data_root,
    )
    _load_json(stats_path, "dataset stats")
    task_instructions = _validate_tasks_jsonl(tasks_path)
    artifacts = [
        config_artifact,
        stats_artifact,
        tasks_artifact,
    ]

    manifest_path = text_root / "manifest.json"
    manifest_artifact = _source_artifact(
        manifest_path,
        f"{repo_prefix}/dataset/text_embedding/manifest.json",
        "text_embedding_manifest",
        root=text_root,
    )
    references = _manifest_references(
        manifest_path,
        expected_model_id=model_id,
        expected_tokenizer_model_id=tokenizer_model_id,
        task_instructions=task_instructions,
    )
    expected_text_names = {"manifest.json", *(path for path, _, _ in references)}
    actual_text_names = set()
    for item in text_root.iterdir():
        if item.is_symlink() or not item.is_file():
            raise UploadValidationError(
                f"text embedding directory contains non-regular/orphan item: {item}"
            )
        actual_text_names.add(item.name)
    if actual_text_names != expected_text_names:
        raise UploadValidationError(
            "text embedding directory does not exactly match manifest references"
        )
    artifacts.append(manifest_artifact)
    for relative_path, expected_size, expected_sha256 in references:
        local_path = text_root.joinpath(*PurePosixPath(relative_path).parts)
        artifact = _source_artifact(
            local_path,
            f"{repo_prefix}/dataset/text_embedding/{relative_path}",
            "text_embedding",
            root=text_root,
        )
        if artifact.size != expected_size or artifact.sha256 != expected_sha256:
            raise UploadValidationError(
                f"manifest size/SHA256 mismatch for embedding: {relative_path}"
            )
        artifacts.append(artifact)

    weights_root = run_root / "checkpoints" / "weights"
    _assert_directory(weights_root, "weights directory")
    if "state" in {part.casefold() for part in weights_root.parts}:
        raise UploadValidationError("weight source must not be located under state")
    for step in step_values:
        name = f"step_{step:06d}.pt"
        artifacts.append(
            _source_artifact(
                weights_root / name,
                f"{repo_prefix}/weights/{name}",
                "model_weights",
                root=run_root,
            )
        )

    readme_path = f"{repo_prefix}/README.md"
    manifest_output_path = f"{repo_prefix}/artifact_manifest.json"
    readme = _bytes_artifact(
        _readme_bytes(
            repo_id,
            repo_prefix,
            revision,
            step_values,
            [artifact.remote_path for artifact in artifacts]
            + [readme_path, manifest_output_path],
        ),
        readme_path,
        "readme",
    )
    artifacts.append(readme)
    manifest = _bytes_artifact(
        _artifact_manifest_bytes(repo_id, repo_prefix, artifacts),
        manifest_output_path,
        "artifact_manifest",
    )
    artifacts.append(manifest)

    paths = [artifact.remote_path for artifact in artifacts]
    if len(paths) != len(set(paths)):
        raise UploadValidationError("duplicate remote artifact paths")
    return UploadPlan(
        repo_id=repo_id,
        repo_prefix=repo_prefix,
        revision=revision,
        staging_parent=run_root.parent,
        steps=step_values,
        artifacts=tuple(sorted(artifacts, key=lambda item: item.remote_path)),
    )


def _assert_artifact_unchanged(artifact: Artifact) -> None:
    if artifact.local_path is None:
        return
    try:
        current = os.lstat(artifact.local_path)
    except OSError as exc:
        raise UploadValidationError(
            f"source disappeared before upload completed: {artifact.local_path}"
        ) from exc
    if stat.S_ISLNK(current.st_mode) or _identity(current) != artifact.identity:
        raise UploadValidationError(f"source changed after planning: {artifact.local_path}")


def _open_verified_artifact(artifact: Artifact):
    if artifact.local_path is None or artifact.identity is None:
        raise UploadValidationError(
            f"local artifact identity is unavailable: {artifact.remote_path}"
        )
    flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
    try:
        fd = os.open(artifact.local_path, flags)
    except OSError as exc:
        raise UploadValidationError(
            f"unable to safely open source: {artifact.local_path}"
        ) from exc
    fileobj = os.fdopen(fd, "rb")
    try:
        if _identity(os.fstat(fileobj.fileno())) != artifact.identity:
            raise UploadValidationError(
                f"source inode changed before upload: {artifact.local_path}"
            )
        digest = hashlib.sha256()
        for block in iter(lambda: fileobj.read(1024 * 1024), b""):
            digest.update(block)
        if digest.hexdigest() != artifact.sha256:
            raise UploadValidationError(
                f"source SHA256 changed before upload: {artifact.local_path}"
            )
        if _identity(os.fstat(fileobj.fileno())) != artifact.identity:
            raise UploadValidationError(
                f"source changed while opening for upload: {artifact.local_path}"
            )
        _assert_artifact_unchanged(artifact)
        fileobj.seek(0)
        return fileobj
    except Exception:
        fileobj.close()
        raise


def _clone_or_copy(source, destination) -> None:
    try:
        fcntl.ioctl(destination.fileno(), FICLONE, source.fileno())
        destination.seek(0)
        return
    except OSError as exc:
        if exc.errno not in {
            errno.EOPNOTSUPP,
            errno.EXDEV,
            errno.ENOTTY,
            errno.EINVAL,
            errno.ENOSYS,
        }:
            raise
    destination.seek(0)
    destination.truncate(0)
    source.seek(0)
    while True:
        block = source.read(1024 * 1024)
        if not block:
            break
        destination.write(block)
    destination.seek(0)


def _cleanup_snapshot_dir(path: Path) -> None:
    if not path.exists():
        return
    for child in path.iterdir():
        child.unlink()
    path.rmdir()


def _snapshot_operation_sources(plan: UploadPlan, stack: ExitStack):
    snapshot_dir = Path(
        tempfile.mkdtemp(prefix=".hf-upload-snapshot-", dir=plan.staging_parent)
    )
    os.chmod(snapshot_dir, 0o700)
    stack.callback(_cleanup_snapshot_dir, snapshot_dir)
    sources = []
    for index, artifact in enumerate(plan.artifacts):
        if artifact.content is not None:
            handle = BytesIO(artifact.content)
            stack.callback(handle.close)
            sources.append((artifact, handle))
            continue
        original = _open_verified_artifact(artifact)
        snapshot_path = snapshot_dir / f"{index:04d}.snapshot"
        try:
            fd = os.open(
                snapshot_path,
                os.O_RDWR | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0),
                0o600,
            )
            with os.fdopen(fd, "w+b") as writable:
                _clone_or_copy(original, writable)
                writable.flush()
                os.fsync(writable.fileno())
                writable.seek(0)
                digest = hashlib.sha256()
                for block in iter(lambda: writable.read(1024 * 1024), b""):
                    digest.update(block)
                if (
                    os.fstat(writable.fileno()).st_size != artifact.size
                    or digest.hexdigest() != artifact.sha256
                ):
                    raise UploadValidationError(
                        f"immutable snapshot mismatch: {artifact.remote_path}"
                    )
            read_fd = os.open(
                snapshot_path, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
            )
            handle = os.fdopen(read_fd, "rb")
            snapshot_path.unlink()
            if not stat.S_ISREG(os.fstat(handle.fileno()).st_mode):
                handle.close()
                raise UploadValidationError(
                    f"snapshot is not regular: {artifact.remote_path}"
                )
            stack.callback(handle.close)
            sources.append((artifact, handle))
        finally:
            original.close()
            if snapshot_path.exists():
                snapshot_path.unlink()
    return sources


def _lfs_sha256(value) -> str | None:
    if value is None:
        return None
    if isinstance(value, dict):
        return value.get("sha256")
    return getattr(value, "sha256", None)


def _hash_downloaded_file(path: Path) -> tuple[int, str]:
    """Hash an HF cache file while allowing its normal pointer symlink."""
    try:
        before = os.stat(path)
    except OSError as exc:
        raise UploadValidationError(f"downloaded cache file is unavailable: {path}") from exc
    if not stat.S_ISREG(before.st_mode):
        raise UploadValidationError(f"downloaded cache target is not regular: {path}")
    with path.open("rb") as fileobj:
        opened = os.fstat(fileobj.fileno())
        if not stat.S_ISREG(opened.st_mode) or _identity(opened) != _identity(before):
            raise UploadValidationError(f"downloaded cache target changed: {path}")
        digest = hashlib.sha256()
        for block in iter(lambda: fileobj.read(1024 * 1024), b""):
            digest.update(block)
        after = os.fstat(fileobj.fileno())
    if _identity(after) != _identity(before) or _identity(os.stat(path)) != _identity(before):
        raise UploadValidationError(f"downloaded cache target changed while hashing: {path}")
    return before.st_size, digest.hexdigest()


def _verify_revision(
    *,
    api,
    plan: UploadPlan,
    revision: str,
    download_file: Callable[..., str],
) -> None:
    remote_files = set(
        api.list_repo_files(plan.repo_id, repo_type=REPO_TYPE, revision=revision)
    )
    managed = {path for path in remote_files if path.startswith(plan.repo_prefix + "/")}
    if managed != plan.expected_paths:
        raise UploadValidationError(
            "remote managed path set mismatch: "
            f"expected={sorted(plan.expected_paths)!r}, actual={sorted(managed)!r}"
        )
    info = api.repo_info(
        plan.repo_id,
        repo_type=REPO_TYPE,
        revision=revision,
        files_metadata=True,
    )
    siblings = {item.rfilename: item for item in info.siblings}
    for artifact in plan.artifacts:
        remote = siblings.get(artifact.remote_path)
        if remote is None or remote.size != artifact.size:
            raise UploadValidationError(
                f"remote size mismatch: {artifact.remote_path}"
            )
        lfs_sha256 = _lfs_sha256(getattr(remote, "lfs", None))
        if artifact.role == "model_weights":
            if lfs_sha256 != artifact.sha256:
                raise UploadValidationError(
                    f"remote LFS SHA256 mismatch: {artifact.remote_path}"
                )
        elif lfs_sha256 is not None and lfs_sha256 != artifact.sha256:
            raise UploadValidationError(
                f"remote LFS SHA256 mismatch: {artifact.remote_path}"
            )
        if artifact.size <= SMALL_FILE_LIMIT:
            downloaded = Path(
                download_file(
                    repo_id=plan.repo_id,
                    filename=artifact.remote_path,
                    repo_type=REPO_TYPE,
                    revision=revision,
                    force_download=True,
                )
            )
            size, sha256 = _hash_downloaded_file(downloaded)
            if size != artifact.size or sha256 != artifact.sha256:
                raise UploadValidationError(
                    f"downloaded remote file mismatch: {artifact.remote_path}"
                )


def upload_plan(
    plan: UploadPlan,
    *,
    confirm_public_repo: str,
    revision: str | None = None,
    commit_message: str = "Upload real-machine experiment",
    api=None,
    download_file: Callable[..., str] | None = None,
) -> str:
    revision = plan.revision if revision is None else revision
    if revision != plan.revision:
        raise UploadValidationError("upload revision differs from planned README revision")
    if not isinstance(commit_message, str) or not commit_message.strip():
        raise UploadValidationError("commit message must be a non-empty string")
    _scan_public_text(commit_message, location="commit message")
    if confirm_public_repo != plan.repo_id:
        raise UploadValidationError(
            "--confirm-public-repo must exactly match --repo-id"
        )
    if api is None or download_file is None:
        from huggingface_hub import HfApi, hf_hub_download

        api = api or HfApi()
        download_file = download_file or hf_hub_download

    initial = api.repo_info(
        plan.repo_id,
        repo_type=REPO_TYPE,
        revision=revision,
        files_metadata=True,
    )
    if getattr(initial, "private", None) is not False:
        raise UploadValidationError("target repository must be confirmed public")
    parent_commit = getattr(initial, "sha", None)
    if not parent_commit:
        raise UploadValidationError("repository HEAD SHA is unavailable")
    existing = {
        item.rfilename
        for item in initial.siblings
        if item.rfilename.startswith(plan.repo_prefix + "/")
    }
    if existing and existing != plan.expected_paths:
        raise UploadValidationError(
            "existing experiment prefix must be empty or exactly match the plan: "
            f"expected={sorted(plan.expected_paths)!r}, actual={sorted(existing)!r}"
        )
    if existing == plan.expected_paths:
        _verify_revision(
            api=api,
            plan=plan,
            revision=parent_commit,
            download_file=download_file,
        )
        for artifact in plan.artifacts:
            _assert_artifact_unchanged(artifact)
        current = api.repo_info(
            plan.repo_id,
            repo_type=REPO_TYPE,
            revision=revision,
            files_metadata=True,
        )
        if getattr(current, "sha", None) != parent_commit:
            raise UploadValidationError("repository HEAD drifted during idempotent check")
        return parent_commit

    for artifact in plan.artifacts:
        _assert_artifact_unchanged(artifact)
    latest = api.repo_info(
        plan.repo_id,
        repo_type=REPO_TYPE,
        revision=revision,
        files_metadata=True,
    )
    if getattr(latest, "sha", None) != parent_commit:
        raise UploadValidationError("repository HEAD drifted before create_commit")

    from huggingface_hub import CommitOperationAdd

    with ExitStack() as stack:
        operation_sources = _snapshot_operation_sources(plan, stack)
        operations = [
            CommitOperationAdd(
                path_in_repo=artifact.remote_path,
                path_or_fileobj=source,
            )
            for artifact, source in operation_sources
        ]
        commit = api.create_commit(
            repo_id=plan.repo_id,
            repo_type=REPO_TYPE,
            revision=revision,
            operations=operations,
            commit_message=commit_message,
            parent_commit=parent_commit,
        )
        oid = getattr(commit, "oid", None) or getattr(commit, "commit_id", None)
        if not oid:
            raise UploadValidationError("Hugging Face did not return a commit oid")
        for artifact, handle in operation_sources:
            if handle.closed:
                raise UploadValidationError(
                    f"snapshot handle closed during upload: {artifact.remote_path}"
                )

    _verify_revision(
        api=api,
        plan=plan,
        revision=oid,
        download_file=download_file,
    )
    current = api.repo_info(
        plan.repo_id,
        repo_type=REPO_TYPE,
        revision=revision,
        files_metadata=True,
    )
    current_sha = getattr(current, "sha", None)
    if current_sha != oid:
        raise UploadValidationError(
            "repository HEAD no longer equals the uploaded commit oid"
        )
    return oid


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=__doc__.splitlines()[0],
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""默认是只读 dry-run。正式公开写入示例：
  %(prog)s --run-dir /path/to/run --steps 10000 20000 \\
    --repo-id owner/repo --task close_drawer --experiment-name experiment \\
    --revision main --upload --confirm-public-repo owner/repo
""",
    )
    parser.add_argument("--run-dir", required=True, help="训练 run 目录")
    parser.add_argument(
        "--steps", required=True, nargs="+", type=int, help="要发布的正整数训练 step"
    )
    parser.add_argument("--repo-id", required=True, help="目标 Hugging Face model repo")
    parser.add_argument("--task", required=True, help="远端一级任务目录名")
    parser.add_argument(
        "--experiment-name", required=True, help="远端二级实验目录名"
    )
    parser.add_argument(
        "--upload", action="store_true", help="执行远端写入；省略时仅 dry-run"
    )
    parser.add_argument(
        "--confirm-public-repo",
        help="正式写入时必须与 --repo-id 完全一致",
    )
    parser.add_argument("--revision", default="main", help="目标 branch，默认 main")
    parser.add_argument(
        "--commit-message",
        default="Upload real-machine experiment",
        help="提交消息；上传前执行秘密扫描",
    )
    return parser


def _print_plan(plan: UploadPlan, *, upload: bool) -> None:
    mode = "UPLOAD" if upload else "DRY-RUN"
    print(f"[{mode}] repo={plan.repo_id} prefix={plan.repo_prefix}")
    for artifact in plan.artifacts:
        print(
            f"{artifact.remote_path}\t{artifact.role}\t"
            f"{artifact.size}\t{artifact.sha256}"
        )
    print(f"files={len(plan.artifacts)} total_bytes={plan.total_size}")


def main(argv: list[str] | None = None) -> int:
    args = _build_parser().parse_args(argv)
    plan = build_upload_plan(
        run_dir=args.run_dir,
        steps=args.steps,
        repo_id=args.repo_id,
        task=args.task,
        experiment_name=args.experiment_name,
        revision=args.revision,
    )
    _print_plan(plan, upload=args.upload)
    if not args.upload:
        return 0
    if args.confirm_public_repo is None:
        raise SystemExit("--upload requires --confirm-public-repo <repo-id>")
    oid = upload_plan(
        plan,
        confirm_public_repo=args.confirm_public_repo,
        revision=args.revision,
        commit_message=args.commit_message,
    )
    print(f"[DONE] revision={oid}")
    return 0


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