Blame for processor/scripts/cleanup_active_media.py

"""Clean up active LeRobot v2.1 media without touching recoverable history.

The first D435 collectors used several names for the same two physical
streams.  This utility normalizes them to ``D435_rgb`` and ``D435_depth`` in
active datasets.  It also removes files advertised as pure depth when ffprobe
shows that they are colour/video pixels rather than metric
``gray12le``/``gray16le`` depth.  Every moved video and every rewritten
metadata file is preserved in ``data/.backups/session-migrations``.

Default mode is a read-only report.  Use ``--apply`` only after reviewing it.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any


APP_ROOT = Path(__file__).resolve().parents[1]
if str(APP_ROOT) not in sys.path:
    sys.path.insert(0, str(APP_ROOT))

from app.lerobot_v21 import (  # noqa: E402
    VIDEO_SUFFIXES,
    canonical_source_key,
    is_depth_source,
    iter_video_streams,
)
from app.project_dataset import (  # noqa: E402
    episode_files,
    episode_row,
    project_episode_rows,
    write_project_episode_index,
)


EPISODE_RE = re.compile(r"episode_(\d+)$")
SOURCE_RENAMES = {
    "D435_depth_rgb": "D435_rgb",
    "D435_head_rgb": "D435_rgb",
    "D435_head_depth": "D435_depth",
}
DEPTH_PIX_FMTS = {"gray12le", "gray16le"}
DEVICE_NAME_ALIASES = {
    "d435_depth": "D435",
    "d435_head": "D435",
}


def _probe(path: Path) -> dict[str, str]:
    try:
        result = subprocess.run(
            ["ffprobe", "-v", "error", "-select_streams", "v:0",
             "-show_entries", "stream=codec_name,pix_fmt,width,height",
             "-of", "json", str(path)],
            capture_output=True, text=True, timeout=30, check=False,
        )
        streams = json.loads(result.stdout or "{}").get("streams") or []
        return {str(k): str(v or "") for k, v in (streams[0] if streams else {}).items()}
    except (OSError, ValueError, TypeError, subprocess.TimeoutExpired):
        return {}


def _copy_or_link(source: Path, target: Path) -> None:
    target.parent.mkdir(parents=True, exist_ok=True)
    if target.exists():
        return
    try:
        os.link(source, target)
    except OSError:
        shutil.copy2(source, target)


def _backup_file(source: Path, project: Path, backup: Path) -> None:
    _copy_or_link(source, backup / project.name / source.relative_to(project))


def _backup_video(source: Path, project: Path, backup: Path) -> None:
    _copy_or_link(source, backup / project.name / source.relative_to(project))


def _metadata_files(project: Path) -> list[Path]:
    meta = project / "meta"
    if not meta.is_dir():
        return []
    return sorted(path for path in meta.rglob("*")
                  if path.is_file() and path.suffix.lower() in {".json", ".jsonl"})


def _merge_values(current: Any, incoming: Any) -> Any:
    """Merge colliding source-key metadata while keeping canonical values."""
    if isinstance(current, dict) and isinstance(incoming, dict):
        merged = dict(current)
        for key, value in incoming.items():
            if key not in merged:
                merged[key] = value
            else:
                merged[key] = _merge_values(merged[key], value)
        return merged
    if isinstance(current, list) and isinstance(incoming, list):
        result = list(current)
        for value in incoming:
            if value not in result:
                result.append(value)
        return result
    return current


def _rename_source_value(value: Any, old: str, new: str) -> Any:
    old_feature = f"observation.images.{old}"
    new_feature = f"observation.images.{new}"
    if isinstance(value, str):
        return value.replace(old_feature, new_feature).replace(old, new)
    if isinstance(value, list):
        return [_rename_source_value(item, old, new) for item in value]
    if not isinstance(value, dict):
        return value
    result: dict[str, Any] = {}
    for key, child in value.items():
        renamed_key = key.replace(old_feature, new_feature).replace(old, new)
        renamed_child = _rename_source_value(child, old, new)
        if renamed_key in result:
            result[renamed_key] = _merge_values(result[renamed_key], renamed_child)
        else:
            result[renamed_key] = renamed_child
    return result


def _rewrite_source_names(project: Path, old: str, new: str,
                          backup: Path | None) -> list[str]:
    changed: list[str] = []
    for path in _metadata_files(project):
        try:
            text = path.read_text(encoding="utf-8")
        except (OSError, UnicodeError):
            continue
        updated = text
        if path.suffix.lower() == ".json":
            try:
                value = json.loads(text)
                updated = json.dumps(
                    _rename_source_value(value, old, new),
                    ensure_ascii=False, indent=2,
                )
            except (ValueError, TypeError):
                updated = text.replace(
                    f"observation.images.{old}",
                    f"observation.images.{new}",
                ).replace(old, new)
        elif path.suffix.lower() == ".jsonl":
            lines = []
            for line in text.splitlines():
                if not line.strip():
                    continue
                try:
                    lines.append(json.dumps(
                        _rename_source_value(json.loads(line), old, new),
                        ensure_ascii=False,
                    ))
                except (ValueError, TypeError):
                    lines.append(line.replace(
                        f"observation.images.{old}",
                        f"observation.images.{new}",
                    ).replace(old, new))
            updated = "\n".join(lines) + ("\n" if lines else "")
        else:
            updated = text.replace(
                f"observation.images.{old}",
                f"observation.images.{new}",
            ).replace(old, new)
        if updated == text:
            continue
        if backup is not None:
            _backup_file(path, project, backup)
        path.write_text(updated, encoding="utf-8")
        changed.append(str(path.relative_to(project)))
    return changed


def _canonical_device_name(value: Any, kind: Any = "", key: Any = "") -> Any:
    if not isinstance(value, str):
        return value
    low = value.strip().lower()
    context = f"{kind} {key}".lower()
    if low in DEVICE_NAME_ALIASES and "d435" in context:
        return DEVICE_NAME_ALIASES[low]
    return value


def _normalise_devices(items: Any) -> Any:
    if not isinstance(items, list):
        return items
    merged: list[Any] = []
    by_key: dict[str, int] = {}
    for raw in items:
        if not isinstance(raw, dict):
            merged.append(raw)
            continue
        item = _normalise_device_metadata(raw)
        key = str(item.get("key") or "").strip()
        if not key or key not in by_key:
            if key:
                by_key[key] = len(merged)
            merged.append(item)
            continue
        index = by_key[key]
        merged[index] = _merge_values(merged[index], item)
    return merged


def _normalise_device_metadata(value: Any, field: str = "",
                               kind: Any = "", key: Any = "") -> Any:
    """Normalize physical D435 names without changing stream keys."""
    if isinstance(value, list):
        return [_normalise_device_metadata(item, field, kind, key)
                for item in value]
    if not isinstance(value, dict):
        if field in {"name", "device", "device_name"}:
            return _canonical_device_name(value, kind, key)
        return value
    local_kind = value.get("kind") or kind
    local_key = value.get("key") or key
    result: dict[str, Any] = {}
    for name, child in value.items():
        if name == "devices":
            child = _normalise_devices(child)
        elif name == "device_names" and isinstance(child, dict):
            child = {
                stream: _canonical_device_name(display, "d435", stream)
                for stream, display in child.items()
            }
        elif name in {"name", "device", "device_name"}:
            child = _canonical_device_name(child, local_kind, local_key)
        else:
            child = _normalise_device_metadata(child, name, local_kind, local_key)
        result[name] = child
    return result


def _rewrite_device_metadata(project: Path, backup: Path | None) -> list[str]:
    changed: list[str] = []
    for path in _metadata_files(project):
        if path.suffix.lower() != ".json":
            continue
        try:
            before = path.read_text(encoding="utf-8")
            value = json.loads(before)
        except (OSError, UnicodeError, ValueError, TypeError):
            continue
        updated_value = _normalise_device_metadata(value)
        updated = json.dumps(updated_value, ensure_ascii=False, indent=2)
        if updated == before:
            continue
        if backup is not None:
            _backup_file(path, project, backup)
        path.write_text(updated, encoding="utf-8")
        changed.append(str(path.relative_to(project)))
    return changed


def _rewrite_video_extensions(project: Path, backup: Path | None) -> list[str]:
    """Make the v2.1 extension index agree with active video files."""
    info_path = project / "meta" / "info.json"
    if not info_path.is_file():
        return []
    actual = {
        source: path.suffix.lower().lstrip(".")
        for source, path in iter_video_streams(project / "videos")
    }
    if not actual:
        return []
    try:
        before = info_path.read_text(encoding="utf-8")
        info = json.loads(before)
    except (OSError, UnicodeError, ValueError, TypeError):
        return []
    if not isinstance(info, dict):
        return []
    current = info.get("video_extensions")
    if current == actual:
        return []
    info["video_extensions"] = dict(sorted(actual.items()))
    if backup is not None:
        _backup_file(info_path, project, backup)
    info_path.write_text(json.dumps(info, ensure_ascii=False, indent=2),
                         encoding="utf-8")
    return [str(info_path.relative_to(project))]


def _rewrite_json(path: Path, transform) -> bool:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeError, ValueError, TypeError):
        return False
    updated = transform(value)
    if updated == value:
        return False
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    temporary.write_text(json.dumps(updated, ensure_ascii=False, indent=2),
                        encoding="utf-8")
    os.replace(temporary, path)
    return True


def _remove_source_from_value(value: Any, source: str) -> Any:
    """Remove a source only from source-bearing metadata fields."""
    source_fields = {"cameras", "device_names"}
    list_fields = {"slots", "depth_sources", "metric_depth_sources", "camera_sources"}
    if isinstance(value, list):
        return [_remove_source_from_value(item, source) for item in value
                if item != source]
    if not isinstance(value, dict):
        return value
    result: dict[str, Any] = {}
    for key, item in value.items():
        if key in source_fields and isinstance(item, dict):
            item = {name: child for name, child in item.items() if name != source}
        elif key in list_fields and isinstance(item, list):
            item = [child for child in item if child != source]
        result[key] = _remove_source_from_value(item, source)
    return result


def _remove_episode_source_metadata(project: Path, episode_id: str,
                                    source: str, backup: Path | None) -> list[str]:
    changed: list[str] = []
    collector = project / "meta" / "collector" / episode_id
    if not collector.is_dir():
        return changed
    for path in sorted(collector.glob("*.json")):
        before = path.read_text(encoding="utf-8", errors="ignore")
        if source not in before:
            continue
        if backup is not None:
            _backup_file(path, project, backup)
        if _rewrite_json(path, lambda value: _remove_source_from_value(value, source)):
            changed.append(str(path.relative_to(project)))
    return changed


def _remove_episode_row_source(project: Path, episode_index: int, source: str,
                               backup: Path | None) -> bool:
    rows = project_episode_rows(project)
    if not rows:
        return False
    changed = False
    prefix = f"videos/observation.images.{source}/"
    updated: list[dict[str, Any]] = []
    for row in rows:
        if int(row.get("episode_index", -1)) == episode_index:
            filtered = {key: value for key, value in row.items()
                        if not str(key).startswith(prefix)}
            changed |= filtered != row
            row = filtered
        updated.append(row)
    if not changed:
        return False
    if backup is not None:
        for path in sorted((project / "meta" / "episodes").rglob("*.parquet")):
            _backup_file(path, project, backup)
    write_project_episode_index(project, updated)
    return True


def _episode_id(project: Path, index: int) -> str:
    for row in project_episode_rows(project):
        if int(row.get("episode_index", -1)) == index:
            return str(row.get("episode_id") or row.get("source_batch") or "")
    return ""


def _active_video_paths(project: Path) -> list[Path]:
    videos = project / "videos"
    if not videos.is_dir():
        return []
    return sorted(path for path in videos.rglob("*.mp4")
                  if path.is_file() and any(part.startswith("chunk-") for part in path.parts))


def clean_project(project: Path, backup: Path | None, apply: bool) -> dict[str, Any]:
    result: dict[str, Any] = {"project": project.name, "renamed": [],
                              "invalid_depth": [], "metadata": []}
    if not project.is_dir():
        return result

    # Normalize historical D435 source names in every chunk.  If the target
    # directory already exists, merge only non-colliding episode files into
    # it; the caller has already verified that episode numbers do not overlap.
    for old_source, new_source in SOURCE_RENAMES.items():
        old_dirs = sorted((project / "videos").rglob(
            f"observation.images.{old_source}"))
        for old_dir in old_dirs:
            if not old_dir.is_dir():
                continue
            new_dir = old_dir.with_name(f"observation.images.{new_source}")
            result["renamed"].append({
                "from": str(old_dir.relative_to(project)),
                "to": str(new_dir.relative_to(project)),
            })
            if not apply:
                continue
            if new_dir.exists():
                conflicts = sorted(
                    item.name for item in old_dir.iterdir()
                    if (new_dir / item.name).exists()
                )
                if conflicts:
                    raise RuntimeError(
                        f"Rename collision: {new_dir} ({', '.join(conflicts)})"
                    )
            if backup is not None:
                for path in old_dir.rglob("*"):
                    if path.is_file():
                        _backup_video(path, project, backup)
            new_dir.mkdir(parents=True, exist_ok=True)
            for item in sorted(old_dir.iterdir()):
                target = new_dir / item.name
                if target.exists():
                    raise RuntimeError(f"Rename collision: {target}")
                item.rename(target)
            old_dir.rmdir()
        if apply and old_dirs:
            result["metadata"].extend(_rewrite_source_names(
                project, old_source, new_source, backup))

    # Merge duplicate physical-device records left by historical uploads.
    # Stream keys remain untouched: D435_rgb and D435_depth are data sources,
    # while D435 is the single hardware name shown by the UI.
    if apply:
        result["metadata"].extend(_rewrite_device_metadata(project, backup))
        result["metadata"].extend(_rewrite_video_extensions(project, backup))

    # A metric depth stream must be lossless gray12/gray16.  Keep the valid
    # files and move only the colour-coded/ordinary-video members out of the
    # active dataset.  This handles mixed historical batches in one source.
    invalid: list[tuple[Path, str, int]] = []
    for path in _active_video_paths(project):
        relative_parts = path.relative_to(project).parts
        source_parts = [part for part in relative_parts
                        if part.startswith("observation.images.")]
        if not source_parts:
            continue
        source = canonical_source_key(source_parts[-1])
        if not is_depth_source(source):
            continue
        probe = _probe(path)
        if probe.get("pix_fmt") in DEPTH_PIX_FMTS:
            continue
        match = EPISODE_RE.search(path.stem)
        if not match:
            continue
        invalid.append((path, source, int(match.group(1))))
    for path, source, episode_index in invalid:
        result["invalid_depth"].append({
            "path": str(path.relative_to(project)),
            "source": source,
            "episode_index": episode_index,
            "probe": _probe(path),
        })
        if not apply:
            continue
        if backup is not None:
            _backup_video(path, project, backup)
        path.unlink()
        episode_id = _episode_id(project, episode_index)
        if episode_id:
            result["metadata"].extend(_remove_episode_source_metadata(
                project, episode_id, source, backup))
        if _remove_episode_row_source(project, episode_index, source, backup):
            result["metadata"].append("meta/episodes/")

    return result


def clean_duplicate_processed_depth(root: Path, backup: Path | None,
                                    apply: bool) -> list[dict[str, Any]]:
    """Move raw PNG sidecars when the same episode has canonical depth video."""
    processed = root / "processed"
    sessions = root / "sessions"
    result: list[dict[str, Any]] = []
    if not processed.is_dir():
        return result

    for raw_root in sorted(processed.rglob("raw_depth")):
        if not raw_root.is_dir():
            continue
        try:
            relative = raw_root.relative_to(processed)
            project_name, episode_name = relative.parts[0], relative.parts[1]
        except (ValueError, IndexError):
            continue
        project = sessions / project_name
        if not project.is_dir():
            continue
        row = episode_row(project, episode_name)
        if not row:
            continue
        try:
            episode_index = int(row["episode_index"])
        except (KeyError, TypeError, ValueError):
            continue
        canonical = {
            source: path for source, path in episode_files(project, episode_index)["videos"]
        }
        for source_dir in sorted(path for path in raw_root.rglob("*") if path.is_dir()):
            pngs = [path for path in source_dir.glob("*.png")
                    if path.is_file() and path.stem.isdigit()]
            if not pngs:
                continue
            source = canonical_source_key(source_dir.name)
            video = canonical.get(source)
            if video is None or not is_depth_source(source):
                continue
            probe = _probe(video)
            if probe.get("pix_fmt") not in DEPTH_PIX_FMTS:
                continue
            item = {
                "path": str(source_dir.relative_to(root)),
                "project": project_name,
                "episode_id": episode_name,
                "episode_index": episode_index,
                "source": source,
                "png_count": len(pngs),
                "canonical_video": str(video.relative_to(root)),
                "canonical_probe": probe,
            }
            result.append(item)
            if apply:
                if backup is None:
                    raise RuntimeError("backup path is required for cleanup")
                target = backup / "processed-legacy" / source_dir.relative_to(root)
                target.parent.mkdir(parents=True, exist_ok=True)
                original_parent = source_dir.parent
                shutil.move(str(source_dir), str(target))
                # Remove only the empty legacy raw_depth shell left behind in
                # the active processed run; never remove a non-empty run or
                # any sibling annotation output.
                parent = original_parent
                while parent != processed and parent != root:
                    try:
                        parent.rmdir()
                    except OSError:
                        break
                    parent = parent.parent
    return result


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, required=True,
                        help="active data root containing sessions/")
    parser.add_argument("--project", action="append",
                        help="limit to one or more project directory names")
    parser.add_argument("--apply", action="store_true",
                        help="perform moves and metadata updates; default is dry-run")
    args = parser.parse_args()

    root = args.root.expanduser().resolve()
    sessions = root / "sessions"
    if not sessions.is_dir():
        raise SystemExit(f"sessions directory not found: {sessions}")
    wanted = set(args.project or [])
    projects = [path for path in sorted(sessions.iterdir())
                if path.is_dir() and (not wanted or path.name in wanted)]
    stamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    backup = root / ".backups" / "session-migrations" / f"media-cleanup-{stamp}" \
        if args.apply else None
    if backup is not None:
        backup.mkdir(parents=True, exist_ok=False)

    report = {"root": str(root), "apply": bool(args.apply),
              "backup": str(backup) if backup else None, "projects": [],
              "legacy_processed_depth": []}
    for project in projects:
        report["projects"].append(clean_project(project, backup, args.apply))
    report["legacy_processed_depth"] = clean_duplicate_processed_depth(
        root, backup, args.apply)

    if args.apply:
        for item in report["projects"]:
            project = sessions / item["project"]
            if backup is None:
                continue
            # Keep maintenance reports outside the canonical LeRobot meta
            # namespace.  ``meta`` should contain dataset metadata only.
            manifest = backup / "reports" / f"{item['project']}.json"
            manifest.parent.mkdir(parents=True, exist_ok=True)
            manifest.write_text(json.dumps(item, ensure_ascii=False, indent=2),
                                encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2))
    return 0


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