"""Example 4 — Operations: inspect the op graph and remove a node (dependency-gated).

Every store is an append-only DAG of operations. ``list_store_operations`` shows
them; ``remove_store_operation`` removes one by its system ``id`` — but removal is
dependency-gated, mirroring "you can't delete an in-between node":

* a **grow op** (add / source_import) is a removable leaf while nothing has
  enriched its samples;
* the **first attach op freezes every prior grow op** — removing one is then a
  409, and the way to discard the data is to delete the whole store.

This runs both cases side by side.

    DATAVO_SERVER=https://dev.datavo.io python examples/04_operations.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 (
    DatavoApiError,
    HasKeys,
    AttachedBy,
    Not,
    SampleStream,
    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)
]


def show_operations(client, store: str) -> list[dict]:
    ops = client.list_store_operations(store)["operations"]
    for o in ops:
        print(f"    - {o['kind']:<13} state={o['state']:<10} id={o['id']}")
    return ops


def attach_f0(client, store: str) -> None:
    """Attach an f0.npy key onto the store's samples (a condensed 03_attach_keys.py)."""
    workset = client.create_workset(
        sources=[SourceSpec(
            sample_store=store,
            keys=["audio.wav"],
            filters=[
                HasKeys(["audio.wav"]),
                Not(AttachedBy("f0_v1", ["f0.npy"])),
            ],
        )],
        shard_plan=ShardPlan(shard_size=2),
    )
    inputs = SampleStream.from_workset(
        workset.workset_id, keys=["audio.wav"], client=client, decode=False
    )
    with tempfile.TemporaryDirectory() as tmp:
        shard = Path(tmp) / "f0.tar"
        with TarWriter(str(shard)) as writer:
            for sample in inputs:
                writer.write({"__key__": sample["__key__"], "f0.npy": np.zeros(4, dtype=np.float32)})
        client.attach_keys(
            store, str(shard),
            keys=["f0.npy"], producer_id="f0_v1",
        )


def main() -> None:
    client = make_client()
    base = unique("example_ops_src")
    removable = unique("example_ops_removable")
    frozen = unique("example_ops_frozen")
    try:
        step(f"seed a base store {base!r} to Add from")
        seed_source_store(client, base, samples=SAMPLES)

        step(f"case 1 — Add into {removable!r}; a clean grow op is a removable leaf")
        client.create_store(name=removable)
        client.add_to_store(removable, operation_id="add-1", from_store=base, key_subset=["audio.wav"])
        client.wait_for_add(removable, "add-1", timeout_seconds=180)
        add_id = next(o["id"] for o in show_operations(client, removable) if o["kind"] == "add")
        client.remove_store_operation(removable, add_id)
        print(f"    ✓ removed add {add_id}; operations now: {client.list_store_operations(removable)['operations']}")

        step(f"case 2 — Add into {frozen!r}, then produce onto it; the Add is now frozen")
        client.create_store(name=frozen)
        client.add_to_store(frozen, operation_id="add-2", from_store=base, key_subset=["audio.wav"])
        client.wait_for_add(frozen, "add-2", timeout_seconds=180)
        attach_f0(client, frozen)
        add_id = next(o["id"] for o in show_operations(client, frozen) if o["kind"] == "add")
        try:
            client.remove_store_operation(frozen, add_id)
            print("    UNEXPECTED: removal succeeded")
        except DatavoApiError as exc:
            print(f"    ✓ removal blocked with {exc.status_code} — an attach op depends on the Add's samples;")
            print("      to discard enriched data, delete the whole store instead.")
    finally:
        cleanup(client, frozen, removable, base)


if __name__ == "__main__":
    main()