"""Example 5 — Read: inspect a store, then select a workset and iterate its samples.

There are two ways to *read* what a store holds, shown side by side:

* **Read the store's shape** — the live, still-growing object. ``get_store`` gives
  its counts and keys; ``list_store_operations`` gives the op graph. This is what the UI and
  the write-side examples (01–04) do. It is not iterable sample bytes — it is the store's
  current shape.
* **Iterate its samples** — pick the keys you want with a **workset** (a live selection
  over the store) and stream decoded samples with ``SampleStream.from_workset``. A workset
  stores your selection rather than its result, so creating one costs nothing and every read
  re-runs it against the store as it is now. It is directly iterable — the bytes don't move
  until you iterate — and it is released once you stop using it.

That's all you need to read. When you want a *stable, reusable* snapshot to train on and
reproduce later, you convert a workset into a **dataset** — that enters in
08_workset_to_dataset.py (the bridge to training).

    DATAVO_SERVER=https://dev.datavo.io python examples/05_read_samples.py
"""

from __future__ import annotations

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

from datavo_sdk import (
    SourceSpec,
    SampleStream,
    KeyGroupSpec,
    ShardPlan,
)

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


def main() -> None:
    client = make_client()
    store = unique("example_read")
    try:
        step(f"seed store {store!r} with 5 samples (see 01_ingest_source.py)")
        seed_source_store(client, store, samples=SAMPLES, default_tags={"language": "en"})

        step("read the store's SHAPE directly — its live counts, keys, op graph")
        detail = client.get_store(store)
        ops = client.list_store_operations(store)["operations"]
        print(f"    store has {detail.sample_count} samples; keys={[k.key for k in detail.keys]}")
        print(f"    operations: {[o['kind'] for o in ops]}")

        step(f"select a workset over {store!r} (the keys you want) — a lightweight, iterable selection")
        workset = client.create_workset(
            sources=[
                SourceSpec(
                    sample_store=store,
                    keys=["audio.wav", "transcript.txt"],
                )
            ],
            key_groups=[
                KeyGroupSpec(name="audio", keys=["audio.wav"]),
                KeyGroupSpec(name="text", keys=["transcript.txt"]),
            ],
            shard_plan=ShardPlan(shard_size=2),
        )
        ready = client.wait_for_workset(workset.workset_id, timeout_seconds=600)
        print(f"    workset state={ready.state!r}, sample_count={ready.sample_count}")

        step("ITERATE the workset directly — open + iterate decoded samples (bytes stream lazily)")
        ds = SampleStream.from_workset(workset.workset_id, keys=["audio.wav", "transcript.txt"], client=client)
        count = 0
        for sample in ds:
            count += 1
            if count == 1:  # show the shape of one decoded sample
                print(
                    f"    sample {sample['__key__']}: "
                    f"audio.wav={type(sample['audio.wav']).__name__}({len(sample['audio.wav'])} bytes), "
                    f"transcript.txt={sample['transcript.txt']!r}"
                )
        print(f"    iterated {count} samples from the workset")
    finally:
        # A ready workset doesn't pin the store (only an in-flight one does); dropping the
        # store cascades its workset cleanup, and an unused workset is released on its own.
        cleanup(client, store)


if __name__ == "__main__":
    main()