"""Example 8 — Workset → dataset: keep a selection, the bridge to training.

Everything so far (05–07) read data through a **workset** — a lightweight, directly
iterable selection over a store. A workset is **live**: it stores your selection, not its
result, so every read re-runs it against the store as it is *right now*. Creating one
costs nothing, and it never goes stale. It is also unnamed and ephemeral — it lives as
long as you are using it and is released once you stop.

That is exactly what you want while exploring and producing, and exactly what you do not
want to train on: a source that can answer differently tomorrow is not reproducible. When
you want a **stable, reusable** source — one you can train on now and re-open unchanged
months later, even as the store keeps growing — you **convert** the workset into a
**dataset**: an immutable, shard-materialized snapshot with a name.

That is the whole lifetime model in one line: **naming a selection is converting it.**
This is the first time a dataset appears, and it is exactly the moment you feel the need —
right before training (09).

This walks the arc: create a workset, read it, watch it follow the store, then convert and
watch the dataset refuse to.

    DATAVO_SERVER=https://dev.datavo.io python examples/08_workset_to_dataset.py
"""

from __future__ import annotations

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

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

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, 6)
]
MORE = [
    {"__key__": f"{i:04d}", "audio.wav": f"audio-{i}".encode(), "transcript.txt": f"utterance {i}"}
    for i in range(6, 9)
]
MORE2 = [
    {"__key__": f"{i:04d}", "audio.wav": f"audio-{i}".encode(), "transcript.txt": f"utterance {i}"}
    for i in range(9, 12)
]


def _import_more(client, store: str, samples: list[dict], *, operation_id: str) -> None:
    """Add more samples to an existing store via a second source import."""
    imp = client.create_source_import(
        SourceImportCreateRequest(sample_store=store, operation_id=operation_id)
    )
    client.upload_source_import_part(imp.source_import_id, 0, build_source_tar(samples, root_key="audio.wav"))
    client.commit_source_import(imp.source_import_id)
    client.wait_for_source_import(imp.source_import_id, timeout_seconds=180)


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

        step("create a WORKSET over the store — a live selection, free to create")
        # No name and no lifetime to declare: you are never asked to predict how long you
        # will need this. It stays alive as long as you keep using it.
        workset = client.create_workset(
            sources=[
                SourceSpec(sample_store=store, keys=KEYS)
            ],
            key_groups=[KeyGroupSpec(name="sample", keys=KEYS)],
            shard_plan=ShardPlan(shard_size=2),
        )
        print(f"    workset={workset.workset_id}, samples={workset.sample_count}")

        step("read it while exploring")
        live = SampleStream.from_workset(workset.workset_id, keys=KEYS, client=client)
        print(f"    iterated {sum(1 for _ in live)} samples from the workset")

        step("grow the store — the workset FOLLOWS it, with no action from you")
        _import_more(client, store, MORE, operation_id=f"{store}-more")
        # Iterate again rather than reading the reported count: iteration re-runs the
        # selection, which is the thing being demonstrated. `sample_count` is a convenience
        # number the server caches briefly (it carries `resolved_at` saying when it was
        # computed), so right after an import it may still show the previous value — that is
        # the cache being cheap, not the workset being stale.
        grown = SampleStream.from_workset(workset.workset_id, keys=KEYS, client=client)
        print(
            f"    same workset id, {sum(1 for _ in grown)} samples now — no new workset, "
            "no refresh call: the selection re-ran against the current store"
        )

        step(f"CONVERT the workset → permanent dataset {kept!r} — the way to KEEP a selection")
        # Convert resolves the selection once, at THIS instant, and freezes that result into
        # the dataset. Sources + key groups carry over, so you don't restate the selection.
        # The workset stays live and usable afterwards; the dataset never moves again. This
        # is the stable source you train on (09).
        client.convert_workset_to_dataset(workset.workset_id, kept)
        ready = client.wait_for_dataset(kept, timeout_seconds=600)
        print(f"    dataset state={ready.lifecycle_state!r}, sample_count={ready.sample_count}")

        step("grow the store AFTER conversion — the dataset is frozen (that's reproducibility)")
        _import_more(client, store, MORE2, operation_id=f"{store}-more2")
        store_now = client.get_store(store).sample_count
        frozen = client.get_dataset(kept)
        print(
            f"    live store now holds {store_now} samples; the dataset still holds "
            f"{frozen.sample_count} — an immutable snapshot, unaffected by later imports"
        )

        step("open the dataset like any other — the reusable source for training (09)")
        ds = SampleStream.from_dataset(kept, keys=KEYS, client=client)
        print(f"    iterated {sum(1 for _ in ds)} samples from the frozen dataset")
    finally:
        cleanup_datasets(client, kept)  # a dataset pins its store → tear down first
        cleanup(client, store)


if __name__ == "__main__":
    main()