"""Example 2 — Add: branch a selection of samples into a new store.

An **Add** is a *grow op*, like a source import, but instead of ingesting raw
bytes it brings an existing selection — a whole store, a filtered/limited subset,
or a dataset — into a new store as a disjoint, zero-copy union: fresh sample ids,
shared (refcounted) blobs, and a stable ``origin_key`` linking each new row back
to the sample it came from.

The canonical shape is two steps: ``create_store`` (empty) then
``add_to_store``. Adds are worker-deferred, so we poll the operation to ``ready``.

    DATAVO_SERVER=https://dev.datavo.io python examples/02_add_to_store.py
"""

from __future__ import annotations

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

from datavo_sdk import SampleStoreCreateRequest

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


def main() -> None:
    client = make_client()
    source = unique("example_add_src")
    branch = unique("example_add")
    try:
        step(f"seed a 6-sample source store {source!r} (see 01_ingest_source.py for ingest details)")
        seed_source_store(client, source, samples=SAMPLES)

        step(f"create the empty target store {branch!r}")
        client.create_store(name=branch, display_name="Add example")

        step(f"Add a 4-sample slice of {source!r} into {branch!r}, projecting only audio.wav")
        client.add_to_store(
            branch,
            operation_id="add-demo",  # caller-owned; pass a stable id for idempotent replay
            from_store=source,
            key_subset=["audio.wav"],  # bring only this key across
            limit=4,  # deterministic subset
            display_name="branch of source",
        )

        step("wait for the worker to materialize the Add")
        client.wait_for_add(branch, "add-demo", timeout_seconds=180)

        src_detail = client.get_store(source)
        branch_detail = client.get_store(branch)
        print(f"\n  source {source!r}: {src_detail.sample_count} samples, keys={[k.key for k in src_detail.keys]}")
        print(f"  branch {branch!r}: {branch_detail.sample_count} samples, keys={[k.key for k in branch_detail.keys]}")
        print("  → disjoint rows (fresh ids, shared blobs), limited to 4 and projected to audio.wav")
    finally:
        cleanup(client, branch, source)


if __name__ == "__main__":
    main()