"""Example 6 — Consume produced features, and compare two producers for one key.

A produced key (see 03_attach_keys.py) lives in the store under a **producer**, and it
is *not* visible from a plain base view — reading the store's ``all`` view surfaces
only ground-truth keys. To read a derived feature you select a **workset** whose
``sample_store`` source lists the key and **pins the producer** that wrote it
(``producer_pins={key: producer_id}``).

Crucially, the pin is on *producer identity*, not the key name — so one store can
hold the **same** semantic key written by **several** producers (e.g. two model
versions), and each consuming workset pins the one it wants. Reading the key
unpinned would be ambiguous; the pin resolves exactly one producer's bytes.

This runs that comparison end to end:

1. produce ``f0.npy`` into one store under two producers, ``f0_v1`` and ``f0_v2``
   (idempotency is per ``(producer_id, key)``, so the second producer sees every
   sample as its own work — the two coexist in the store);
2. select one workset per producer, each pinning ``producer_pins={"f0.npy": <pid>}``;
3. read both and line up ``f0.npy`` for the same samples — same key, different bytes.

See docs/guide/09_consume_derived_keys.md ("Two producers, one key").

    DATAVO_SERVER=https://dev.datavo.io python examples/06_consume_produced_features.py
"""

from __future__ import annotations

import tempfile
from pathlib import Path

import numpy as np

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

from datavo_sdk import (
    SourceSpec,
    HasKeys,
    AttachedBy,
    Not,
    SampleStream,
    KeyGroupSpec,
    ShardPlan,
    TarWriter,
)

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

FEATURE_KEY = "f0.npy"
F0_DIM = 8
# Two producers writing the SAME key f0.npy — stand-ins for two f0 model versions.
# The constant fill is just so the two outputs are visibly different when compared.
PRODUCERS = [("f0_v1", 1.0), ("f0_v2", 2.0)]


def attach_f0(client, store: str, producer_id: str, fill: float) -> None:
    """Attach ``f0.npy`` under ``producer_id``, filled with ``fill`` (condensed 03_attach_keys.py).

    A real producer computes f0 from ``audio.wav``; here each "model version" emits a
    distinct constant so the comparison at the end is legible.
    """
    # The samples this producer has not attached FEATURE_KEY to yet. Idempotency is per
    # (producer_id, key), so the second producer sees every sample as its own work.
    workset = client.create_workset(
        sources=[SourceSpec(
            sample_store=store,
            keys=["audio.wav"],
            filters=[
                HasKeys(["audio.wav"]),
                Not(AttachedBy(producer_id, [FEATURE_KEY])),
            ],
        )],
        shard_plan=ShardPlan(shard_size=2),
    )
    if not workset.sample_count:
        return  # nothing missing — idempotent no-op
    inputs = SampleStream.from_workset(
        workset.workset_id, keys=["audio.wav"], client=client, decode=False
    )
    with tempfile.TemporaryDirectory() as tmp:
        shard = Path(tmp) / f"{producer_id}.tar"
        with TarWriter(str(shard)) as writer:
            for sample in inputs:  # keyed by canonical sample_id
                writer.write({"__key__": sample["__key__"], FEATURE_KEY: np.full(F0_DIM, fill, dtype=np.float32)})
        client.attach_keys(
            store, str(shard),
            keys=[FEATURE_KEY], producer_id=producer_id,
            # Caller-recorded producer fields are how you tell the two producers apart for
            # reproducibility — they ride the producer's generic ``metadata``, not a
            # datavo-specific column.
            metadata={"producer_kind": "f0", "maker": "example", "model_version": producer_id},
        )


def workset_pinned_to(client, store: str, producer_id: str) -> str:
    """A workset over ``store`` that pins ``FEATURE_KEY`` to ``producer_id``; returns its id."""
    workset = client.create_workset(
        sources=[
            SourceSpec(
                sample_store=store,
                keys=["audio.wav", FEATURE_KEY],
                # base audio is unpinned (resolves from the store); the derived key
                # is pinned so it resolves ONLY this producer's bytes:
                producer_pins={FEATURE_KEY: producer_id},
            )
        ],
        key_groups=[
            KeyGroupSpec(name="inputs", keys=["audio.wav"]),
            KeyGroupSpec(name="features", keys=[FEATURE_KEY]),
        ],
        shard_plan=ShardPlan(shard_size=2),
    )
    client.wait_for_workset(workset.workset_id, timeout_seconds=600)
    return workset.workset_id


def main() -> None:
    client = make_client()
    store = unique("example_compare")
    try:
        step(f"seed store {store!r} with 4 audio samples")
        seed_source_store(client, store, samples=SAMPLES)

        step(f"produce {FEATURE_KEY} into the SAME store under two producers: {[p for p, _ in PRODUCERS]}")
        for producer_id, fill in PRODUCERS:
            attach_f0(client, store, producer_id, fill)
            coverage = client.producer_coverage(store, producer_id)
            print(f"    producer {producer_id!r} (fill={fill}) coverage complete={coverage.get('complete')}")
        # The store now holds f0.npy twice over — one revision per producer. A plain
        # read of the store can't say which; the pin disambiguates.

        step("select one workset per producer, each pinning f0.npy to its own producer")
        worksets = {}
        for producer_id, _ in PRODUCERS:
            worksets[producer_id] = workset_pinned_to(client, store, producer_id)
            print(f"    workset for {producer_id!r} → pins {FEATURE_KEY} = {producer_id!r}")

        step("read both worksets and compare f0.npy for the same samples")
        by_producer = {
            producer_id: {
                s["__key__"]: s[FEATURE_KEY]
                for s in SampleStream.from_workset(ws_id, keys=[FEATURE_KEY], client=client)
            }
            for producer_id, ws_id in worksets.items()
        }
        sample_ids = sorted(next(iter(by_producer.values())))
        for sample_id in sample_ids[:3]:
            means = "  ".join(f"{pid}={by_producer[pid][sample_id].mean():.1f}" for pid, _ in PRODUCERS)
            print(f"    sample {sample_id}: {FEATURE_KEY} mean  {means}")
        print("    → same key f0.npy, same samples, different bytes — the pin selects the producer")
    finally:
        # Ready worksets don't block store teardown (dropping the store cascades them);
        # an unused workset is released on its own.
        cleanup(client, store)


if __name__ == "__main__":
    main()