HVideoTool/hvideotool/core/project.py

"""A HVideoTool *project*: a self-contained folder on disk.

This replaces the old "open a bare folder of images" model. A project is a folder
that holds:

```
MyProject/
├── project.json     # version, name, created, source, settings{detector/model/threshold/restore}
├── frames/          # the images (what used to be "the folder")
├── detections.json  # the detection cache (basename -> detections, tagged with detector key)
└── collections/     # curation sub-folders (created lazily)
```

The project file carries the **per-project** settings (detector, model, overlay
threshold, restore engine). Global ``settings.json`` only seeds the defaults for
*new* projects; once a project exists it remembers how it was last inspected.
"""

from __future__ import annotations

import json
import os
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path

from ..config import AppConfig

PROJECT_FILE = "project.json"
FRAMES_DIR = "frames"
CACHE_FILE = "detections.json"
COLLECTIONS_DIR = "collections"
FAVORITES_DIR = "Избранное"  # the single default collection ("в избранное")
RESTORED_DIR = "restored"    # batch "расцензурить все" output (kept out of frames/)
_VERSION = 1

# The subset of AppConfig fields a project remembers (mirrored to/from project.json).
_SETTING_KEYS = (
    "detector",
    "detector_models",
    "default_threshold",
    "model_thresholds",
    "cross_model_nms",
    "nms_iou",
    "restorer",
    "dm_dir",
    "dm_model",
    "dm_gpu",
    "dm_feed_restored",
    "diff_backend",
    "diff_url",
    "diff_model",
    "diff_prompt",
    "diff_negative",
    "diff_steps",
    "diff_cfg",
    "diff_denoise",
    "diff_seed",
    "diff_mask_dilate",
    "diff_mask_blur",
)


@dataclass
class Project:
    """A HVideoTool project rooted at ``root`` — the single source of truth for its
    on-disk layout and its per-project settings."""

    root: Path
    name: str
    settings: dict = field(default_factory=dict)  # subset of AppConfig fields
    source: str | None = None                     # originating video/folder, if any
    created: str | None = None

    # ----------------------------------------------------------------- paths
    @property
    def project_file(self) -> Path:
        return self.root / PROJECT_FILE

    @property
    def frames_dir(self) -> Path:
        return self.root / FRAMES_DIR

    @property
    def cache_path(self) -> Path:
        return self.root / CACHE_FILE

    @property
    def collections_dir(self) -> Path:
        return self.root / COLLECTIONS_DIR

    @property
    def favorites_dir(self) -> Path:
        """The single default collection — frames moved "to favorites" land here."""
        return self.collections_dir / FAVORITES_DIR

    @property
    def restored_dir(self) -> Path:
        """Batch restoration output ("расцензурить все") — mirrors frame basenames.
        Kept out of ``frames/`` so results aren't listed/re-detected/re-restored."""
        return self.root / RESTORED_DIR

    # ------------------------------------------------------------- lifecycle
    @classmethod
    def create(
        cls,
        root: Path | str,
        name: str | None = None,
        settings: dict | None = None,
        source: str | None = None,
    ) -> Project:
        """Create a new project folder (with ``frames/``) and write ``project.json``."""
        root = Path(root)
        proj = cls(
            root=root,
            name=name or root.name,
            settings={k: v for k, v in (settings or {}).items() if k in _SETTING_KEYS},
            source=source,
            created=datetime.now().isoformat(timespec="seconds"),
        )
        proj.frames_dir.mkdir(parents=True, exist_ok=True)
        proj.save()
        return proj

    @classmethod
    def load(cls, path: Path | str) -> Project:
        """Load a project from its folder or directly from its ``project.json``."""
        path = Path(path)
        root = path.parent if path.name == PROJECT_FILE else path
        data = json.loads((root / PROJECT_FILE).read_text(encoding="utf-8"))
        return cls(
            root=root,
            name=data.get("name", root.name),
            settings={k: v for k, v in data.get("settings", {}).items() if k in _SETTING_KEYS},
            source=data.get("source"),
            created=data.get("created"),
        )

    @staticmethod
    def is_project(path: Path | str) -> bool:
        """True if ``path`` is a project folder (or a ``project.json``)."""
        path = Path(path)
        if path.name == PROJECT_FILE:
            return path.is_file()
        return (path / PROJECT_FILE).is_file()

    def save(self) -> None:
        """Write ``project.json``."""
        payload = {
            "version": _VERSION,
            "name": self.name,
            "created": self.created,
            "source": self.source,
            "settings": self.settings,
        }
        self.root.mkdir(parents=True, exist_ok=True)
        # Crash-safe write (sibling .tmp + atomic os.replace): never truncate a good
        # project.json if the app/PC dies mid-save.
        text = json.dumps(payload, ensure_ascii=False, indent=2)
        tmp = self.project_file.with_name(PROJECT_FILE + ".tmp")
        try:
            tmp.write_text(text, encoding="utf-8")
            os.replace(tmp, self.project_file)
        finally:
            if tmp.exists():
                try:
                    tmp.unlink()
                except OSError:
                    pass

    # ----------------------------------------------------- settings <-> config
    def apply_to_config(self, cfg: AppConfig) -> None:
        """Overlay this project's stored settings onto ``cfg`` (mutates it)."""
        for key in _SETTING_KEYS:
            if key in self.settings:
                setattr(cfg, key, self.settings[key])

    def update_from_config(self, cfg: AppConfig) -> None:
        """Capture the per-project settings from ``cfg`` into ``self.settings``."""
        self.settings = {key: getattr(cfg, key) for key in _SETTING_KEYS}