adam/dataset_registry.py

from __future__ import annotations

import hashlib
import json
import threading
from dataclasses import asdict, dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, TYPE_CHECKING
from uuid import uuid4

from adam.dataset_lab import IMAGE_EXTENSIONS, TEXT_EXTENSIONS, VIDEO_EXTENSIONS

if TYPE_CHECKING:
    from adam.assets import Asset, AssetRegistry


DATASET_MARKERS = {
    "dataset_manifest.json",
    "metadata",
    "frames",
    "videos",
    "captions",
    "actions.jsonl",
    "actions.csv",
}
SCAN_LIMIT = 20_000
ASYNC_REFRESH_AFTER = timedelta(minutes=30)
_scan_lock = threading.Lock()
_active_scans: set[str] = set()


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _path_key(path: str | Path) -> str:
    resolved = str(Path(path).expanduser().resolve())
    return hashlib.sha1(resolved.casefold().encode("utf-8")).hexdigest()[:16]


def _atomic_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2), encoding="utf-8")
    temporary.replace(path)


def _date(value: str) -> datetime | None:
    try:
        return datetime.fromisoformat(value)
    except (TypeError, ValueError):
        return None


@dataclass(slots=True)
class DatasetLocation:
    id: str
    name: str
    path: str
    source: str = "user"
    created_at: str = field(default_factory=_now)
    last_seen_at: str = ""
    exists: bool = True

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "DatasetLocation":
        return cls(
            id=str(payload.get("id") or _path_key(str(payload.get("path", "")))),
            name=str(payload.get("name") or Path(str(payload.get("path", ""))).name or "Datasets"),
            path=str(payload.get("path", "")),
            source=str(payload.get("source") or "user"),
            created_at=str(payload.get("created_at") or _now()),
            last_seen_at=str(payload.get("last_seen_at") or ""),
            exists=bool(payload.get("exists", True)),
        )


@dataclass(slots=True)
class DatasetRecord:
    id: str
    name: str
    path: str
    source: str = "asset"
    location_id: str = ""
    favorite: bool = False
    last_used_at: str = ""
    discovered_at: str = field(default_factory=_now)
    scanned_at: str = ""
    exists: bool = True
    item_count: int = 0
    image_count: int = 0
    video_count: int = 0
    caption_count: int = 0
    missing_caption_count: int = 0
    sample_image: str = ""
    dataset_format: str = "Unknown"
    warnings: list[str] = field(default_factory=list)

    @classmethod
    def from_dict(cls, payload: dict[str, Any]) -> "DatasetRecord":
        return cls(
            id=str(payload.get("id") or _path_key(str(payload.get("path", "")))),
            name=str(payload.get("name") or Path(str(payload.get("path", ""))).name or "Dataset"),
            path=str(payload.get("path", "")),
            source=str(payload.get("source") or "asset"),
            location_id=str(payload.get("location_id") or ""),
            favorite=bool(payload.get("favorite", False)),
            last_used_at=str(payload.get("last_used_at") or ""),
            discovered_at=str(payload.get("discovered_at") or _now()),
            scanned_at=str(payload.get("scanned_at") or ""),
            exists=bool(payload.get("exists", True)),
            item_count=int(payload.get("item_count", 0) or 0),
            image_count=int(payload.get("image_count", 0) or 0),
            video_count=int(payload.get("video_count", 0) or 0),
            caption_count=int(payload.get("caption_count", 0) or 0),
            missing_caption_count=int(payload.get("missing_caption_count", 0) or 0),
            sample_image=str(payload.get("sample_image") or ""),
            dataset_format=str(payload.get("dataset_format") or "Unknown"),
            warnings=[str(item) for item in payload.get("warnings", []) if str(item)],
        )


class DatasetRegistry:
    """Persistent ADAM-aware index of known dataset locations and datasets."""

    def __init__(self, root: Path, config: Any | None = None) -> None:
        self.root = root.resolve()
        self.path = self.root / "data" / "dataset_registry.json"
        self.config = config
        self.locations: list[DatasetLocation] = []
        self.datasets: dict[str, DatasetRecord] = {}
        self.load()

    def load(self) -> None:
        try:
            payload = json.loads(self.path.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            payload = {}
        self.locations = [
            DatasetLocation.from_dict(item)
            for item in payload.get("locations", [])
            if isinstance(item, dict)
        ]
        self.datasets = {
            str(key): DatasetRecord.from_dict(value)
            for key, value in (payload.get("datasets", {}) or {}).items()
            if isinstance(value, dict)
        }

    def save(self) -> None:
        _atomic_json(
            self.path,
            {
                "locations": [asdict(item) for item in self.locations],
                "datasets": {key: asdict(value) for key, value in self.datasets.items()},
            },
        )

    def register_location(self, path: str | Path, *, name: str = "", source: str = "user") -> DatasetLocation:
        resolved = Path(path).expanduser().resolve()
        if not resolved.is_dir():
            raise ValueError("Choose an existing dataset folder.")
        key = _path_key(resolved)
        existing = next((item for item in self.locations if item.id == key), None)
        if existing is None:
            existing = DatasetLocation(
                id=key,
                name=name.strip() or resolved.name or str(resolved),
                path=str(resolved),
                source=source,
            )
            self.locations.insert(0, existing)
        else:
            existing.name = name.strip() or existing.name
            existing.path = str(resolved)
            existing.source = source or existing.source
        existing.exists = True
        existing.last_seen_at = _now()
        self.save()
        return existing

    def remove_location(self, location_id: str) -> bool:
        before = len(self.locations)
        self.locations = [item for item in self.locations if item.id != location_id]
        changed = len(self.locations) != before
        if changed:
            self.save()
        return changed

    def known_locations(self) -> list[DatasetLocation]:
        locations = list(self.locations)
        by_path = {Path(item.path).expanduser().resolve(): item for item in locations if item.path}
        for path, name, source in self._automatic_location_candidates():
            try:
                resolved = path.expanduser().resolve()
            except OSError:
                continue
            if resolved in by_path:
                continue
            locations.append(
                DatasetLocation(
                    id=_path_key(resolved),
                    name=name or resolved.name or str(resolved),
                    path=str(resolved),
                    source=source,
                    exists=resolved.is_dir(),
                    last_seen_at=_now() if resolved.is_dir() else "",
                )
            )
        return locations

    def discover_into_assets(self, assets: "AssetRegistry", *, persist: bool = False, update_cache: bool = True) -> list["Asset"]:
        discovered: list[Asset] = []
        records = self.discover(asset_registry=assets, refresh_missing=False, persist=update_cache)
        for record in records:
            if not record.exists:
                continue
            asset = assets.register(
                kind="dataset",
                name=record.name,
                path=record.path,
                metadata={
                    "dataset_registry_source": record.source,
                    "dataset_location_id": record.location_id,
                },
                persist=False,
            )
            discovered.append(asset)
        if persist and discovered:
            assets.save()
        return discovered

    def discover(
        self,
        *,
        asset_registry: "AssetRegistry | None" = None,
        refresh_missing: bool = True,
        persist: bool = True,
    ) -> list[DatasetRecord]:
        self.load()
        changed = False
        locations = self.known_locations()
        known_location_ids = {location.id for location in locations}
        for key, record in list(self.datasets.items()):
            if record.source in {"adam", "tool"} and record.location_id and record.location_id not in known_location_ids:
                del self.datasets[key]
                changed = True
        for location in locations:
            exists = Path(location.path).is_dir()
            if location.source == "user":
                stored = next((item for item in self.locations if item.id == location.id), None)
                if stored:
                    stored.exists = exists
                    stored.last_seen_at = _now() if exists else stored.last_seen_at
                    changed = True
            if not exists:
                continue
            for candidate in self._dataset_candidates(Path(location.path)):
                record = self._cached_or_sampled(candidate, source=location.source, location_id=location.id)
                self.datasets[record.id] = record
                changed = True
        if asset_registry is not None:
            for asset in getattr(asset_registry, "assets", []):
                if getattr(asset, "kind", "") != "dataset":
                    continue
                record = self._cached_or_sampled(Path(asset.path), source="asset", location_id="")
                record.name = asset.name or record.name
                self.datasets[record.id] = record
                changed = True
        for record in self.datasets.values():
            record.exists = Path(record.path).is_dir()
            if refresh_missing and record.exists and self._needs_refresh(record):
                self.refresh_async(record.path, source=record.source, location_id=record.location_id)
        if changed and persist:
            self.save()
        return self.sorted_records()

    def sorted_records(self) -> list[DatasetRecord]:
        records = list(self.datasets.values())
        records.sort(
            key=lambda item: (
                not item.favorite,
                not bool(item.last_used_at),
                item.last_used_at or item.discovered_at,
                item.name.casefold(),
            ),
            reverse=False,
        )
        favorites = sorted([item for item in records if item.favorite], key=lambda item: item.name.casefold())
        recent = sorted(
            [item for item in records if not item.favorite and item.last_used_at],
            key=lambda item: item.last_used_at,
            reverse=True,
        )
        others = sorted(
            [item for item in records if not item.favorite and not item.last_used_at],
            key=lambda item: item.discovered_at,
            reverse=True,
        )
        return [*favorites, *recent, *others]

    def record_for_path(self, path: str | Path) -> DatasetRecord:
        key = _path_key(path)
        record = self.datasets.get(key)
        if record is None:
            record = self._cached_or_sampled(Path(path), source="asset", location_id="")
            self.datasets[key] = record
            self.save()
        return record

    def favorite(self, path: str | Path, enabled: bool) -> DatasetRecord:
        record = self.record_for_path(path)
        record.favorite = bool(enabled)
        self.save()
        return record

    def touch(self, path: str | Path) -> DatasetRecord:
        record = self.record_for_path(path)
        record.last_used_at = _now()
        self.save()
        return record

    def refresh_async(self, path: str | Path, *, source: str = "asset", location_id: str = "") -> None:
        resolved = str(Path(path).expanduser().resolve())
        key = _path_key(resolved)
        with _scan_lock:
            if key in _active_scans:
                return
            _active_scans.add(key)

        def worker() -> None:
            try:
                record = self._scan(Path(resolved), source=source, location_id=location_id, limit=SCAN_LIMIT)
                fresh = DatasetRegistry(self.root, self.config)
                current = fresh.datasets.get(record.id)
                if current:
                    record.favorite = current.favorite
                    record.last_used_at = current.last_used_at
                    record.discovered_at = current.discovered_at
                fresh.datasets[record.id] = record
                fresh.save()
            finally:
                with _scan_lock:
                    _active_scans.discard(key)

        threading.Thread(target=worker, name="ADAMDatasetScan", daemon=True).start()

    def _automatic_location_candidates(self) -> list[tuple[Path, str, str]]:
        candidates: list[tuple[Path, str, str]] = [(self.root / "ADAM_Datasets", "ADAM Datasets", "adam")]
        folders = self.config.get("tool_folders", {}) if self.config is not None else {}
        folders = folders if isinstance(folders, dict) else {}
        collector_raw = str(folders.get("dataset_collector", "")).strip()
        if collector_raw:
            candidates.append((Path(collector_raw) / "Datasets", "Dataset Collector", "tool"))
        for tool_id in ("ddpm_trainer", "lora_trainer", "flow_trainer", "oasis_trainer"):
            raw_root = str(folders.get(tool_id, "")).strip()
            if not raw_root:
                continue
            root = Path(raw_root)
            label = tool_id.removesuffix("_trainer").upper()
            for child in ("Datasets", "datasets", "OldDatasets"):
                candidates.append((root / child, f"{label} {child}", "tool"))
        external_tools = self.root / "config" / "external_tools.json"
        try:
            payload = json.loads(external_tools.read_text(encoding="utf-8"))
        except (OSError, ValueError, TypeError, json.JSONDecodeError):
            payload = {}
        for entry in payload.get("tools", []) if isinstance(payload, dict) else []:
            if not isinstance(entry, dict):
                continue
            backend = entry.get("backend", {})
            if not isinstance(backend, dict):
                continue
            raw_root = str(backend.get("root", "")).strip()
            if raw_root:
                root = Path(raw_root)
                candidates.append((root / "Datasets", f"{entry.get('name', 'External tool')} Datasets", "tool"))
                candidates.append((root / "OldDatasets", f"{entry.get('name', 'External tool')} OldDatasets", "tool"))
        return candidates

    def _dataset_candidates(self, location: Path) -> list[Path]:
        candidates: list[Path] = []
        try:
            children = [item for item in location.iterdir() if item.is_dir()]
        except OSError:
            children = []
        for child in children[:1000]:
            if self._looks_like_dataset(child):
                candidates.append(child)
        if candidates:
            return candidates
        if self._looks_like_dataset(location):
            candidates.append(location)
        return candidates

    def _looks_like_dataset(self, folder: Path) -> bool:
        if not folder.is_dir():
            return False
        try:
            names = {item.name for item in folder.iterdir()}
        except OSError:
            return False
        if names & DATASET_MARKERS:
            return True
        checked = 0
        for item in folder.rglob("*"):
            if checked >= 200:
                break
            checked += 1
            if item.is_file() and item.suffix.casefold() in IMAGE_EXTENSIONS | VIDEO_EXTENSIONS:
                return True
        return False

    def _cached_or_sampled(self, path: Path, *, source: str, location_id: str) -> DatasetRecord:
        key = _path_key(path)
        current = self.datasets.get(key)
        if current is not None:
            current.exists = path.is_dir()
            current.source = current.source or source
            current.location_id = current.location_id or location_id
            return current
        return self._scan(path, source=source, location_id=location_id, limit=800)

    def _needs_refresh(self, record: DatasetRecord) -> bool:
        if not record.scanned_at:
            return True
        scanned = _date(record.scanned_at)
        return scanned is None or datetime.now(timezone.utc) - scanned > ASYNC_REFRESH_AFTER

    def _scan(self, path: Path, *, source: str, location_id: str, limit: int) -> DatasetRecord:
        resolved = path.expanduser().resolve()
        record = DatasetRecord(
            id=_path_key(resolved),
            name=resolved.name or "Dataset",
            path=str(resolved),
            source=source,
            location_id=location_id,
            exists=resolved.is_dir(),
            scanned_at=_now(),
        )
        if not resolved.is_dir():
            record.warnings.append("Dataset folder is unavailable.")
            return record
        files_seen = 0
        capped = False
        try:
            iterator = resolved.rglob("*")
            for item in iterator:
                if not item.is_file():
                    continue
                files_seen += 1
                suffix = item.suffix.casefold()
                if suffix in IMAGE_EXTENSIONS:
                    record.image_count += 1
                    if not record.sample_image:
                        record.sample_image = str(item)
                    if not any(item.with_suffix(ext).is_file() for ext in (".txt", ".caption")):
                        record.missing_caption_count += 1
                elif suffix in VIDEO_EXTENSIONS:
                    record.video_count += 1
                if suffix in TEXT_EXTENSIONS and suffix in {".txt", ".caption"}:
                    record.caption_count += 1
                if files_seen >= limit:
                    capped = True
                    break
        except OSError as exc:
            record.warnings.append(f"Dataset could not be scanned: {exc}")
        record.item_count = record.image_count + record.video_count
        record.dataset_format = self._format_label(resolved, record)
        if capped:
            record.warnings.append(f"Counts are sampled from the first {limit:,} files.")
        return record

    @staticmethod
    def _format_label(folder: Path, record: DatasetRecord) -> str:
        if (folder / "actions.jsonl").is_file() or (folder / "actions.csv").is_file():
            return "Oasis action dataset"
        if record.video_count:
            return "Video dataset"
        if record.image_count and record.caption_count:
            return "Captioned image dataset"
        if record.image_count:
            return "Image dataset"
        return "Dataset folder"