"""
DatasetPipeline is a user-space preprocessing layer.

Design intent:
- Keep row normalization/selection/sampling/splitting in-process for fast iteration.
- Keep execution-contract and manifest generation close to data-shaping code.
- Hand off to FlowNet Dataset facade (`to_dataset().map_batches(...)`) when runtime
  execution/distribution is required.
"""

from __future__ import annotations

from collections import deque
from copy import deepcopy
from dataclasses import dataclass
from pathlib import Path
import random
from typing import Any, Callable, Iterable, Iterator, Mapping

from .adapters import SageDataMappingContext, iter_sagedata_rows
from .connectors import (
    DEFAULT_MAX_JSONL_LINE_BYTES,
    DEFAULT_MAX_PANDAS_FALLBACK_BYTES,
    ConnectorCheckpoint,
    ConnectorQualityStats,
    iter_image_rows,
    iter_jsonl,
    iter_parquet,
    iter_tfrecord_rows,
)
from .contract_injection import ExperimentRowMapper
from .contracts import RunManifest, stable_json_fingerprint
from .manifest import build_run_manifest, compute_operator_chain_fingerprint

RowLike = Mapping[str, Any]
RowMapper = Callable[[RowLike], RowLike]
RowPredicate = Callable[[RowLike], bool]
RowFlatMapper = Callable[[RowLike], Iterable[RowLike]]
RowWindowMapper = Callable[[list[dict[str, Any]]], RowLike | Iterable[RowLike]]
ConnectorCheckpointHook = Callable[[ConnectorCheckpoint], None]
GroupReducer = Callable[[list[dict[str, Any]]], Any]
StreamBuilder = Callable[[Any], Any]
StreamBridgeCheckpointHook = Callable[["StreamBridgeCheckpoint"], None]
Operator = Callable[[Iterator[dict[str, Any]]], Iterator[dict[str, Any]]]


@dataclass(frozen=True)
class _SourceConfig:
    kind: str
    detail: dict[str, Any]


@dataclass(frozen=True)
class StreamBridgeCheckpoint:
    flow_name: str
    emit_mode: str
    cursor: int
    emitted_payloads: int
    emitted_requests: int
    resume_offset: int
    request_chunk_size: int | None

    def as_dict(self) -> dict[str, Any]:
        return {
            "flow_name": self.flow_name,
            "emit_mode": self.emit_mode,
            "cursor": int(self.cursor),
            "emitted_payloads": int(self.emitted_payloads),
            "emitted_requests": int(self.emitted_requests),
            "resume_offset": int(self.resume_offset),
            "request_chunk_size": (
                int(self.request_chunk_size) if self.request_chunk_size is not None else None
            ),
        }


def _coerce_row(row: RowLike) -> dict[str, Any]:
    if not isinstance(row, Mapping):
        raise TypeError(f"rows must be mapping-like dicts, got {type(row)}")
    return dict(row)


def _callable_name(fn: Any) -> str:
    module = getattr(fn, "__module__", None)
    qualname = getattr(fn, "__qualname__", None)
    if isinstance(module, str) and isinstance(qualname, str):
        return f"{module}.{qualname}"
    name = getattr(fn, "__name__", None)
    if isinstance(name, str):
        return name
    return repr(fn)


class DatasetPipeline:
    """
    User-space dataset preprocessing pipeline for FlowNet experiment workloads.

    This is an in-process row pipeline that composes stable preprocessing steps and
    materializes rows for `fn.read_rows(...)` execution.

    Note:
    - This layer itself does not schedule distributed execution.
    - Distributed/runtime execution starts when users hand off to
      `Dataset.map_batches(...)` from the API facade.
    """

    def __init__(
        self,
        *,
        source_factory: Callable[[], Iterator[dict[str, Any]]],
        source_config: _SourceConfig,
        operators: tuple[Operator, ...] = (),
        operator_specs: tuple[dict[str, Any], ...] = (),
        default_manifest_fields: dict[str, Any] | None = None,
    ):
        self._source_factory = source_factory
        self._source_config = source_config
        self._operators = operators
        self._operator_specs = operator_specs
        self._default_manifest_fields = dict(default_manifest_fields or {})

    @classmethod
    def from_rows(
        cls,
        rows: Iterable[RowLike] | Callable[[], Iterable[RowLike]],
        *,
        source_name: str = "python_rows",
    ) -> "DatasetPipeline":
        source_factory = _build_rows_factory(rows)
        return cls(
            source_factory=source_factory,
            source_config=_SourceConfig(kind="rows", detail={"source_name": str(source_name)}),
        )

    @classmethod
    def from_jsonl(
        cls,
        path: str | Path,
        *,
        quality: ConnectorQualityStats | None = None,
        drop_invalid_rows: bool = False,
        max_line_bytes: int = DEFAULT_MAX_JSONL_LINE_BYTES,
        resume_offset: int = 0,
        checkpoint_every: int | None = None,
        on_checkpoint: ConnectorCheckpointHook | None = None,
    ) -> "DatasetPipeline":
        resolved_path = str(path)

        def _factory() -> Iterator[dict[str, Any]]:
            yield from iter_jsonl(
                resolved_path,
                quality=quality,
                drop_invalid_rows=drop_invalid_rows,
                max_line_bytes=max_line_bytes,
                resume_offset=resume_offset,
                checkpoint_every=checkpoint_every,
                on_checkpoint=on_checkpoint,
            )

        return cls(
            source_factory=_factory,
            source_config=_SourceConfig(
                kind="jsonl",
                detail={
                    "path": resolved_path,
                    "drop_invalid_rows": bool(drop_invalid_rows),
                    "max_line_bytes": int(max_line_bytes),
                    "resume_offset": int(resume_offset),
                    "checkpoint_every": (
                        int(checkpoint_every) if checkpoint_every is not None else None
                    ),
                },
            ),
        )

    @classmethod
    def from_parquet(
        cls,
        path: str | Path,
        *,
        batch_size: int = 2048,
        quality: ConnectorQualityStats | None = None,
        max_pandas_fallback_bytes: int | None = DEFAULT_MAX_PANDAS_FALLBACK_BYTES,
        resume_offset: int = 0,
        checkpoint_every: int | None = None,
        on_checkpoint: ConnectorCheckpointHook | None = None,
    ) -> "DatasetPipeline":
        resolved_path = str(path)

        def _factory() -> Iterator[dict[str, Any]]:
            yield from iter_parquet(
                resolved_path,
                batch_size=batch_size,
                quality=quality,
                max_pandas_fallback_bytes=max_pandas_fallback_bytes,
                resume_offset=resume_offset,
                checkpoint_every=checkpoint_every,
                on_checkpoint=on_checkpoint,
            )

        return cls(
            source_factory=_factory,
            source_config=_SourceConfig(
                kind="parquet",
                detail={
                    "path": resolved_path,
                    "batch_size": int(batch_size),
                    "max_pandas_fallback_bytes": (
                        int(max_pandas_fallback_bytes)
                        if max_pandas_fallback_bytes is not None
                        else None
                    ),
                    "resume_offset": int(resume_offset),
                    "checkpoint_every": (
                        int(checkpoint_every) if checkpoint_every is not None else None
                    ),
                },
            ),
        )

    @classmethod
    def from_images(
        cls,
        path: str | Path,
        *,
        mode: str = "RGB",
        quality: ConnectorQualityStats | None = None,
        drop_bad_rows: bool = True,
        resume_offset: int = 0,
        checkpoint_every: int | None = None,
        on_checkpoint: ConnectorCheckpointHook | None = None,
    ) -> "DatasetPipeline":
        resolved_path = str(path)

        def _factory() -> Iterator[dict[str, Any]]:
            yield from iter_image_rows(
                resolved_path,
                mode=mode,
                quality=quality,
                drop_bad_rows=drop_bad_rows,
                resume_offset=resume_offset,
                checkpoint_every=checkpoint_every,
                on_checkpoint=on_checkpoint,
            )

        return cls(
            source_factory=_factory,
            source_config=_SourceConfig(
                kind="images",
                detail={
                    "path": resolved_path,
                    "mode": mode,
                    "drop_bad_rows": bool(drop_bad_rows),
                    "resume_offset": int(resume_offset),
                    "checkpoint_every": (
                        int(checkpoint_every) if checkpoint_every is not None else None
                    ),
                },
            ),
        )

    @classmethod
    def from_tfrecords(
        cls,
        path: str | Path,
        *,
        quality: ConnectorQualityStats | None = None,
        resume_offset: int = 0,
        checkpoint_every: int | None = None,
        on_checkpoint: ConnectorCheckpointHook | None = None,
    ) -> "DatasetPipeline":
        resolved_path = str(path)

        def _factory() -> Iterator[dict[str, Any]]:
            yield from iter_tfrecord_rows(
                resolved_path,
                quality=quality,
                resume_offset=resume_offset,
                checkpoint_every=checkpoint_every,
                on_checkpoint=on_checkpoint,
            )

        return cls(
            source_factory=_factory,
            source_config=_SourceConfig(
                kind="tfrecords",
                detail={
                    "path": resolved_path,
                    "resume_offset": int(resume_offset),
                    "checkpoint_every": (
                        int(checkpoint_every) if checkpoint_every is not None else None
                    ),
                },
            ),
        )

    @classmethod
    def from_sagedata(
        cls,
        rows: Iterable[RowLike] | Callable[[], Iterable[RowLike]],
        *,
        context: SageDataMappingContext,
        mapper: ExperimentRowMapper | None = None,
    ) -> "DatasetPipeline":
        raw_rows_factory = _build_rows_factory(rows)

        def _factory() -> Iterator[dict[str, Any]]:
            for row in iter_sagedata_rows(raw_rows_factory(), context=context, mapper=mapper):
                yield row.to_dict()

        source_config = _SourceConfig(
            kind="sagedata",
            detail={
                "dataset_ref": context.dataset_ref,
                "usage": context.usage,
                "profile": context.profile,
                "source_version": context.source_version,
                "snapshot_hash": context.snapshot_hash,
                "source_format": context.source_format,
            },
        )
        defaults = {
            "dataset_ref": context.dataset_ref,
            "usage": context.usage,
            "profile": context.profile,
            "source_version": context.source_version,
            "snapshot_hash": context.snapshot_hash,
        }
        return cls(
            source_factory=_factory,
            source_config=source_config,
            default_manifest_fields=defaults,
        )

    def map(self, fn: RowMapper) -> "DatasetPipeline":
        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            for row in rows:
                mapped = fn(dict(row))
                if not isinstance(mapped, Mapping):
                    raise TypeError(f"map function must return mapping-like row, got {type(mapped)}")
                yield dict(mapped)

        return self._with_operator(_op, {"op": "map", "fn": _callable_name(fn)})

    def filter(self, predicate: RowPredicate) -> "DatasetPipeline":
        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            for row in rows:
                if bool(predicate(dict(row))):
                    yield row

        return self._with_operator(_op, {"op": "filter", "predicate": _callable_name(predicate)})

    def flat_map(self, fn: RowFlatMapper) -> "DatasetPipeline":
        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            for row in rows:
                mapped_rows = fn(dict(row))
                for mapped in mapped_rows:
                    if not isinstance(mapped, Mapping):
                        raise TypeError(
                            f"flat_map function must emit mapping-like rows, got {type(mapped)}"
                        )
                    yield dict(mapped)

        return self._with_operator(_op, {"op": "flat_map", "fn": _callable_name(fn)})

    def select(self, *fields: str, allow_missing: bool = False) -> "DatasetPipeline":
        names = [str(name) for name in fields if str(name)]
        if not names:
            raise ValueError("select requires at least one field.")

        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            for row in rows:
                selected: dict[str, Any] = {}
                for name in names:
                    if name not in row:
                        if allow_missing:
                            continue
                        raise KeyError(f"select missing required field: {name!r}")
                    selected[name] = row[name]
                yield selected

        return self._with_operator(
            _op,
            {"op": "select", "fields": list(names), "allow_missing": bool(allow_missing)},
        )

    def dedupe(
        self,
        *,
        key: str | Callable[[RowLike], Any],
        max_keys: int | None = None,
    ) -> "DatasetPipeline":
        if isinstance(key, str):
            key_name = key

            def _key_fn(row: RowLike) -> Any:
                if key_name not in row:
                    raise KeyError(f"dedupe key field missing: {key_name!r}")
                return row[key_name]

            key_desc = f"field:{key_name}"
        else:
            _key_fn = key
            key_desc = _callable_name(key)

        budget: int | None = None
        if max_keys is not None:
            budget = max(1, int(max_keys))

        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            seen: set[str] = set()
            order: deque[str] = deque()
            for row in rows:
                key_value = _key_fn(dict(row))
                dedupe_key = stable_json_fingerprint({"key": key_value})
                if dedupe_key in seen:
                    continue
                seen.add(dedupe_key)
                order.append(dedupe_key)
                if budget is not None:
                    while len(order) > budget:
                        oldest = order.popleft()
                        seen.discard(oldest)
                yield row

        return self._with_operator(_op, {"op": "dedupe", "key": key_desc, "max_keys": budget})

    def sample(
        self,
        *,
        n: int | None = None,
        frac: float | None = None,
        seed: int = 0,
    ) -> "DatasetPipeline":
        if (n is None) == (frac is None):
            raise ValueError("sample requires exactly one of n or frac.")

        if n is not None:
            sample_n = int(n)
            if sample_n < 0:
                raise ValueError("sample n must be >= 0.")

            def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
                if sample_n == 0:
                    return
                rng = random.Random(seed)
                reservoir: list[dict[str, Any]] = []
                for index, row in enumerate(rows):
                    if index < sample_n:
                        reservoir.append(row)
                        continue
                    choice = rng.randint(0, index)
                    if choice < sample_n:
                        reservoir[choice] = row
                for row in reservoir:
                    yield row

            return self._with_operator(
                _op,
                {"op": "sample", "mode": "n", "n": sample_n, "seed": int(seed)},
            )

        assert frac is not None
        sample_frac = float(frac)
        if sample_frac < 0.0 or sample_frac > 1.0:
            raise ValueError("sample frac must be within [0.0, 1.0].")

        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            rng = random.Random(seed)
            for row in rows:
                if rng.random() < sample_frac:
                    yield row

        return self._with_operator(
            _op,
            {"op": "sample", "mode": "frac", "frac": sample_frac, "seed": int(seed)},
        )

    def split(
        self,
        ratios: Mapping[str, float],
        *,
        key: str | Callable[[RowLike], Any] = "sample_id",
        seed: int = 0,
    ) -> dict[str, "DatasetPipeline"]:
        if not ratios:
            raise ValueError("split ratios must not be empty.")

        names: list[str] = []
        values: list[float] = []
        for name, value in ratios.items():
            split_name = str(name).strip()
            if not split_name:
                raise ValueError("split name must be non-empty.")
            ratio = float(value)
            if ratio < 0:
                raise ValueError("split ratio must be >= 0.")
            names.append(split_name)
            values.append(ratio)

        total = sum(values)
        if total <= 0:
            raise ValueError("split ratios sum must be > 0.")

        boundaries: list[tuple[str, float]] = []
        running = 0.0
        for name, value in zip(names, values):
            running += value / total
            boundaries.append((name, running))
        boundaries[-1] = (boundaries[-1][0], 1.0)

        if isinstance(key, str):
            key_name = key

            def _key_fn(row: RowLike) -> Any:
                if key_name not in row:
                    raise KeyError(f"split key field missing: {key_name!r}")
                return row[key_name]

        else:
            _key_fn = key

        def _bucket(row: RowLike) -> str:
            key_value = _key_fn(dict(row))
            digest = stable_json_fingerprint({"seed": int(seed), "key": key_value})
            bucket_value = int(digest, 16) / float(16**len(digest))
            for name, upper in boundaries:
                if bucket_value < upper:
                    return name
            return boundaries[-1][0]

        result: dict[str, DatasetPipeline] = {}
        for split_name in names:
            name_copy = split_name
            result[name_copy] = self.filter(lambda row, _name=name_copy: _bucket(row) == _name)
        return result

    def join(
        self,
        right_rows: "DatasetPipeline | Iterable[RowLike] | Callable[[], Iterable[RowLike]]",
        *,
        left_on: str | Callable[[RowLike], Any],
        right_on: str | Callable[[RowLike], Any] | None = None,
        how: str = "inner",
        right_prefix: str = "right_",
    ) -> "DatasetPipeline":
        join_mode = str(how).strip().lower()
        if join_mode not in {"inner", "left"}:
            raise ValueError("join how must be 'inner' or 'left'.")

        left_key_fn = _resolve_join_key_fn(left_on)
        if right_on is None:
            right_key_fn = _resolve_join_key_fn(left_on)
        else:
            right_key_fn = _resolve_join_key_fn(right_on)

        right_factory = _build_join_rows_factory(right_rows)
        right_prefix = str(right_prefix)

        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            index: dict[str, list[dict[str, Any]]] = {}
            for right in right_factory():
                right_row = _coerce_row(right)
                key_value = right_key_fn(right_row)
                key_fp = stable_json_fingerprint({"key": key_value})
                index.setdefault(key_fp, []).append(right_row)

            for left_row in rows:
                key_value = left_key_fn(left_row)
                key_fp = stable_json_fingerprint({"key": key_value})
                matches = index.get(key_fp, ())
                if matches:
                    for right_row in matches:
                        yield _merge_join_rows(left_row, right_row, right_prefix=right_prefix)
                    continue
                if join_mode == "left":
                    yield dict(left_row)

        return self._with_operator(
            _op,
            {
                "op": "join",
                "how": join_mode,
                "left_on": _join_key_desc(left_on),
                "right_on": _join_key_desc(right_on if right_on is not None else left_on),
                "right_prefix": right_prefix,
            },
        )

    def window_map(
        self,
        *,
        size: int,
        fn: RowWindowMapper,
        step: int | None = None,
        emit_partial: bool = False,
    ) -> "DatasetPipeline":
        window_size = int(size)
        if window_size <= 0:
            raise ValueError("window_map size must be > 0.")
        window_step = window_size if step is None else int(step)
        if window_step <= 0:
            raise ValueError("window_map step must be > 0.")

        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            buffer: list[dict[str, Any]] = []
            for row in rows:
                buffer.append(row)
                while len(buffer) >= window_size:
                    window = [dict(item) for item in buffer[:window_size]]
                    yield from _coerce_window_output(fn(window))
                    del buffer[:window_step]

            if emit_partial and buffer:
                window = [dict(item) for item in buffer]
                yield from _coerce_window_output(fn(window))

        return self._with_operator(
            _op,
            {
                "op": "window_map",
                "size": window_size,
                "step": window_step,
                "emit_partial": bool(emit_partial),
                "fn": _callable_name(fn),
            },
        )

    def group_by(
        self,
        *,
        key: str | Callable[[RowLike], Any],
        reducers: Mapping[str, GroupReducer],
        include_group_key: bool = True,
    ) -> "DatasetPipeline":
        reducer_items = [(str(name), reducer) for name, reducer in reducers.items()]
        if not reducer_items:
            raise ValueError("group_by requires at least one reducer.")
        for name, reducer in reducer_items:
            if not name:
                raise ValueError("group_by reducer name must be non-empty.")
            if not callable(reducer):
                raise TypeError(f"group_by reducer {name!r} must be callable.")

        key_fn = _resolve_join_key_fn(key)

        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            groups: dict[str, tuple[Any, list[dict[str, Any]]]] = {}
            for row in rows:
                key_value = key_fn(row)
                key_fp = stable_json_fingerprint({"key": key_value})
                bucket = groups.get(key_fp)
                if bucket is None:
                    groups[key_fp] = (key_value, [dict(row)])
                    continue
                bucket[1].append(dict(row))

            for key_value, group_rows in groups.values():
                out: dict[str, Any] = {}
                if include_group_key:
                    if isinstance(key, str):
                        out[key] = key_value
                    else:
                        out["group_key"] = key_value
                materialized_group = [dict(item) for item in group_rows]
                for field_name, reducer in reducer_items:
                    out[field_name] = reducer(materialized_group)
                yield out

        return self._with_operator(
            _op,
            {
                "op": "group_by",
                "key": _join_key_desc(key),
                "reducers": {name: _callable_name(reducer) for name, reducer in reducer_items},
                "include_group_key": bool(include_group_key),
            },
        )

    def schema_validate(
        self,
        *,
        required_fields: Iterable[str] = (),
        field_types: Mapping[str, type | tuple[type, ...]] | None = None,
    ) -> "DatasetPipeline":
        required = [str(name) for name in required_fields if str(name)]
        type_rules = dict(field_types or {})

        def _op(rows: Iterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
            for row in rows:
                for name in required:
                    if name not in row:
                        raise KeyError(f"schema_validate missing required field: {name!r}")
                for name, expected_type in type_rules.items():
                    if name not in row:
                        continue
                    if not isinstance(row[name], expected_type):
                        raise TypeError(
                            f"schema_validate field {name!r} expected {expected_type}, "
                            f"got {type(row[name])}"
                        )
                yield row

        return self._with_operator(
            _op,
            {
                "op": "schema_validate",
                "required_fields": list(required),
                "field_types": {name: _type_repr(tp) for name, tp in type_rules.items()},
            },
        )

    def iter_rows(self) -> Iterator[dict[str, Any]]:
        # Operator chain runs in the local process by design.
        stream: Iterator[dict[str, Any]] = self._source_factory()
        for op in self._operators:
            stream = op(stream)
        for row in stream:
            yield _coerce_row(row)

    def to_rows(self) -> list[dict[str, Any]]:
        return list(self.iter_rows())

    def to_dataset(self):
        from flownet.api.dataset import read_rows

        # Boundary between local preprocessing and FlowNet Dataset execution.
        return read_rows(lambda: self.iter_rows())

    def to_stream_flow(
        self,
        builder: StreamBuilder,
        *,
        flow_name: str = "dataset_pipeline_stream",
    ):
        """
        Compile a DataStream flow template for row-wise execution.

        The returned flow accepts one row per write:
        - each emitted row is dispatched as one request item (`pipe.request().write(row)`)
        - user-provided stream builder defines downstream map/filter/pipe operators
        """
        if not callable(builder):
            raise TypeError("to_stream_flow builder must be callable.")

        from flownet.compiler.flow_def import FlowDef

        def _flow(init_stream):
            return builder(init_stream)

        _flow.__name__ = str(flow_name)
        return FlowDef(_flow, static=True)

    def run_stream(
        self,
        builder: StreamBuilder,
        *,
        flow_name: str = "dataset_pipeline_stream",
        emit_mode: str = "row",
        emit_batch_size: int = 256,
        request_chunk_size: int | None = None,
        resume_offset: int = 0,
        checkpoint_every: int | None = None,
        on_checkpoint: StreamBridgeCheckpointHook | None = None,
        collect_timeout: float = 60.0,
        late_drain_timeout: float = 1.0,
    ) -> list[Any]:
        """
        Execute current pipeline rows through a DataStream flow and collect outputs.

        This is the minimal bridge for mixed batch/stream composition:
        local preprocessing -> row-wise stream operators -> collected stream outputs.

        Args:
        - emit_mode:
          - `row`: write one row dict per request.write(...)
          - `batch`: write list[dict] batches per request.write(...)
        - emit_batch_size:
          - batch size used when `emit_mode="batch"`.
        - request_chunk_size:
          - max emitted items per request before starting a new request.
            This provides coarse in-flight pressure control.
        - resume_offset / checkpoint_every / on_checkpoint:
          - replay and progress hooks at emitted-payload cursor granularity.
        """
        normalized_emit_mode = _normalize_stream_emit_mode(emit_mode)
        normalized_emit_batch_size = _normalize_stream_emit_batch_size(
            emit_batch_size,
            emit_mode=normalized_emit_mode,
        )
        normalized_request_chunk_size = _normalize_request_chunk_size(request_chunk_size)
        tracker = _StreamBridgeCheckpointTracker(
            flow_name=str(flow_name),
            emit_mode=normalized_emit_mode,
            resume_offset=resume_offset,
            checkpoint_every=checkpoint_every,
            on_checkpoint=on_checkpoint,
            request_chunk_size=normalized_request_chunk_size,
        )

        flow = self.to_stream_flow(builder, flow_name=flow_name)
        payload_iter = _iter_stream_payloads(
            self.iter_rows(),
            emit_mode=normalized_emit_mode,
            emit_batch_size=normalized_emit_batch_size,
        )

        try:
            with flow.open() as pipe:
                if normalized_request_chunk_size is None:
                    outputs, writes = _collect_stream_request(
                        pipe,
                        _iter_resumed_payloads(payload_iter, tracker=tracker),
                        collect_timeout=collect_timeout,
                        late_drain_timeout=late_drain_timeout,
                    )
                    tracker.observe_request_emitted(writes)
                    return outputs

                outputs: list[Any] = []
                chunk: list[Any] = []
                for payload in payload_iter:
                    if not tracker.allow_emit_payload():
                        continue
                    chunk.append(payload)
                    if len(chunk) >= normalized_request_chunk_size:
                        chunk_outputs, writes = _collect_stream_request(
                            pipe,
                            chunk,
                            collect_timeout=collect_timeout,
                            late_drain_timeout=late_drain_timeout,
                        )
                        tracker.observe_request_emitted(writes)
                        outputs.extend(chunk_outputs)
                        chunk = []
                if chunk:
                    chunk_outputs, writes = _collect_stream_request(
                        pipe,
                        chunk,
                        collect_timeout=collect_timeout,
                        late_drain_timeout=late_drain_timeout,
                    )
                    tracker.observe_request_emitted(writes)
                    outputs.extend(chunk_outputs)
                return outputs
        finally:
            tracker.finalize()

    def operator_chain_spec(self) -> dict[str, Any]:
        return {
            "source": {"kind": self._source_config.kind, **deepcopy(self._source_config.detail)},
            "ops": [deepcopy(spec) for spec in self._operator_specs],
        }

    def operator_chain_fingerprint(self) -> str:
        return compute_operator_chain_fingerprint(self.operator_chain_spec())

    def build_run_manifest(
        self,
        *,
        run_id: str,
        dataset_ref: str | None = None,
        snapshot_hash: str | None = None,
        usage: str | None = None,
        profile: str | None = None,
        source_version: str | None = None,
        connector_policy: Mapping[str, Any] | None = None,
        created_at: str | None = None,
    ) -> RunManifest:
        resolved_dataset_ref = dataset_ref or self._default_manifest_fields.get("dataset_ref")
        resolved_snapshot_hash = snapshot_hash or self._default_manifest_fields.get("snapshot_hash")
        if not resolved_dataset_ref:
            raise ValueError("dataset_ref is required to build run manifest.")
        if not resolved_snapshot_hash:
            raise ValueError("snapshot_hash is required to build run manifest.")
        return build_run_manifest(
            run_id=run_id,
            dataset_ref=str(resolved_dataset_ref),
            usage=usage if usage is not None else self._default_manifest_fields.get("usage"),
            profile=profile if profile is not None else self._default_manifest_fields.get("profile"),
            source_version=(
                source_version
                if source_version is not None
                else self._default_manifest_fields.get("source_version")
            ),
            snapshot_hash=str(resolved_snapshot_hash),
            connector_policy=connector_policy,
            operator_chain=self.operator_chain_spec(),
            created_at=created_at,
        )

    def _with_operator(self, op: Operator, spec: dict[str, Any]) -> "DatasetPipeline":
        return DatasetPipeline(
            source_factory=self._source_factory,
            source_config=self._source_config,
            operators=self._operators + (op,),
            operator_specs=self._operator_specs + (dict(spec),),
            default_manifest_fields=self._default_manifest_fields,
        )


def _build_rows_factory(
    rows: Iterable[RowLike] | Callable[[], Iterable[RowLike]],
) -> Callable[[], Iterator[dict[str, Any]]]:
    if callable(rows):
        rows_fn = rows

        def _factory() -> Iterator[dict[str, Any]]:
            for row in rows_fn():
                yield _coerce_row(row)

        return _factory

    snapshot = [_coerce_row(row) for row in rows]

    def _factory() -> Iterator[dict[str, Any]]:
        for row in snapshot:
            yield dict(row)

    return _factory


def _build_join_rows_factory(
    rows: DatasetPipeline | Iterable[RowLike] | Callable[[], Iterable[RowLike]],
) -> Callable[[], Iterator[dict[str, Any]]]:
    if isinstance(rows, DatasetPipeline):
        return lambda: rows.iter_rows()
    return _build_rows_factory(rows)


def _resolve_join_key_fn(key: str | Callable[[RowLike], Any]) -> Callable[[RowLike], Any]:
    if isinstance(key, str):
        key_name = key

        def _fn(row: RowLike) -> Any:
            if key_name not in row:
                raise KeyError(f"join key field missing: {key_name!r}")
            return row[key_name]

        return _fn
    return key


def _join_key_desc(key: str | Callable[[RowLike], Any]) -> str:
    if isinstance(key, str):
        return f"field:{key}"
    return _callable_name(key)


def _merge_join_rows(
    left_row: Mapping[str, Any],
    right_row: Mapping[str, Any],
    *,
    right_prefix: str,
) -> dict[str, Any]:
    merged = dict(left_row)
    for key, value in right_row.items():
        if key not in merged:
            merged[key] = value
            continue
        prefixed = f"{right_prefix}{key}" if right_prefix else key
        if prefixed == key:
            prefixed = f"right_{key}"
        merged[prefixed] = value
    return merged


def _coerce_window_output(
    output: Mapping[str, Any] | Iterable[Mapping[str, Any]],
) -> Iterator[dict[str, Any]]:
    if isinstance(output, Mapping):
        yield dict(output)
        return
    for item in output:
        if not isinstance(item, Mapping):
            raise TypeError(f"window_map function must emit mapping-like rows, got {type(item)}")
        yield dict(item)


def _type_repr(tp: Any) -> str:
    if isinstance(tp, tuple):
        return "|".join(_type_repr(item) for item in tp)
    name = getattr(tp, "__name__", None)
    if isinstance(name, str):
        return name
    return str(tp)


def _normalize_stream_emit_mode(emit_mode: str) -> str:
    normalized = str(emit_mode).strip().lower()
    if normalized not in {"row", "batch"}:
        raise ValueError("run_stream emit_mode must be 'row' or 'batch'.")
    return normalized


def _normalize_stream_emit_batch_size(emit_batch_size: int, *, emit_mode: str) -> int:
    value = int(emit_batch_size)
    if emit_mode == "batch" and value <= 0:
        raise ValueError("run_stream emit_batch_size must be > 0 when emit_mode='batch'.")
    return max(1, value)


def _normalize_request_chunk_size(request_chunk_size: int | None) -> int | None:
    if request_chunk_size is None:
        return None
    value = int(request_chunk_size)
    if value <= 0:
        raise ValueError("run_stream request_chunk_size must be > 0 when provided.")
    return value


def _normalize_resume_offset(resume_offset: int) -> int:
    value = int(resume_offset)
    if value < 0:
        raise ValueError("run_stream resume_offset must be >= 0.")
    return value


def _normalize_checkpoint_every(checkpoint_every: int | None) -> int | None:
    if checkpoint_every is None:
        return None
    value = int(checkpoint_every)
    if value <= 0:
        raise ValueError("run_stream checkpoint_every must be > 0 when provided.")
    return value


class _StreamBridgeCheckpointTracker:
    def __init__(
        self,
        *,
        flow_name: str,
        emit_mode: str,
        resume_offset: int,
        checkpoint_every: int | None,
        on_checkpoint: StreamBridgeCheckpointHook | None,
        request_chunk_size: int | None,
    ):
        self._flow_name = flow_name
        self._emit_mode = emit_mode
        self._resume_offset = _normalize_resume_offset(resume_offset)
        self._checkpoint_every = _normalize_checkpoint_every(checkpoint_every)
        self._on_checkpoint = on_checkpoint
        self._request_chunk_size = request_chunk_size
        self._cursor = 0
        self._emitted_payloads = 0
        self._emitted_requests = 0
        self._last_checkpoint_cursor = -1

    def allow_emit_payload(self) -> bool:
        self._cursor += 1
        self._emit_checkpoint_if_needed()
        if self._cursor <= self._resume_offset:
            return False
        self._emitted_payloads += 1
        return True

    def observe_request_emitted(self, writes: int) -> None:
        if int(writes) > 0:
            self._emitted_requests += 1

    def finalize(self) -> None:
        if self._on_checkpoint is None:
            return
        if self._cursor == self._last_checkpoint_cursor:
            return
        self._on_checkpoint(self._build_checkpoint())
        self._last_checkpoint_cursor = self._cursor

    def _emit_checkpoint_if_needed(self) -> None:
        if self._on_checkpoint is None or self._checkpoint_every is None:
            return
        if self._cursor <= 0 or self._cursor % self._checkpoint_every != 0:
            return
        self._on_checkpoint(self._build_checkpoint())
        self._last_checkpoint_cursor = self._cursor

    def _build_checkpoint(self) -> StreamBridgeCheckpoint:
        return StreamBridgeCheckpoint(
            flow_name=self._flow_name,
            emit_mode=self._emit_mode,
            cursor=self._cursor,
            emitted_payloads=self._emitted_payloads,
            emitted_requests=self._emitted_requests,
            resume_offset=self._resume_offset,
            request_chunk_size=self._request_chunk_size,
        )


def _iter_stream_payloads(
    rows: Iterable[dict[str, Any]],
    *,
    emit_mode: str,
    emit_batch_size: int,
) -> Iterator[Any]:
    if emit_mode == "row":
        for row in rows:
            yield dict(row)
        return

    batch: list[dict[str, Any]] = []
    for row in rows:
        batch.append(dict(row))
        if len(batch) >= emit_batch_size:
            yield list(batch)
            batch = []
    if batch:
        yield list(batch)


def _iter_resumed_payloads(
    payloads: Iterable[Any],
    *,
    tracker: _StreamBridgeCheckpointTracker,
) -> Iterator[Any]:
    for payload in payloads:
        if not tracker.allow_emit_payload():
            continue
        yield payload


def _collect_stream_request(
    pipe: Any,
    payloads: Iterable[Any],
    *,
    collect_timeout: float,
    late_drain_timeout: float,
) -> tuple[list[Any], int]:
    wrote = False
    writes = 0
    with pipe.request() as request:
        for payload in payloads:
            request.write(payload)
            wrote = True
            writes += 1
    if not wrote:
        return [], 0
    return request.collect(timeout=collect_timeout, late_drain_timeout=late_drain_timeout), writes


__all__ = ["DatasetPipeline", "StreamBridgeCheckpoint"]
