"""Example 12 — Shard caching: the second read is free (local disk, no server).

Every shard datavo streams is **content-addressed** — named by a hash of its bytes. So the
first time you read a dataset its shards are fetched from the server and dropped into a
local **cache tier**; a second pass — same process, a new process, or another machine
sharing a tier — is served from local disk and never touches the server. That is what makes
multi-epoch training and re-opening a frozen dataset cheap. You do not turn it on: every
read in 05–11 already went through it.

This builds a small dataset, reads it COLD (the fetch populates the cache), then reads it
again through a FRESH cache object over the same directory — as a new process would — and
shows the second pass downloads nothing. It uses an explicit on-disk cache dir so the demo
is self-contained; a local user just relies on the default cache (shown at the end).

    DATAVO_SERVER=https://dev.datavo.io python examples/12_warm_and_read_cache.py
"""

from __future__ import annotations

import os
import shutil
import tempfile
from pathlib import Path

from _common import build_dataset, cleanup, cleanup_datasets, make_client, seed_source_store, step, unique

from datavo_sdk import (
    SourceSpec,
    SampleStream,
    KeyGroupSpec,
    ShardPlan,
    ShardCacheHierarchy,
    ShardCacheTier,
    get_default_shard_cache,
)

KEYS = ["audio.wav", "transcript.txt"]
SAMPLES = [
    {"__key__": f"{i:04d}", "audio.wav": f"audio-{i}".encode(), "transcript.txt": f"utterance {i}"}
    for i in range(1, 13)
]


def _folder_cache(root: Path) -> ShardCacheHierarchy:
    """A one-tier local cache under ``root`` — the zero-config shape, made explicit so the
    demo controls where bytes land (a real local user just relies on the default cache)."""
    return ShardCacheHierarchy([ShardCacheTier(root, name="local", max_bytes=2 * 1024**3)])


def _read_all(dataset: str, *, client, cache: ShardCacheHierarchy) -> SampleStream:
    """Open ``dataset`` through ``cache`` and iterate it fully; return the stream so the
    caller can read its ``download_stats`` (downloads vs cache hits)."""
    ds = SampleStream.from_dataset(dataset, keys=KEYS, client=client, cache=cache)
    count = sum(1 for _ in ds)
    print(f"    iterated {count} samples — {ds.download_stats.summary_line()}")
    return ds


def main() -> None:
    client = make_client()
    store = unique("example_cache")
    dataset = unique("example_cache_ds")
    cache_dir = Path(tempfile.mkdtemp(prefix="datavo-cache-"))
    try:
        step(f"seed store {store!r} and freeze a dataset {dataset!r} (see 08_workset_to_dataset.py)")
        seed_source_store(client, store, samples=SAMPLES)
        build_dataset(
            client,
            dataset,
            sources=[SourceSpec(sample_store=store, keys=KEYS)],
            key_groups=[KeyGroupSpec(name="sample", keys=KEYS)],
            shard_plan=ShardPlan(shard_size=4),  # 12 samples → a few shards to fetch
        )

        step("COLD read — the cache is empty, so shards are fetched from the server and cached")
        cold = _read_all(dataset, client=client, cache=_folder_cache(cache_dir))
        assert cold.download_stats.shards_downloaded > 0, "cold read should fetch shards"
        cached_tars = list(cache_dir.rglob("*.tar"))
        print(f"    cache now holds {len(cached_tars)} shard tars under {cache_dir}")

        step("WARM read — a FRESH cache over the same dir (as a new process has) serves from disk")
        # Nothing about the client changes. Shards are content-addressed, so a brand-new cache
        # object pointed at the same directory finds them by content hash — no server round-trip.
        warm = _read_all(dataset, client=client, cache=_folder_cache(cache_dir))
        assert warm.download_stats.cache_hits > 0, "warm read should hit the cache"
        assert warm.download_stats.shards_downloaded == 0, "warm read should download nothing"
        print("    second pass fetched zero shards — the payoff of the cache")

        step("inspect the hierarchy (fastest-first tiers; here just one folder tier)")
        for tier in _folder_cache(cache_dir).tiers:
            budget = "unbounded" if tier.max_bytes is None else f"{tier.max_bytes / 1024**3:.0f} GiB"
            print(f"    tier {tier.name!r}: direct_read={tier.supports_direct_read}, budget={budget}")

        step("zero-config default — get_default_shard_cache() needs no wiring; DATAVO_SHARD_CACHE_DIR points it")
        os.environ["DATAVO_SHARD_CACHE_DIR"] = str(cache_dir / "default")  # keep the demo self-contained
        default = get_default_shard_cache()
        tier = default.tiers[0]
        budget = "unbounded" if tier.max_bytes is None else f"{tier.max_bytes / 1024**3:.0f} GiB"
        print(
            f"    default cache → 1 folder tier at {tier.path} ({budget} LRU). A local user gets this "
            "for free; full topology (cloud tiers, shared machines): docs/reference/shard_cache.md"
        )
    finally:
        cleanup_datasets(client, dataset)  # a dataset pins its store → tear down first
        cleanup(client, store)
        shutil.rmtree(cache_dir, ignore_errors=True)
        print(f"  ✓ cleaned up local cache dir {cache_dir}")


if __name__ == "__main__":
    main()