"""The whole Crello corpus, served from a Hugging Face bucket.

    FineEnvs/crello-bucket                     (cyberagent/crello 5.1.0, copied as is)
      data/<split>-000NN-of-000MM.parquet      38 shards, 100-design row groups, 18 GB
      resources/fonts.pickle
      openenv/index/<snapshot>/catalog.sqlite  built here, published next to the data

**Index once.** `build_index` converts every design exactly as the pack builder
does, keeps the ones that make a playable task, and stores the finished task
record (ground truth, asset list, brief, difficulty features) plus where the
design lives (shard, row group, row). That is the whole catalog: listing tasks,
`get_task` and reset never touch the images.

**Fetch on demand.** Asset pixels are only needed once an episode renders. The
first time a design is used its row group is read (from a mount at `/corpus`
on a Space, over HTTP locally), the design is converted again and checked
against the indexed record, and its assets land in a byte-bounded disk cache.

**Pinned identity.** The snapshot id hashes the index version, the conversion
rules and each shard's content hash. A shard that changed since indexing is
refused rather than silently serving different ground truth.
"""

from __future__ import annotations

import fcntl
import hashlib
import json
import os
import shutil
import sqlite3
import threading
import time
import zlib
from collections import OrderedDict
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from contextlib import closing
from pathlib import Path

from .difficulty import Scorer, features, shape
from .pack import SPLITS, TaskPack, accept, make_task, write_design

INDEX_VERSION = 1
DEFAULT_BUCKET = "FineEnvs/crello-bucket"
DEFAULT_MAX_ELEMENTS = 30
INDEX_PREFIX = "openenv/index"
FONTS_PATH = "resources/fonts.pickle"
SOURCE = {"repo": "cyberagent/crello", "revision": "5.1.0",
          "license": "CDLA-Permissive-2.0; templates (c) VistaCreate, see dataset card"}

DDL = """
CREATE TABLE tasks(
    task_id TEXT PRIMARY KEY, split TEXT NOT NULL, position INTEGER NOT NULL,
    crello_id TEXT NOT NULL UNIQUE, file TEXT NOT NULL, row_group INTEGER NOT NULL,
    row INTEGER NOT NULL, public TEXT NOT NULL, payload BLOB NOT NULL,
    features TEXT NOT NULL, score REAL NOT NULL, tier TEXT NOT NULL, shape TEXT NOT NULL,
    doc_sha256 TEXT NOT NULL, UNIQUE(split, position));
CREATE TABLE rejected(crello_id TEXT, split TEXT, reason TEXT);
CREATE TABLE meta(key TEXT PRIMARY KEY, value TEXT NOT NULL);
"""


def canonical(obj) -> str:
    return json.dumps(obj, sort_keys=True, separators=(",", ":"), ensure_ascii=False)


def sha256_bytes(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def sha256_file(path: Path) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as fh:
        for chunk in iter(lambda: fh.read(1 << 22), b""):
            h.update(chunk)
    return h.hexdigest()


def doc_digest(doc: dict) -> str:
    return sha256_bytes(canonical(doc).encode())


# ------------------------------------------------------------------ source


class BucketSource:
    """Read-only access to the bucket, from a mounted/copied directory or over HTTP."""

    def __init__(self, bucket_id: str = DEFAULT_BUCKET, root: str | Path | None = None):
        self.bucket_id = os.environ.get("DESIGN_CANVAS_BUCKET_ID") or bucket_id
        self.root = Path(root) if root else None
        self._fs = None

    @property
    def fs(self):
        if self._fs is None:
            from huggingface_hub import HfFileSystem

            self._fs = HfFileSystem()
        return self._fs

    def url(self, path: str) -> str:
        return f"hf://buckets/{self.bucket_id}/{path}"

    def open(self, path: str):
        if self.root:
            return open(self.root / path, "rb")
        return self.fs.open(self.url(path), "rb")

    def inventory(self) -> list[dict]:
        """Every parquet shard with its size and content hash (the bucket's xet hash)."""
        if self.root:
            return [{"path": str(p.relative_to(self.root)), "size": p.stat().st_size,
                     "sha256": sha256_file(p)}
                    for p in sorted((self.root / "data").glob("*.parquet"))]
        from huggingface_hub import HfApi

        return [{"path": f.path, "size": f.size, "xet_hash": f.xet_hash}
                for f in sorted(HfApi().list_bucket_tree(self.bucket_id, prefix="data", recursive=True),
                                key=lambda f: f.path)
                if f.path.endswith(".parquet")]

    def verify(self, spec: dict) -> None:
        """Refuse a shard whose content no longer matches the indexed one."""
        if self.root:
            path = self.root / spec["path"]
            if path.stat().st_size != spec["size"]:
                raise ValueError(f"{spec['path']} changed since it was indexed (size)")
            return
        if not spec.get("xet_hash"):
            raise ValueError(f"{spec['path']} was indexed from a local copy; serve it with a source root")
        from huggingface_hub import HfApi

        rows = list(HfApi().get_bucket_paths_info(self.bucket_id, [spec["path"]]))
        if len(rows) != 1 or rows[0].xet_hash != spec["xet_hash"] or rows[0].size != spec["size"]:
            raise ValueError(f"{spec['path']} changed in {self.bucket_id} since it was indexed; "
                             "rebuild the index rather than serve different ground truth")

    def fetch(self, path: str, dest: Path) -> Path:
        dest.parent.mkdir(parents=True, exist_ok=True)
        tmp = dest.with_suffix(dest.suffix + f".{os.getpid()}.part")  # workers may fetch the same file at once
        with self.open(path) as src, open(tmp, "wb") as out:
            shutil.copyfileobj(src, out, 1 << 22)
        tmp.replace(dest)
        return dest


def split_of(path: str) -> str:
    name = Path(path).name
    for split in SPLITS:
        if name.startswith(split + "-"):
            return split
    raise ValueError(f"cannot tell the split of {path}")


def crello_features(parquet_file):
    import datasets

    return datasets.Features.from_arrow_schema(parquet_file.schema_arrow)


# ------------------------------------------------------------------ index build


def _index_shard(args) -> list[dict]:
    """Worker: convert every design in one shard; return one record per design."""
    import pyarrow.parquet as pq

    from .crello import convert_row
    from .fonts import FontRegistry

    bucket_id, root, path, fonts_dir, max_elements = args
    source = BucketSource(bucket_id, root)
    fonts = FontRegistry(fonts_dir)
    split = split_of(path)
    out = []
    with source.open(path) as fh:
        pf = pq.ParquetFile(fh, pre_buffer=True)
        feats = crello_features(pf)
        cols = [n for n in pf.schema_arrow.names if n != "preview"]  # 29% of the bytes, not needed here
        for rg in range(pf.num_row_groups):
            for row_i, row in enumerate(pf.read_row_group(rg, columns=cols).to_pylist()):
                design = convert_row(row, feats, fonts)
                rec = {"crello_id": design.crello_id, "split": split, "file": path,
                       "row_group": rg, "row": row_i}
                reason = accept(design, max_elements)
                if reason:
                    out.append({**rec, "rejected": reason})
                    continue
                task = make_task(design, split, 0)
                for a in task["assets"]:
                    img = design.assets[a["id"]]
                    a["sha256"] = sha256_bytes(img.tobytes() + repr(img.size).encode())
                f = features(design.doc)
                out.append({**rec, "task": task, "features": f,
                            "cluster": row.get("cluster_index"),
                            "group": feats["group"].names[row["group"]] if row.get("group") is not None else None})
    return out


def build_index(out: Path | str, source: BucketSource, *, max_elements: int = DEFAULT_MAX_ELEMENTS,
                workers: int = 8, shards: list[str] | None = None, log=print) -> dict:
    """Convert the whole corpus into `out/catalog.sqlite` + `out/manifest.json`."""
    from .fonts import extract_crello_fonts

    out = Path(out)
    out.mkdir(parents=True, exist_ok=True)
    t0 = time.time()
    inventory = [s for s in source.inventory() if shards is None or s["path"] in shards]
    if not inventory:
        raise ValueError("no parquet shards found in the source")
    fonts_dir = out / "fonts"
    if not (fonts_dir / "manifest.json").exists():
        pickle_path = source.fetch(FONTS_PATH, out / "fonts.pickle") if not source.root \
            else source.root / FONTS_PATH
        extract_crello_fonts(pickle_path, fonts_dir)
    log(f"indexing {len(inventory)} shards with {workers} workers")
    jobs = [(source.bucket_id, str(source.root) if source.root else None, s["path"], str(fonts_dir),
             max_elements) for s in inventory]
    records = []
    with ProcessPoolExecutor(workers) as ex:
        for i, recs in enumerate(ex.map(_index_shard, jobs), 1):
            records.extend(recs)
            log(f"  {i}/{len(jobs)} shards, {len(records)} designs, {time.time() - t0:.0f}s")

    kept = [r for r in records if "task" in r]
    scorer = Scorer([r["features"] for r in kept])
    counts, rejected = {}, {}
    db_path = out / "catalog.sqlite"
    db_path.unlink(missing_ok=True)
    with closing(sqlite3.connect(db_path)) as db:
        db.executescript(DDL)
        for split in SPLITS:
            mine = sorted((r for r in kept if r["split"] == split),
                          key=lambda r: (r["file"], r["row_group"], r["row"]))
            counts[split] = len(mine)
            for pos, r in enumerate(mine):
                task = r["task"]
                task.update(task_id=f"{split}-{pos:05d}", index=pos)
                score = scorer.score(r["features"])
                tier = scorer.tier(score)
                task["difficulty"] = tier
                task["difficulty_score"] = score
                task["meta"] = {"cluster": r["cluster"], "group": r["group"]}
                pub = {**TaskPack.public(task), "difficulty_score": score,
                       "shape": shape(r["features"]["aspect"])}
                db.execute("INSERT INTO tasks VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?)", (
                    task["task_id"], split, pos, r["crello_id"], r["file"], r["row_group"], r["row"],
                    canonical(pub), zlib.compress(canonical(task).encode(), 6),
                    canonical({**r["features"], "cluster": r["cluster"], "group": r["group"],
                               "category": task["category"], "format": task["format"]}),
                    score, tier, pub["shape"], doc_digest(task["doc"])))
        for r in records:
            if "rejected" in r:
                db.execute("INSERT INTO rejected VALUES(?,?,?)", (r["crello_id"], r["split"], r["rejected"]))
                rejected[r["rejected"]] = rejected.get(r["rejected"], 0) + 1
        db.commit()
        db.execute("VACUUM")
    config = {"max_elements": max_elements, "index_version": INDEX_VERSION}
    snapshot = sha256_bytes(canonical([INDEX_VERSION, config, scorer.describe(),
                                       [[s["path"], s["size"], s.get("xet_hash") or s.get("sha256")]
                                        for s in inventory]]).encode())
    fonts_sha = sha256_file(out / "fonts.pickle") if (out / "fonts.pickle").exists() else \
        sha256_file(source.root / FONTS_PATH)
    manifest = {
        "status": "ready",
        "storage": "bucket-parquet",
        "index_version": INDEX_VERSION,
        "snapshot_id": snapshot,
        "created": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "bucket_id": source.bucket_id,
        "source": SOURCE,
        "config": config,
        "difficulty": scorer.describe(),
        "splits": counts,
        "rejected": dict(sorted(rejected.items(), key=lambda kv: -kv[1])),
        "fonts": {"path": FONTS_PATH, "sha256": fonts_sha},
        "catalog": {"path": f"{INDEX_PREFIX}/{snapshot[:16]}/catalog.sqlite",
                    "size": db_path.stat().st_size, "sha256": sha256_file(db_path)},
        "source_files": inventory,
    }
    (out / "manifest.json").write_text(json.dumps(manifest, indent=2) + "\n")
    log(f"index ready: {counts}, rejected {sum(rejected.values())}, {time.time() - t0:.0f}s")
    return manifest


def publish_index(out: Path | str, bucket_id: str | None = None) -> str:
    """Upload catalog + manifest next to the data. Returns the manifest's hf:// URL."""
    from huggingface_hub import HfApi

    out = Path(out)
    manifest = json.loads((out / "manifest.json").read_text())
    bucket = bucket_id or manifest["bucket_id"]
    prefix = manifest["catalog"]["path"].rsplit("/", 1)[0]
    HfApi().batch_bucket_files(bucket, add=[
        (str(out / "catalog.sqlite"), manifest["catalog"]["path"]),
        (str(out / "manifest.json"), f"{prefix}/manifest.json"),
    ])
    return f"hf://buckets/{bucket}/{prefix}/manifest.json"


# ------------------------------------------------------------------ disk cache


class DiskLRU:
    """Byte-bounded directory cache. Entries used in the last `min_age` seconds are never evicted.

    Safe to share between server replicas on one machine: an entry is filled in a private
    temporary directory and renamed into place (the first finished copy wins), and "last
    used" is the mtime of the entry's `.ok` marker, so a replica never evicts an entry that
    another one touched recently.
    """

    _TMP = ".tmp-"

    def __init__(self, root: Path, budget: int, min_age: float = 900.0):
        self.root = Path(root)
        self.root.mkdir(parents=True, exist_ok=True)
        self.budget = budget
        self.min_age = min_age
        self._lock = threading.Lock()
        self._keys: dict[str, threading.Lock] = {}
        self._sizes: OrderedDict[str, int] = OrderedDict()
        for p in sorted(self.root.iterdir(), key=lambda p: p.stat().st_mtime):
            if p.name.startswith(self._TMP):
                if time.time() - p.stat().st_mtime > 3600:  # abandoned by a crashed writer
                    shutil.rmtree(p, ignore_errors=True)
            elif (p / ".ok").exists():
                self._sizes[p.name] = int((p / ".ok").read_text() or 0)
            elif p.is_dir():
                shutil.rmtree(p, ignore_errors=True)  # a write from before fills were atomic

    def _hit(self, key: str, path: Path) -> bool:
        ok = path / ".ok"
        try:
            size = int(ok.read_text() or 0)
            os.utime(ok)  # last used, visible to every replica
        except (OSError, ValueError):
            return False
        with self._lock:
            self._sizes[key] = size
            self._sizes.move_to_end(key)
        return True

    def get(self, key: str, fill) -> Path:
        path = self.root / key
        with self._lock:
            lock = self._keys.setdefault(key, threading.Lock())
        with lock:
            if self._hit(key, path):
                return path
            tmp = self.root / f"{self._TMP}{key}-{os.getpid()}-{threading.get_ident()}"
            shutil.rmtree(tmp, ignore_errors=True)
            tmp.mkdir(parents=True)
            try:
                fill(tmp)
                size = sum(f.stat().st_size for f in tmp.rglob("*") if f.is_file())
                (tmp / ".ok").write_text(str(size))
                if path.exists() and not (path / ".ok").exists():
                    shutil.rmtree(path, ignore_errors=True)  # a pre-atomic leftover
                try:
                    tmp.rename(path)
                except OSError:  # another replica finished first: use its copy
                    if not self._hit(key, path):
                        raise
                    return path
            finally:
                shutil.rmtree(tmp, ignore_errors=True)
            with self._lock:
                self._sizes[key] = size
                self._evict(keep=key)
            return path

    def _evict(self, keep: str) -> None:
        now = time.time()
        total = sum(self._sizes.values())
        if total <= self.budget:
            return
        used = {}
        for key in list(self._sizes):
            try:
                used[key] = (self.root / key / ".ok").stat().st_mtime
            except OSError:  # already evicted by another replica
                total -= self._sizes.pop(key)
        for key in sorted(used, key=used.get):  # least recently used by any replica first
            if total <= self.budget:
                break
            if key == keep or now - used[key] < self.min_age:
                continue
            total -= self._sizes.pop(key)
            trash = self.root / f"{self._TMP}evict-{key}-{os.getpid()}"
            try:
                (self.root / key).rename(trash)  # atomic: readers see it whole or not at all
            except OSError:
                continue
            shutil.rmtree(trash, ignore_errors=True)

    def stats(self) -> dict:
        with self._lock:
            return {"entries": len(self._sizes), "bytes": sum(self._sizes.values()), "budget": self.budget}


# ------------------------------------------------------------------ serving


def _read_manifest(value: str | Path) -> tuple[dict, str]:
    """Manifest + the directory (local or hf://) it came from."""
    value = str(value)
    if value.startswith("hf://"):
        from huggingface_hub import HfFileSystem

        with HfFileSystem().open(value, "rb") as fh:
            return json.loads(fh.read()), value.rsplit("/", 1)[0]
    path = Path(value).expanduser()
    if path.is_dir():
        path = path / "manifest.json"
    return json.loads(path.read_text()), str(path.parent)


class CorpusPack(TaskPack):
    """The TaskPack interface over the full bucket-backed corpus.

    Everything the server calls works unchanged: task records come from the
    index, and any call that needs pixels (`asset_paths`, `preview_path`,
    renders, grading) materialises the design first.
    """

    def __init__(self, manifest: str | Path, cache_dir: str | Path | None = None, *,
                 source_root: str | Path | None = None, design_cache_bytes: int = 4 << 30,
                 group_cache_bytes: int = 2 << 30):
        from .fonts import FontRegistry, extract_crello_fonts
        from .render import Renderer

        self.manifest, where = _read_manifest(manifest)
        if self.manifest.get("storage") != "bucket-parquet" or \
                self.manifest.get("index_version") != INDEX_VERSION:
            raise ValueError("not a supported corpus manifest; rebuild the index")
        self.snapshot_id = self.manifest["snapshot_id"]
        cache = Path(cache_dir or os.environ.get("DESIGN_CANVAS_CACHE")
                     or Path.home() / ".cache/design_canvas/corpus").expanduser() / self.snapshot_id[:16]
        cache.mkdir(parents=True, exist_ok=True)
        self.cache_dir = cache
        root = source_root or os.environ.get("DESIGN_CANVAS_SOURCE_ROOT")
        self.source = BucketSource(self.manifest["bucket_id"], root)
        self._specs = {s["path"]: s for s in self.manifest["source_files"]}

        fonts_dir = cache / "fonts"
        # the multi-worker server starts several processes on one cache: the first copies the catalog and extracts
        # the fonts while the others wait, then they all find both in place
        with open(cache / ".setup.lock", "w") as lock:
            fcntl.flock(lock, fcntl.LOCK_EX)
            self._catalog = self._local_catalog(where)
            if not (fonts_dir / "manifest.json").exists():
                pickle_path = self.source.root / FONTS_PATH if self.source.root else \
                    self.source.fetch(FONTS_PATH, cache / "fonts.pickle")
                if sha256_file(pickle_path) != self.manifest["fonts"]["sha256"]:
                    raise ValueError("fonts.pickle does not match the indexed corpus")
                extract_crello_fonts(pickle_path, fonts_dir)
        self.root = cache  # task["assets"][i]["file"] and task["preview"] resolve here
        self.fonts = FontRegistry(fonts_dir)
        self.renderer = Renderer(self.fonts)
        self._docs = OrderedDict()
        self._renders = OrderedDict()
        self._tasks_cache: OrderedDict[str, dict] = OrderedDict()
        self._lock = threading.Lock()
        self._db_local = threading.local()
        self.designs = DiskLRU(cache / "designs", design_cache_bytes)
        self.groups = DiskLRU(cache / "groups", group_cache_bytes, min_age=60.0)
        self._features = None
        self._prefetch = ThreadPoolExecutor(2, thread_name_prefix="design-canvas-prefetch")

    def _local_catalog(self, where: str) -> Path:
        """SQLite is copied to local disk and checksummed; it is never queried over a mount."""
        spec = self.manifest["catalog"]
        local = self.cache_dir / "catalog.sqlite"
        if local.exists() and local.stat().st_size == spec["size"]:
            return local
        name = spec["path"].rsplit("/", 1)[1]
        if where.startswith("hf://"):
            from huggingface_hub import HfFileSystem

            tmp = local.with_suffix(".part")
            HfFileSystem().get_file(f"{where}/{name}", str(tmp))
        else:
            candidates = [Path(where) / name]
            if self.source.root:
                candidates.append(self.source.root / spec["path"])
            src = next((p for p in candidates if p.exists()), None)
            if src is None:
                return self._local_catalog(f"hf://buckets/{self.source.bucket_id}/{spec['path'].rsplit('/', 1)[0]}")
            tmp = local.with_suffix(".part")
            shutil.copyfile(src, tmp)
        if sha256_file(tmp) != spec["sha256"]:
            tmp.unlink()
            raise ValueError("catalog.sqlite does not match its manifest")
        tmp.replace(local)
        return local

    @property
    def db(self) -> sqlite3.Connection:
        conn = getattr(self._db_local, "conn", None)
        if conn is None:
            conn = sqlite3.connect(f"file:{self._catalog}?mode=ro", uri=True)
            self._db_local.conn = conn
        return conn

    # ------------------------------------------------------------ index

    def splits(self) -> list[str]:
        return [s for s in SPLITS if self.manifest["splits"].get(s)]

    def count(self, split: str) -> int:
        if split not in self.manifest["splits"]:
            raise KeyError(f"unknown split {split!r}; have {', '.join(self.splits())}")
        return int(self.manifest["splits"][split])

    def _task(self, row) -> dict:
        task_id, payload = row
        with self._lock:
            task = self._tasks_cache.get(task_id)
            if task is not None:
                self._tasks_cache.move_to_end(task_id)
                return task
        task = json.loads(zlib.decompress(payload))
        with self._lock:
            self._tasks_cache[task_id] = task
            while len(self._tasks_cache) > 4096:
                self._tasks_cache.popitem(last=False)
            while len(self._docs) > 1024:
                self._docs.popitem(last=False)
        return task

    def at(self, split: str, index: int) -> dict:
        n = self.count(split)
        if not 0 <= index < n:
            raise IndexError(f"{split} has {n} tasks; index {index} is out of range")
        row = self.db.execute("SELECT task_id, payload FROM tasks WHERE split=? AND position=?",
                              (split, index)).fetchone()
        return self._task(row)

    def get(self, task_id: str) -> dict:
        row = self.db.execute("SELECT task_id, payload FROM tasks WHERE task_id=?", (task_id,)).fetchone()
        if row is None:
            raise KeyError(f"unknown task_id {task_id!r}")
        return self._task(row)

    def public(self, task: dict) -> dict:
        row = self.db.execute("SELECT public FROM tasks WHERE task_id=?", (task["task_id"],)).fetchone()
        return json.loads(row[0])

    def task_range(self, split: str, start=None, stop=None) -> list[dict]:
        n = self.count(split)
        lo, hi, _ = slice(start, stop).indices(n)
        rows = self.db.execute("SELECT public FROM tasks WHERE split=? AND position>=? AND position<? "
                               "ORDER BY position", (split, lo, hi)).fetchall()
        return [json.loads(r[0]) for r in rows]

    def rows(self, split: str | None = None) -> list[dict]:
        """Every indexed design's selection features (for building evaluation sets)."""
        q = "SELECT task_id, split, crello_id, features, score, tier, shape, doc_sha256 FROM tasks"
        args = ()
        if split:
            q += " WHERE split=?"
            args = (split,)
        return [{"task_id": t, "split": s, "crello_id": c, **json.loads(f), "score": sc,
                 "tier": tier, "shape": sh, "doc_sha256": d}
                for t, s, c, f, sc, tier, sh, d in self.db.execute(q + " ORDER BY split, position", args)]

    # --------------------------------------------------------- materialise

    def _location(self, task_id: str):
        return self.db.execute("SELECT file, row_group, row, crello_id, doc_sha256 FROM tasks "
                               "WHERE task_id=?", (task_id,)).fetchone()

    def _group(self, file: str, rg: int) -> Path:
        spec = self._specs[file]
        key = f"{(spec.get('xet_hash') or spec.get('sha256'))[:16]}-{rg}"

        def fill(dest: Path):
            import pyarrow as pa
            import pyarrow.parquet as pq

            self.source.verify(spec)
            with self.source.open(file) as fh:
                pf = pq.ParquetFile(fh, pre_buffer=True)
                # Crello's own preview is 29% of the bytes and only the playground shows
                # it, so it is fetched separately, on request (`preview_path`).
                cols = [n for n in pf.schema_arrow.names if n != "preview"]
                table = pf.read_row_group(rg, columns=cols)
            with pa.OSFile(str(dest / "group.arrow"), "wb") as sink, \
                    pa.ipc.new_file(sink, table.schema) as writer:
                writer.write_table(table)

        return self.groups.get(key, fill) / "group.arrow"

    def materialize(self, task: dict) -> Path:
        """Make sure this design's assets and preview are on disk; return its directory."""
        file, rg, row_i, crello_id, digest = self._location(task["task_id"])

        def fill(dest: Path):
            import pyarrow as pa

            from .crello import convert_row

            with pa.memory_map(str(self._group(file, rg))) as src:
                table = pa.ipc.open_file(src).read_all()
            row = table.slice(row_i, 1).to_pylist()[0]
            if str(row["id"]) != crello_id:
                raise ValueError(f"{task['task_id']}: row {file}:{rg}:{row_i} is not design {crello_id}")
            if self._features is None:
                import datasets

                self._features = datasets.Features.from_arrow_schema(table.schema)
            design = convert_row(row, self._features, self.fonts)
            if doc_digest(design.doc.model_dump(mode="json", exclude_defaults=False)) != digest:
                raise ValueError(f"{task['task_id']}: converted design differs from the index; "
                                 "the converter changed since indexing, rebuild the index")
            for a in task["assets"]:
                img = design.assets[a["id"]]
                if sha256_bytes(img.tobytes() + repr(img.size).encode()) != a["sha256"]:
                    raise ValueError(f"{task['task_id']}: asset {a['id']} differs from the index")
            staging = dest / "_stage"
            write_design(staging, design)  # -> _stage/designs/<id>/{assets/, preview.jpg}
            for p in (staging / "designs" / crello_id).iterdir():
                p.replace(dest / p.name)
            shutil.rmtree(staging)

        # The design cache lives at <root>/designs/<crello_id>, which is where
        # the task record's relative asset paths point.
        return self.designs.get(crello_id, fill)

    def prefetch(self, tasks: list[dict]) -> None:
        for task in tasks:
            self._prefetch.submit(self.materialize, task)

    def warm(self, splits=("train", "validation", "test"), workers: int = 4, progress=None) -> dict:
        """Materialise every design in `splits`, one row group at a time, so each ~50 MB group is
        downloaded once. Afterwards resets never touch the network: run it before training or a
        large eval so a burst of cold resets does not queue behind bucket downloads."""
        by_group: dict[tuple, list[str]] = {}
        for split in splits:
            for task_id, file, rg in self.db.execute(
                    "SELECT task_id, file, row_group FROM tasks WHERE split=? ORDER BY file, row_group, row", (split,)):
                by_group.setdefault((file, rg), []).append(task_id)
        done = {"groups": 0, "designs": 0, "failed": []}
        lock = threading.Lock()

        def one(ids):
            for task_id in ids:
                try:
                    self.materialize(self.get(task_id))
                except Exception as exc:  # noqa: BLE001 - report, keep warming the rest
                    with lock:
                        done["failed"].append(f"{task_id}: {type(exc).__name__}: {exc}"[:300])
            with lock:
                done["groups"] += 1
                done["designs"] += len(ids)
                if progress:
                    progress(done["groups"], len(by_group), done["designs"])

        with ThreadPoolExecutor(workers, thread_name_prefix="design-canvas-warm") as pool:
            list(pool.map(one, by_group.values()))
        return {"warmed_row_groups": done["groups"], "warmed_designs": done["designs"], "failed": done["failed"],
                "cache": self.stats()}

    def asset_paths(self, task: dict) -> dict[str, Path]:
        self.materialize(task)
        return super().asset_paths(task)

    def preview_path(self, task: dict) -> Path:
        path = super().preview_path(task)
        if not path.exists():
            import pyarrow.parquet as pq

            self.materialize(task)
            file, rg, row_i, crello_id, _ = self._location(task["task_id"])
            with self.source.open(file) as fh:
                col = pq.ParquetFile(fh, pre_buffer=True).read_row_group(rg, columns=["id", "preview"])
            row = col.slice(row_i, 1).to_pylist()[0]
            if str(row["id"]) != crello_id or not row["preview"]:
                raise FileNotFoundError(f"{task['task_id']} has no Crello preview")
            from .crello import _to_pil

            tmp = path.with_suffix(".part")
            _to_pil(row["preview"]).convert("RGB").save(tmp, "JPEG", quality=88)
            tmp.replace(path)
        return path

    def asset_similarity(self, task: dict):
        self.materialize(task)
        return super().asset_similarity(task)

    def stats(self) -> dict:
        return {"snapshot_id": self.snapshot_id, "designs": self.designs.stats(),
                "row_groups": self.groups.stats(),
                "transport": "mounted-filesystem" if self.source.root else "http"}