Source code for zephon.ops.children

# Copyright 2025 DatologyAI
# SPDX-License-Identifier: Apache-2.0

"""Helpers for child/contributor-aware metadata construction."""

from collections.abc import Iterable
from typing import Any

from zephon.types import ContributorRef, SampleCursor, SampleMeta, SampleRecord




def spawn_child(
    parent: SampleMeta,
    child_idx: int,
    *,
    is_last_child: bool = False,
    tags: dict[str, Any] | None = None,
) -> SampleMeta:
    """Create metadata for a child derived from ``parent``.

    - Assigns deterministic child lineage for identity/replay.
    - Contributors:
        * If the parent has no contributors (single-base record), emit one
          contributor tied to the child cursor and honor ``is_last_child``.
        * If the parent already has contributors (e.g., a child or packed input),
          preserve their cursors. Only the final child inherits their closing
          flags; every earlier child gets non-closing refs.
    - ``is_last_child`` identifies the final child of this immediate parent.
      An inherited closing flag is responsibility to complete a base offset,
      not evidence that it has already been delivered and closed.
    - Raises if you try to spawn a closing child from a tombstone.
    """
    if is_last_child and parent.tombstone:
        raise ValueError("Cannot spawn a closing child from tombstone metadata")

    new_lineage = parent.lineage + (int(child_idx),)
    if parent.contributors:
        if is_last_child:
            # Preserve the parent's contributors, including their closing flags.
            refs = parent.contributors
        else:
            # Preserve every contributor's identity, but clear its closing flag.
            refs = tuple(
                ContributorRef(cursor=ref.cursor, is_last_child=False)
                for ref in parent.contributors
            )
    else:
        cursor = SampleCursor(
            parent.chunk_id, parent.chunk_offset, parent.sample_id, new_lineage
        )
        refs = (ContributorRef(cursor=cursor, is_last_child=is_last_child),)

    child_tags = dict(parent.tags) if tags is None else dict(tags)
    child_tags.pop("_tombstone", None)
    meta = SampleMeta(
        sample_id=parent.sample_id,
        lane_id=parent.lane_id,
        chunk_id=parent.chunk_id,
        chunk_offset=parent.chunk_offset,
        # Preserve component contribution tracking from parent.
        # Children inherit the same component membership as their parent.
        component_sample_counts=dict(parent.component_sample_counts),
        component_token_counts=(
            dict(parent.component_token_counts)
            if parent.component_token_counts is not None
            else None
        ),
        lineage=new_lineage,
        tags=child_tags,
    ).with_contributors(refs)
    return meta





def pack_meta(
    primary_cursor: SampleCursor,
    contributors: Iterable[ContributorRef],
    *,
    lane_id: int,
    component_sample_counts: dict[int, int],
    component_token_counts: dict[int, int] | None = None,
    tags: dict[str, Any] | None = None,
) -> SampleMeta:
    """Build metadata for a packed record that merges multiple contributors.

    ``primary_cursor`` is the replay identity for the packed record and must be
    unique per lane. ``contributors`` lists all contributors included in the pack;
    any contributor that completes a base offset must set ``is_last_child=True``.

    ``component_sample_counts`` aggregates how many original samples from each
    component are included in this pack. For example, if packing 3 samples from
    component 0 and 2 from component 1, this would be ``{0: 3, 1: 2}``.

    ``component_token_counts`` optionally provides token counts per component,
    computed when packing happens after tokenization. If packing before tokenize,
    pass None and ensure_mixture will fall back to distributing by sample counts.
    """
    tags = {} if tags is None else dict(tags)
    tags.pop("_tombstone", None)
    meta = SampleMeta(
        sample_id=primary_cursor.sample_id,
        lane_id=lane_id,
        chunk_id=primary_cursor.chunk_id,
        chunk_offset=primary_cursor.chunk_offset,
        component_sample_counts=component_sample_counts,
        component_token_counts=component_token_counts,
        lineage=primary_cursor.lineage,
        tags=tags,
    ).with_contributors(tuple(contributors))
    return meta





def tombstone_meta(ref: ContributorRef, lane_id: int) -> SampleMeta:
    """Emit a tombstone that marks a base offset complete without payload."""
    if not ref.is_last_child:
        raise ValueError(
            "tombstone_meta requires a closing contributor (is_last_child=True)"
        )
    meta = SampleMeta(
        sample_id=ref.cursor.sample_id,
        lane_id=lane_id,
        chunk_id=ref.cursor.chunk_id,
        chunk_offset=ref.cursor.chunk_offset,
        lineage=ref.cursor.lineage,
        tags={"_tombstone": True},
    ).with_contributors((ref,))
    return meta





def tombstones_for_record(record: SampleRecord) -> list[SampleRecord]:
    """Emit tombstone records for every closing contributor in *record*.

    Each record tracks which chunk offsets it contributes to via
    ``contribution_refs()``.  Only refs with ``is_last_child=True`` need a
    tombstone — intermediate children (from split/spawn) don't close the
    base offset, so the engine doesn't need a signal for them.
    """
    tombstones: list[SampleRecord] = []
    for ref in record.meta.contribution_refs():
        if ref.is_last_child:
            tombstones.append(
                SampleRecord(
                    meta=tombstone_meta(ref, record.meta.lane_id),
                    payload=None,
                )
            )
    return tombstones



__all__ = [
    "pack_meta",
    "spawn_child",
    "tombstone_meta",
    "tombstones_for_record",
]