"""Example 3 — Attach: enrich a store with a derived key group under a producer.

An **attach** op writes new keys onto a store's *existing* samples, under a named
producer — attach is the act, the producer is the identity that performed it. The
samples to attach into are a **workset** whose row filter carries a
``Not(AttachedBy(...))`` predicate; the producer identity lives on the attach, not the
workset:

1. ``create_workset`` with ``Not(AttachedBy(producer, output_keys))`` in its filter list —
   selects exactly the samples that still need this output.
2. ``SampleStream.from_workset(..., decode=False)`` — stream the raw payloads for those
   samples, in-process. No ``missing_sample_ids`` list to enumerate.
3. ``attach_keys`` — reserve → upload → commit → poll-to-terminal. The first attach
   **creates the producer** (recording the keys it wrote + caller ``metadata``); no
   separate ``register_producer`` call is needed.

Both idempotency mechanisms are visible here, and they do different jobs:

* **convergence** — step 1 subtracts what this producer already wrote, so doneness is
  simply an **empty workset** (``sample_count == 0``) and a half-finished run resumes
  with exactly the remainder.
* **replay** — step 3 passes a stable ``operation_id``, so a retried attach returns the
  ingest already opened for it rather than uploading a second time.

    DATAVO_SERVER=https://dev.datavo.io python examples/03_attach_keys.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 (
    SampleStream,
    HasKeys,
    AttachedBy,
    Not,
    ShardPlan,
    SourceSpec,
    TarWriter,
)

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

PRODUCER_ID = "f0_v1"
OUTPUT_KEY = "f0.npy"
INPUT_KEY = "audio.wav"


def _attach_workset(client, store: str):
    """An attach workset = the not-yet-attached samples that have the required input."""
    return client.create_workset(
        # One SourceSpec per store — the only way a workset says where samples come from.
        # Store, keys and row filter all describe the same source, so they live together.
        sources=[SourceSpec(
            sample_store=store,
            keys=[INPUT_KEY],
            # A row filter is a list of predicates, each contributing a set of samples.
            # `AttachedBy` names the samples this producer has already covered; `Not`
            # subtracts that set instead of intersecting it — and wraps any predicate.
            filters=[
                HasKeys([INPUT_KEY]),
                Not(AttachedBy(PRODUCER_ID, [OUTPUT_KEY])),
            ],
        )],
        shard_plan=ShardPlan(shard_size=2),
    )


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

        step(f"create an attach workset: samples missing {OUTPUT_KEY} from producer {PRODUCER_ID!r}")
        workset = _attach_workset(client, store)
        print(f"  workset={workset.workset_id}, state={workset.state!r}, samples={workset.sample_count}")
        if not workset.sample_count:
            # CONVERGENCE: the workset subtracts what this producer already wrote,
            # so it drains to empty and a re-run computes nothing.
            print("  nothing to compute — the attach workset has drained (convergence).")
            return

        step(f"stream inputs with from_workset(..., decode=False) and compute {OUTPUT_KEY}")
        # decode=False: a producer forwards payloads rather than interpreting them.
        inputs = SampleStream.from_workset(
            workset.workset_id, keys=[INPUT_KEY], client=client, decode=False
        )
        with tempfile.TemporaryDirectory() as tmp:
            shard_path = Path(tmp) / "f0_shard.tar"
            with TarWriter(str(shard_path)) as writer:
                for sample in inputs:
                    _audio = sample[INPUT_KEY]  # a real producer computes f0 from these bytes
                    f0 = np.zeros(8, dtype=np.float32)
                    writer.write({"__key__": sample["__key__"], OUTPUT_KEY: f0})
            # The attach registers producer f0_v1 (declared output keys + metadata) if new.
            result = client.attach_keys(
                store,
                str(shard_path),
                keys=[OUTPUT_KEY],
                producer_id=PRODUCER_ID,
                # REPLAY, the other mechanism: a stable id means a retry of this attach
                # returns the ingest already opened for it rather than uploading twice.
                # Omit it and the SDK mints a fresh one, so nothing replays by accident.
                operation_id=f"{store}-{PRODUCER_ID}",
                metadata={"producer_kind": "f0", "maker": "example"},
            )
        print(f"  attach state={result.state!r}, sample_count={result.sample_count}")

        step("re-plan the attach workset → now empty (doneness = an empty workset)")
        again = _attach_workset(client, store)
        print(f"  samples={again.sample_count} (0 = every input-complete sample has the key)")
    finally:
        cleanup(client, store)


if __name__ == "__main__":
    main()