"""Example 14 — Manual split sets: materialize a partition you already have.

Sometimes the split is decided outside datavo — a published benchmark's own train/test
lists, a hand-curated eval set, an external split you must reproduce exactly. A **manual
split set** takes that partition directly: every split carries an ``id_list`` selector and
you upload each split's ``sample_id``s, so datavo materializes what you hand it instead of
computing an assignment. It bypasses the engine of example 09 entirely (no ratios, no
rules, no preview).

This script is four moves:

    1. SEED + FREEZE   a parent dataset (splits are cut over a dataset, never a workset)
    2. PARTITION       decide the membership yourself — including a stale id and some
                       samples left in no split (manual sets may cover the parent partially)
    3. UPLOAD + CREATE upload each split's ids (multi-part, tagged by split), then create
    4. READ            confirm the children, and read the drop/overlap diagnostics

You own the partition; datavo enforces one thing — a split draws from its parent. An
uploaded id not in the parent is dropped (reported, not an error), and overlap is
reported, not gated.

    DATAVO_SERVER=https://dev.datavo.io python examples/14_manual_split.py
"""

from __future__ import annotations

from _common import (
    build_dataset,
    cleanup,
    cleanup_datasets,
    cleanup_split_sets,
    make_client,
    seed_source_store,
    step,
    unique,
    wait_for_split_set,
)

from datavo_sdk import (
    DatasetSplitSetCreateRequest,
    KeyGroupSpec,
    SampleStream,
    ShardPlan,
    SourceSpec,
    SplitIdListSelector,
    SplitOutputSpec,
)

# A plain parent — manual splits care only about identity, so the samples carry just one key.
SAMPLES = [{"__key__": f"utt-{i:03d}", "audio.wav": f"audio-{i}".encode()} for i in range(20)]


def parent_ids(client, dataset: str) -> list[str]:
    """The canonical sample ids in the parent, in membership order. ``decode=False``: we want
    which samples are members (their ids), not their bytes."""
    return [
        sample["__key__"]
        for sample in SampleStream.from_dataset(dataset, keys=["audio.wav"], client=client, decode=False)
    ]


def child_ids(client, child_dataset: str) -> set[str]:
    return {
        sample["__key__"]
        for sample in SampleStream.from_dataset(child_dataset, keys=["audio.wav"], client=client, decode=False)
    }


def main() -> None:
    client = make_client()
    store = unique("example_manual_store")
    dataset = unique("example_manual_ds")
    split_set = unique("benchmark")

    try:
        # 1. SEED + FREEZE — the same setup as any split: a store, then a frozen dataset.
        step(f"seed {store!r} ({len(SAMPLES)} samples) and freeze it into dataset {dataset!r}")
        seed_source_store(client, store, samples=SAMPLES)
        build_dataset(
            client,
            dataset,
            sources=[SourceSpec(sample_store=store, keys=["audio.wav"])],
            key_groups=[KeyGroupSpec(name="audio", keys=["audio.wav"])],
            shard_plan=ShardPlan(shard_size=8),
        )

        # 2. PARTITION — decide membership yourself. Here: the first 12 for train, the next 6
        #    for test (plus one STALE id that is not in this parent), and the last 2 left in NO
        #    split. A manual split set may cover the parent partially — that is legal.
        ids = parent_ids(client, dataset)
        train_ids = ids[:12]
        test_ids = ids[12:18] + ["utt-STALE-not-in-parent"]
        step(f"partition by hand: train={len(train_ids)}, test={len(test_ids)} (1 stale), "
             f"{len(ids) - 18} left uncovered")

        # 3. UPLOAD + CREATE — upload each split's ids (test spread across TWO files, to show a
        #    large split need not be one file), get the membership_token, then create the set.
        step("upload the membership (multi-part, tagged by split) and create the manual split set")
        token = client.upload_split_membership(
            dataset,
            {
                "train": train_ids,                    # a flat list → one part
                "test": [test_ids[:3], test_ids[3:]],  # a list of lists → one part each
            },
        )
        client.create_dataset_split_set(
            dataset,
            DatasetSplitSetCreateRequest(
                name=split_set,
                membership_token=token,
                splits=[
                    SplitOutputSpec(name="train", selector=SplitIdListSelector()),
                    SplitOutputSpec(name="test", selector=SplitIdListSelector()),
                ],
                shard_plan=ShardPlan(shard_size=8),
            ),
        )

        # 4. READ — wait for materialization, then confirm the children and the diagnostics.
        detail = wait_for_split_set(client, dataset, split_set)
        sizes = {child.split_name: child.sample_count for child in detail.datasets}
        print(f"    → child sizes {sizes} (test is 6, not 7: the stale id was dropped)")
        for report in detail.manual.splits:
            print(f"      {report.name:<6} staged={report.staged} materialized={report.materialized} "
                  f"dropped={report.dropped}")
        print(f"      overlap: {detail.manual.overlap.count}")

        # The children are ordinary datasets — read one back like any other.
        assert child_ids(client, f"{split_set}/train") == set(train_ids)
        print(f"    → {split_set}/train holds exactly the 12 ids you uploaded")
    finally:
        # Teardown: split set (archive → delete, cascading children) → parent dataset → store.
        cleanup_split_sets(client, dataset, split_set)
        cleanup_datasets(client, dataset)
        cleanup(client, store)


if __name__ == "__main__":
    main()