"""Example 7 — Union several stores into one workset, and reconcile differing schemas.

**Spanning stores is adding a source.** ``sources=[…]`` is the union; each source
carries its own store, keys, row filter and producer pins, and the membership is the
row-wise union of them. Over keys the stores agree on, that is the whole story — no
``key_groups`` needed, since one ``all`` group over the sources' keys is derived.

Where it gets interesting is when the **schemas differ**, and an ``any`` key group lets
one logical field be satisfied by **either** of several keys per sample — e.g. audio
delivered as ``audio.wav`` in one batch and ``audio.flac`` in another.

The catch the consumer must handle: an ``any`` group guarantees **at least one**
member is present per sample, but not **which** one, and datavo does **not** coalesce
them. So a row carries ``audio.wav`` OR ``audio.flac`` (never assume a fixed key) and
you read whichever is there:

    audio = row.get("audio.wav") or row.get("audio.flac")

Here two 3-sample stores — one all-``audio.wav``, one all-``audio.flac`` — union twice:
first over ``transcript.txt``, the key both carry, which needs nothing but the two
sources; then over the audio, with an ``any`` group. Iterating the second, each row
resolves through exactly one of the two audio keys.

    DATAVO_SERVER=https://dev.datavo.io python examples/07_union_sources.py
"""

from __future__ import annotations

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

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

# Same logical content, different audio encodings per batch.
SAMPLES_WAV = [
    {"__key__": f"w{i:03d}", "audio.wav": f"wav-{i}".encode(), "transcript.txt": f"wav batch {i}"}
    for i in range(1, 4)
]
SAMPLES_FLAC = [
    {"__key__": f"f{i:03d}", "audio.flac": f"flac-{i}".encode(), "transcript.txt": f"flac batch {i}"}
    for i in range(1, 4)
]

AUDIO_KEYS = ["audio.wav", "audio.flac"]


def main() -> None:
    client = make_client()
    wav_store = unique("example_or_wav")
    flac_store = unique("example_or_flac")
    try:
        step(f"seed {wav_store!r} (3× audio.wav) and {flac_store!r} (3× audio.flac)")
        seed_source_store(client, wav_store, samples=SAMPLES_WAV, root_key="audio.wav")
        seed_source_store(client, flac_store, samples=SAMPLES_FLAC, root_key="audio.flac")

        step("union both stores over the key they share — two sources, nothing else")
        shared = client.create_workset(
            # This is the whole of "a workset spanning stores": one SourceSpec per store.
            # `key_groups` is omitted, so one `all` group over the sources' keys is derived —
            # here just `transcript.txt`, which both stores carry.
            sources=[
                SourceSpec(sample_store=wav_store, keys=["transcript.txt"]),
                SourceSpec(sample_store=flac_store, keys=["transcript.txt"]),
            ],
            shard_plan=ShardPlan(shard_size=2),
        )
        ready_shared = client.wait_for_workset(shared.workset_id, timeout_seconds=600)
        assert ready_shared.sample_count == 6, ready_shared.sample_count
        print(f"    workset state={ready_shared.state!r}, sample_count={ready_shared.sample_count} (= 3 + 3)")
        # The create response echoes the composition it resolved — both stores, and the
        # key group that was derived — so there is no guessing what a union turned into.
        print("    → composition on the create response:")
        print(f"      sources={[s.sample_store for s in shared.sources]}, "
              f"key_groups={[(g.name, g.mode, g.keys) for g in shared.key_groups]}")

        step("union the audio too — the names differ, so an 'any' group (audio.wav OR audio.flac)")
        workset = client.create_workset(
            sources=[
                # Each source contributes only the audio key its store actually has;
                # the union is the row-wise union of both memberships.
                SourceSpec(
                    sample_store=wav_store,
                    keys=["audio.wav", "transcript.txt"],
                ),
                SourceSpec(
                    sample_store=flac_store,
                    keys=["audio.flac", "transcript.txt"],
                ),
            ],
            key_groups=[
                # mode="any": at least one of these keys per sample (never both, here).
                KeyGroupSpec(name="audio", keys=AUDIO_KEYS, mode="any"),
                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} (= 3 + 3)")

        step("read + iterate — the consumer must coalesce: row has audio.wav OR audio.flac")
        ds = SampleStream.from_workset(workset.workset_id, keys=[*AUDIO_KEYS, "transcript.txt"], client=client)
        by_encoding = {"audio.wav": 0, "audio.flac": 0}
        for sample in ds:
            present = [k for k in AUDIO_KEYS if sample.get(k) is not None]
            assert len(present) == 1, f"'any' group should yield exactly one audio key, got {present}"
            # Deal with it during iteration — do NOT assume a fixed key name:
            audio = sample.get("audio.wav") or sample.get("audio.flac")
            by_encoding[present[0]] += 1
            assert audio is not None
        print(f"    iterated {sum(by_encoding.values())} samples — "
              f"{by_encoding['audio.wav']} via audio.wav, {by_encoding['audio.flac']} via audio.flac")
        print("    → an 'any' group guarantees ≥1 member, not a specific one; read with row.get(...)")
    finally:
        cleanup(client, wav_store, flac_store)


if __name__ == "__main__":
    main()