ComfyAssets-ComfyUI_PromptManager/py/api/thumbnail_move_plan.py

"""Move thumbnails, the review (#320 Part 3 phase 1): what a move would copy,
remove and leave, checked before anything is written."""

import hashlib
import json
import os
import secrets
import shutil
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path

from aiohttp import web

from . import preview_location, preview_store
from . import thumbnail_move_files as files
from .job_routes import busy_response
from .jobs import THUMBNAILS_JOB_KIND, JobBusy
from .thumbnail_location import default_path

try:
    from ...utils import library_location
except ImportError:  # tests import the package top-level
    from utils import library_location

MOVE_TITLE = "Move thumbnails"
FINISHING_TITLE = "Move thumbnails (finishing)"
PRUNE_TITLE = "Remove thumbnails of removed sources"
MOVE_GROUPS = ("moved", "skipped", "left", "errors")
PRUNE_GROUPS = ("removed", "left", "errors")
PLAN_TTL = 600.0
JOURNAL_NAME = "thumbnail-move.json"
SPACE_MARGIN = 64 * 2**20
SPACE_SHARE = 0.05
LISTED = 5
SENTENCE = (
    "Every thumbnail is copied and checked before Prompt Manager switches. "
    "Scans, deletes and tagging wait until it finishes; closing this page does "
    "not stop it."
)


def _number(value):
    return f"{value:,}"


def _plural(count, one, many=None):
    return f"{_number(count)} {one if count == 1 else (many or one + 's')}"


def _size(value):
    size = float(value)
    for unit in ("B", "KB", "MB", "GB"):
        if size < 1024 or unit == "GB":
            return f"{size:.1f} {unit}" if unit != "B" else f"{int(size)} B"
        size /= 1024
    return f"{size:.1f} GB"


def _refusal(code, status=None, **extra):
    known_status, text = preview_location.MESSAGES.get(code, (400, code))
    return web.json_response(
        {"success": False, "code": code, "error": text, **extra},
        status=status or known_status,
    )


def journal_path():
    """``<data folder>/data/thumbnail-move.json`` (only while a move copies)."""
    return library_location.cache_dir(JOURNAL_NAME)


def _read_json(path):
    try:
        with open(path, "r", encoding="utf-8") as handle:
            data = json.load(handle)
    except (OSError, ValueError):
        return None
    return data if isinstance(data, dict) else None


def _write_json(path, data):
    os.makedirs(os.path.dirname(path), exist_ok=True)
    temp = f"{path}.{os.getpid()}.tmp"
    with open(temp, "w", encoding="utf-8") as handle:
        json.dump(data, handle, indent=2)
        handle.flush()
        os.fsync(handle.fileno())
    os.replace(temp, path)


def _remove(path):
    try:
        os.unlink(path)
    except FileNotFoundError:
        pass  # already gone: removing it is the goal


def _is_network(path):
    text = str(path)
    if text.startswith(("\\\\", "//")):
        return True
    if os.name == "nt":
        try:
            import ctypes

            drive = os.path.splitdrive(os.path.abspath(text))[0] + "\\"
            return ctypes.windll.kernel32.GetDriveTypeW(drive) == 4  # DRIVE_REMOTE
        except (AttributeError, OSError):
            return False
    return False


def _volume(path):
    current = Path(os.path.abspath(path))
    while not current.exists() and current.parent != current:
        current = current.parent
    try:
        return os.stat(current).st_dev
    except OSError:
        return None


def _free(path):
    current = Path(os.path.abspath(path))
    while not current.exists() and current.parent != current:
        current = current.parent
    try:
        return shutil.disk_usage(current).free
    except OSError:
        return None


# ── the plan ─────────────────────────────────────────────────────────


@dataclass
class SourcePlan:
    """One listed source: its old store, its new store and what moves."""

    root: str
    label: str
    key: str | None
    old_store: str | None
    new_store: str
    scan: files.StoreScan | None
    available: bool = True
    too_long: list = field(default_factory=list)


@dataclass
class MovePlan:
    """A reviewed move: from, to, the sources and the removals."""

    token: str
    created: float
    revision: str
    old: preview_location.ParsedLocation
    target: dict  # {"mode", "path"}
    target_id: str | None
    target_exists: bool
    target_default: bool
    install: str
    sources: list
    unlisted: list  # (key, store path, scan)
    needs_bytes: int
    confirm_required: bool

    @property
    def to_folder(self):
        return self.target["mode"] == preview_store.FOLDER_MODE

    @property
    def from_folder(self):
        return self.old.mode == preview_store.FOLDER_MODE


def _roots_revision(roots):
    text = json.dumps(sorted(str(root) for root in roots))
    return hashlib.sha256(text.encode("utf-8")).hexdigest()[:16]


class ThumbnailMovePlanMixin:
    """The move's review and its routes."""

    _move_plans = None

    # ── routes ────────────────────────────────────────────────────────

    def _register_thumbnail_move_routes(self, routes):
        @routes.post("/prompt_manager/images/thumbnails/plan")
        async def thumbnail_move_plan_route(request):
            return await self.plan_thumbnail_move(request)

        @routes.post("/prompt_manager/images/thumbnails/location/recreate")
        async def thumbnail_location_recreate_route(request):
            return await self.recreate_thumbnail_folder(request)

        @routes.post("/prompt_manager/images/thumbnails/location/adopt")
        async def thumbnail_location_adopt_route(request):
            return await self.adopt_thumbnail_folder(request)

    # ── review ────────────────────────────────────────────────────────

    async def plan_thumbnail_move(self, request):
        try:
            body = await request.json()
        except (ValueError, UnicodeDecodeError):
            body = None
        if not isinstance(body, dict) or body.get("mode") != "move":
            return web.json_response(
                {
                    "success": False,
                    "code": "invalid-request",
                    "error": "mode must be move",
                },
                status=400,
            )
        target = body.get("location")
        if not isinstance(target, dict) or target.get("mode") not in (
            preview_store.SOURCE_MODE,
            preview_store.FOLDER_MODE,
        ):
            return web.json_response(
                {
                    "success": False,
                    "code": "invalid-request",
                    "error": "location must name a mode",
                },
                status=400,
            )
        blocker = self._jobs.blocker(THUMBNAILS_JOB_KIND, admission="guarded")
        if blocker is not None:
            return busy_response(JobBusy(blocker))
        try:
            outcome = await self._run_in_executor(self._build_move_plan, target)
        except Exception:
            self.logger.exception("The thumbnail move could not be planned")
            return web.json_response(
                {
                    "success": False,
                    "code": None,
                    "error": "The move could not be planned",
                },
                status=500,
            )
        if isinstance(outcome, str):
            return _refusal(outcome)
        plan, payload = outcome
        if payload.get("insufficient"):
            return _refusal(
                "insufficient-space",
                status=507,
                error_detail=payload["insufficient"],
                plan=payload,
            )
        self._store_move_plan(plan)
        return web.json_response({"success": True, **payload})

    def _store_move_plan(self, plan):
        if self._move_plans is None:
            self._move_plans = {}
        now = time.monotonic()
        self._move_plans = {
            token: kept
            for token, kept in self._move_plans.items()
            if now - kept.created < PLAN_TTL
        }
        self._move_plans[plan.token] = plan

    def _build_move_plan(self, target):
        """``(plan, payload)`` or a refusal code (blocking)."""
        old = preview_store.setting()
        location = preview_store.current()
        install = old.install or str(uuid.uuid4())
        target_path = None
        default = default_path()
        if target["mode"] == preview_store.FOLDER_MODE:
            raw = target.get("path")
            if not isinstance(raw, str):
                return "location-not-absolute"
            context = self._thumbnail_location_context()
            code = preview_location.location_problem(
                raw, purpose="move", context=context
            )
            if code is not None:
                return code
            target_path = os.path.abspath(raw)
        elif old.mode == preview_store.SOURCE_MODE and not old.problem:
            return "location-same"
        target_default = bool(default) and os.path.normcase(
            target_path or ""
        ) == os.path.normcase(os.path.abspath(default))
        target_exists = bool(target_path) and os.path.isdir(target_path)
        target_marker = (
            preview_location.read_marker(target_path) if target_exists else None
        )
        roots = [str(root) for root in self._get_all_output_dirs()]
        labels = self._source_labels(roots)
        old_keys = {}
        if old.mode == preview_store.FOLDER_MODE and location.available:
            marker = preview_location.read_marker(old.path) or {}
            old_keys = marker.get("keys") or {}
        new_keys = dict((target_marker or {}).get("keys") or {})
        taken = {key: entry.get("identity") for key, entry in new_keys.items()}
        sources = []
        for root in roots:
            identity = os.path.normcase(os.path.realpath(root))
            dev_ino = preview_location.file_identity(root)
            old_key = None
            old_store = None
            available = True
            if old.mode == preview_store.SOURCE_MODE:
                old_store = str(preview_store.source_folder(root))
            elif location.available:
                old_key = preview_location.find_key(old_keys, root, identity, dev_ino)
                if old_key:
                    old_store = os.path.join(old.path, old_key)
            else:
                available = False
            key = None
            if target_path:
                key = old_key or preview_location.find_key(
                    new_keys, root, identity, dev_ino
                )
                if key is None or taken.get(key, identity) != identity:
                    key = preview_location.source_key(root, identity, taken)
                taken[key] = identity
                new_store = os.path.join(target_path, key)
            else:
                new_store = str(preview_store.source_folder(root))
            scan = None
            if old_store:
                scan = files.scan_store(old_store, root)
            plan = SourcePlan(
                root, labels.get(root, root), key, old_store, new_store, scan, available
            )
            if scan is not None and target_path:
                plan.too_long = [
                    item[0]
                    for item in scan.copy
                    if files.too_long(Path(new_store).joinpath(*item[0]))
                ]
            sources.append(plan)
        unlisted = self._unlisted_keys(old, location, old_keys, install)
        needs = sum(source.scan.bytes for source in sources if source.scan is not None)
        plan = MovePlan(
            token=secrets.token_urlsafe(16),
            created=time.monotonic(),
            revision=_roots_revision(roots),
            old=old,
            target={"mode": target["mode"], "path": target_path},
            target_id=(target_marker or {}).get("id"),
            target_exists=target_exists,
            target_default=target_default,
            install=install,
            sources=sources,
            unlisted=unlisted,
            needs_bytes=needs,
            confirm_required=bool(target_path) and not target_default,
        )
        return plan, self._plan_payload(plan, location)

    def _unlisted_keys(self, old, location, keys, install):
        """This install's keys whose source no longer is in any root list
        (``(key, store, scan)``), in the folder in use."""
        if old.mode != preview_store.FOLDER_MODE or not location.available:
            return []
        listed = {
            os.path.normcase(os.path.realpath(root))
            for root in self._all_source_roots()
        }
        found = []
        for key, entry in sorted(keys.items()):
            if not preview_location.valid_key(key) or not isinstance(entry, dict):
                continue  # a marker is data: never a path out of the folder
            if entry.get("install") not in (install, None):
                continue
            if entry.get("identity") in listed:
                continue
            store = os.path.join(old.path, key)
            found.append(
                (key, store, files.scan_store(store, entry.get("path") or store))
            )
        return found

    def _source_labels(self, roots):
        from . import sources

        labels = sources._comfyui_labels()
        return {root: sources.source_label(root, labels) for root in roots}

    def _plan_payload(self, plan, location):
        sources = plan.sources
        copies = sum(len(s.scan.copy) for s in sources if s.scan)
        orphans = sum(len(s.scan.orphans) for s in sources if s.scan)
        temps = sum(len(s.scan.temps) for s in sources if s.scan)
        links = sum(s.scan.links for s in sources if s.scan)
        left = [
            f"{s.label} · {'/'.join(item)}"
            for s in sources
            if s.scan
            for item in s.scan.left
        ]
        too_long = [
            f"{s.label} · {'/'.join(item)}" for s in sources for item in s.too_long
        ]
        unlisted_count = sum(
            len(scan.copy) + len(scan.orphans) for _k, _s, scan in plan.unlisted
        )
        free = None
        insufficient = None
        if plan.needs_bytes:
            volumes = {}
            for source in sources:
                if source.scan is None or not source.scan.bytes:
                    continue
                volume = _volume(source.new_store)
                entry = volumes.setdefault(volume, [source.new_store, 0])
                entry[1] += source.scan.bytes
            for where, needed in volumes.values():
                available = _free(where)
                if available is None:
                    continue
                free = available if free is None else min(free, available)
                margin = max(SPACE_MARGIN, int(needed * SPACE_SHARE))
                if needed + margin > available:
                    insufficient = {
                        "needed": needed + margin,
                        "free": available,
                        "text": (
                            f"There is not enough space for the thumbnails: "
                            f"{_size(available)} free, {_size(needed + margin)} "
                            "needed. Free some space or choose another folder."
                        ),
                    }
        payload = {
            "mode": "move",
            "from": {
                "mode": plan.old.mode,
                "path": plan.old.path,
                "available": location.available,
            },
            "to": {
                "mode": plan.target["mode"],
                "path": plan.target["path"],
                "exists": plan.target_exists,
                "will_create": bool(plan.target["path"]) and not plan.target_exists,
                "default": plan.target_default,
            },
            "sources": [
                {
                    "label": s.label,
                    "key": s.key,
                    "files": len(s.scan.copy) if s.scan else 0,
                    "bytes": s.scan.bytes if s.scan else 0,
                    "too_long": len(s.too_long),
                    "links": s.scan.links if s.scan else 0,
                    "orphans": len(s.scan.orphans) if s.scan else 0,
                    "left_in_place": len(s.scan.left) if s.scan else 0,
                    "available": s.available,
                }
                for s in sources
            ],
            "totals": {
                "sources": len(sources),
                "files": copies,
                "bytes": plan.needs_bytes,
                "orphans": orphans,
                "temps": temps,
                "links": links,
                "left_in_place": len(left),
                "too_long": len(too_long),
                "unlisted_sources": len(plan.unlisted),
                "unlisted_files": unlisted_count,
            },
            "too_long_listed": too_long[:LISTED],
            "left_listed": left[:LISTED],
            "unlisted": [
                {"key": key, "files": len(scan.copy) + len(scan.orphans)}
                for key, _store, scan in plan.unlisted
            ],
            "needs_bytes": plan.needs_bytes,
            "free_bytes": free,
            "confirm_required": plan.confirm_required,
            "token": plan.token,
            "sentence": SENTENCE,
        }
        if insufficient:
            payload["insufficient"] = insufficient
        return payload