# DOCS_REF:
# - docs/user/reference/api-surface.md
# - docs/user/12-ray-facade-and-benchmark-mapping.md
# DOCS_SYNC:
# - If Dataset facade signatures or fallback semantics change,
#   update the referenced docs and benchmark migration notes.

from __future__ import annotations

from collections import deque
from dataclasses import dataclass
import hashlib
import importlib.util
import json
from pathlib import Path
from typing import Any, Callable, Iterable, Iterator, Mapping

from flownet.api.actor_method_ref import ActorMethodRef
from flownet.core.stateless import StatelessOpRef
from flownet.compiler.flow_def import FlowDef
from flownet.runtime.flows.replay_checkpoint import parse_replay_checkpoint


SINK_IDEMPOTENCY_META_FIELD = "__flownet_sink_meta"
_SUPPORTED_BATCH_FORMATS = {"python", "numpy", "pandas", "pyarrow"}
_SUPPORTED_FALLBACK_MODES = {"raise", "fallback_to_python"}


@dataclass(frozen=True)
class DatasetWriteResult:
    output_path: str
    rows_written: int
    duplicates_dropped: int
    emitted_requests: int
    idempotent_write: bool
    metadata_field: str
    execution_metrics: "DatasetExecutionMetrics | None" = None


@dataclass(frozen=True)
class DatasetExecutionMetrics:
    source_batches: int
    source_rows: int
    emitted_requests: int
    output_batches: int
    output_rows: int
    map_op_count: int
    map_requested_batch_formats: tuple[str, ...]
    map_effective_batch_format: str | None
    map_fallback_to_python_ops: int
    map_native_pyarrow_ops: int

    def as_dict(self) -> dict[str, Any]:
        return {
            "source_batches": int(self.source_batches),
            "source_rows": int(self.source_rows),
            "emitted_requests": int(self.emitted_requests),
            "output_batches": int(self.output_batches),
            "output_rows": int(self.output_rows),
            "map_op_count": int(self.map_op_count),
            "map_requested_batch_formats": list(self.map_requested_batch_formats),
            "map_effective_batch_format": self.map_effective_batch_format,
            "map_fallback_to_python_ops": int(self.map_fallback_to_python_ops),
            "map_native_pyarrow_ops": int(self.map_native_pyarrow_ops),
        }


@dataclass(frozen=True)
class _MapBatchesOp:
    target: ActorMethodRef | StatelessOpRef
    batch_size: int | None
    requested_batch_format: str
    effective_batch_format: str
    fallback_to_python: bool


class Dataset:
    """
    Minimal Dataset facade for benchmark-oriented read/map_batches/write flows.

    Notes:
    - `map_batches` executes with Python list[dict] batches by default.
    - `batch_format="pyarrow"` can execute natively when the target declares
      pyarrow support and runtime dependency is available.
    - Other non-`python` batch formats can either fail (`raise`) or fall back
      to Python list batches (`fallback_to_python`).
    - Without map ops, Dataset iteration is local/in-process.
    - With map ops, Dataset compiles a FlowDef and executes through runtime pipe,
      which is the execution/distribution boundary for this facade.
    """

    def __init__(
        self,
        *,
        source_factory: Callable[[], Iterator[dict[str, Any]]],
        map_ops: tuple[_MapBatchesOp, ...] = (),
    ):
        self._source_factory = source_factory
        self._map_ops = map_ops
        self._last_execution_metrics: DatasetExecutionMetrics | None = None

    def map_batches(
        self,
        fn_object: ActorMethodRef | StatelessOpRef,
        *,
        batch_size: int | None = None,
        batch_format: str = "python",
        on_unsupported_batch_format: str = "fallback_to_python",
    ) -> "Dataset":
        target = _ensure_dataset_target(fn_object)
        normalized_batch_size = _normalize_batch_size(batch_size)

        requested_format = _normalize_batch_format(batch_format)
        fallback_mode = _normalize_fallback_mode(on_unsupported_batch_format)

        effective_format, fallback_to_python = _resolve_effective_map_batch_format(
            target,
            requested_format=requested_format,
            fallback_mode=fallback_mode,
        )

        if self._map_ops:
            existing_effective = self._map_ops[0].effective_batch_format
            if effective_format != existing_effective:
                if fallback_mode == "raise":
                    raise ValueError(
                        "map_batches in one Dataset chain must use one effective batch format. "
                        f"Existing: {existing_effective!r}, new: {effective_format!r}."
                    )
                effective_format = existing_effective
                fallback_to_python = effective_format == "python"

        op = _MapBatchesOp(
            target=target,
            batch_size=normalized_batch_size,
            requested_batch_format=requested_format,
            effective_batch_format=effective_format,
            fallback_to_python=fallback_to_python,
        )
        return Dataset(
            source_factory=self._source_factory,
            map_ops=self._map_ops + (op,),
        )

    def iter_rows(
        self,
        *,
        batch_size: int | None = None,
        timeout: float = 60.0,
    ) -> Iterator[dict[str, Any]]:
        for _, rows in self._iter_batches_with_request_id(batch_size=batch_size, timeout=timeout):
            for row in rows:
                yield row

    def iter_batches(
        self,
        *,
        batch_size: int | None = None,
        batch_format: str = "python",
        on_unsupported_batch_format: str = "fallback_to_python",
        timeout: float = 60.0,
    ) -> Iterator[Any]:
        requested_format = _normalize_batch_format(batch_format)
        fallback_mode = _normalize_fallback_mode(on_unsupported_batch_format)

        for _, rows in self._iter_batches_with_request_id(batch_size=batch_size, timeout=timeout):
            yield _rows_to_batch(
                rows,
                requested_format,
                fallback_mode=fallback_mode,
            )

    def write_parquet(
        self,
        output_path: str | Path,
        *,
        batch_size: int | None = None,
        idempotent_write: bool = False,
        stage_id: str = "flownet-dataset-sink",
        metadata_field: str = SINK_IDEMPOTENCY_META_FIELD,
        passthrough_idempotency_metadata: bool = True,
        max_dedupe_keys: int = 100_000,
        timeout: float = 60.0,
    ) -> DatasetWriteResult:
        try:
            import pyarrow as pa
            import pyarrow.parquet as pq
        except Exception as exc:  # pragma: no cover
            raise RuntimeError("write_parquet requires pyarrow in the runtime environment.") from exc

        dedupe_budget = max(1, int(max_dedupe_keys))
        seen_dedupe_keys: set[str] = set()
        dedupe_order: deque[str] = deque()

        output_path = str(output_path)
        writer = None
        rows_written = 0
        duplicates_dropped = 0
        emitted_requests = 0
        seen_requests: set[str] = set()

        try:
            for request_id, rows in self._iter_batches_with_request_id(
                batch_size=batch_size,
                timeout=timeout,
            ):
                if request_id not in seen_requests:
                    emitted_requests += 1
                    seen_requests.add(request_id)

                candidate_rows = [dict(row) for row in rows]
                if idempotent_write:
                    candidate_rows = _ensure_idempotency_metadata(
                        candidate_rows,
                        request_id=request_id,
                        stage_id=stage_id,
                        metadata_field=metadata_field,
                        passthrough_idempotency_metadata=passthrough_idempotency_metadata,
                    )

                deduped_rows: list[dict[str, Any]] = []
                for row in candidate_rows:
                    if not idempotent_write:
                        deduped_rows.append(row)
                        continue

                    meta = row.get(metadata_field)
                    row_dedupe_key = meta.get("row_dedupe_key") if isinstance(meta, dict) else None
                    if not row_dedupe_key:
                        deduped_rows.append(row)
                        continue
                    if row_dedupe_key in seen_dedupe_keys:
                        duplicates_dropped += 1
                        continue
                    seen_dedupe_keys.add(row_dedupe_key)
                    dedupe_order.append(row_dedupe_key)
                    while len(dedupe_order) > dedupe_budget:
                        oldest = dedupe_order.popleft()
                        seen_dedupe_keys.discard(oldest)
                    deduped_rows.append(row)

                if not deduped_rows:
                    continue

                table = pa.Table.from_pylist(deduped_rows)
                if writer is None:
                    writer = pq.ParquetWriter(output_path, table.schema)
                writer.write_table(table)
                rows_written += len(deduped_rows)
        finally:
            if writer is not None:
                writer.close()

        return DatasetWriteResult(
            output_path=output_path,
            rows_written=rows_written,
            duplicates_dropped=duplicates_dropped,
            emitted_requests=emitted_requests,
            idempotent_write=bool(idempotent_write),
            metadata_field=metadata_field,
            execution_metrics=self.last_execution_metrics(),
        )

    def last_execution_metrics(self) -> DatasetExecutionMetrics | None:
        return self._last_execution_metrics

    def _iter_batches_with_request_id(
        self,
        *,
        batch_size: int | None,
        timeout: float,
    ) -> Iterator[tuple[str, list[dict[str, Any]]]]:
        self._last_execution_metrics = None
        metrics = _DatasetExecutionMetricsBuilder(self._map_ops)

        try:
            effective_batch_size = self._resolve_input_batch_size(batch_size)
            row_iter = self._source_factory()

            if not self._map_ops:
                # Local passthrough path: no flow compilation/runtime dispatch.
                for index, batch in enumerate(_chunked(row_iter, effective_batch_size), start=1):
                    rows = [_coerce_row(item) for item in batch]
                    if not rows:
                        continue
                    metrics.observe_input_rows(len(rows))
                    metrics.observe_emitted_request()
                    metrics.observe_output_rows(len(rows))
                    yield (f"local:{index}", rows)
                return

            input_format = self._map_ops[0].effective_batch_format
            output_format = self._map_ops[-1].effective_batch_format
            fallback_mode = (
                "fallback_to_python" if any(op.fallback_to_python for op in self._map_ops) else "raise"
            )

            # Runtime execution path: compile map chain to a static flow and run via pipe.
            pipeline = self._compile_pipeline()
            pipe = pipeline.open()

            for batch in _chunked(row_iter, effective_batch_size):
                rows = [_coerce_row(item) for item in batch]
                if not rows:
                    continue
                metrics.observe_input_rows(len(rows))
                request_id = pipe.write(
                    _rows_to_batch(
                        rows,
                        input_format,
                        fallback_mode=fallback_mode,
                    )
                )
                metrics.observe_emitted_request()
                while True:
                    out = pipe.read(timeout=timeout, request_id=request_id)
                    if _is_run_output_done(out):
                        break
                    if _is_run_output_error(out):
                        error_type = getattr(out, "error_type", "UnknownError")
                        message = getattr(out, "message", "")
                        traceback_text = getattr(out, "traceback_text", "")
                        raise RuntimeError(
                            f"Dataset map_batches flow failed: {error_type}: {message}\n{traceback_text}"
                        )
                    out_rows = _batch_to_rows(
                        out,
                        expected_format=output_format,
                        fallback_mode=fallback_mode,
                    )
                    if out_rows:
                        metrics.observe_output_rows(len(out_rows))
                        yield request_id, out_rows
        finally:
            self._last_execution_metrics = metrics.build()

    def _resolve_input_batch_size(self, override_batch_size: int | None) -> int:
        if override_batch_size is not None:
            return _normalize_batch_size(override_batch_size)
        for op in self._map_ops:
            if op.batch_size is not None:
                return op.batch_size
        return 256

    def _compile_pipeline(self) -> FlowDef:
        map_ops = tuple(self._map_ops)

        def _pipeline(init_stream):
            stream = init_stream
            for op in map_ops:
                stream = stream.map(op.target)
            return stream

        return FlowDef(_pipeline, static=True)


def read_rows(
    rows: Iterable[Mapping[str, Any]] | Callable[[], Iterable[Mapping[str, Any]]],
) -> Dataset:
    if callable(rows):
        source_factory = rows

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

        return Dataset(source_factory=_factory)

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

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

    return Dataset(source_factory=_factory)


def read_parquet(path: str | Path, *, columns: list[str] | None = None) -> Dataset:
    parquet_path = Path(path)
    if not parquet_path.exists():
        raise FileNotFoundError(f"Parquet path not found: {path}")

    selected_columns = list(columns) if columns is not None else None

    def _factory() -> Iterator[dict[str, Any]]:
        try:
            import pandas as pd

            frame = pd.read_parquet(parquet_path, columns=selected_columns)
            for row in frame.to_dict(orient="records"):
                yield _coerce_row(row)
            return
        except Exception:
            pass

        try:
            import pyarrow.parquet as pq
        except Exception as exc:  # pragma: no cover
            raise RuntimeError("read_parquet fallback requires pyarrow.") from exc

        table = pq.read_table(parquet_path, columns=selected_columns)
        for row in table.to_pylist():
            yield _coerce_row(row)

    return Dataset(source_factory=_factory)


def read_images(
    path: str | Path,
    *,
    mode: str = "RGB",
    on_error: str = "skip",
) -> Dataset:
    normalized_on_error = str(on_error).strip().lower()
    if normalized_on_error not in {"skip", "raise"}:
        raise ValueError("read_images on_error must be 'skip' or 'raise'.")

    input_path = Path(path)
    if input_path.is_dir():
        files = sorted(x for x in input_path.rglob("*") if x.is_file())
    elif input_path.is_file():
        files = [input_path]
    else:
        raise FileNotFoundError(f"Image path not found: {path}")

    def _factory() -> Iterator[dict[str, Any]]:
        try:
            import numpy as np
            from PIL import Image
        except Exception as exc:  # pragma: no cover
            raise RuntimeError("read_images requires pillow and numpy.") from exc

        for image_file in files:
            try:
                with Image.open(image_file) as img:
                    arr = np.array(img.convert(mode))
                yield {
                    "path": str(image_file),
                    "image": arr,
                }
            except Exception as exc:
                if normalized_on_error == "raise":
                    raise RuntimeError(f"Failed to decode image: {image_file}") from exc
                continue

    return Dataset(source_factory=_factory)


def read_tfrecords(path: str | Path) -> Dataset:
    input_path = Path(path)
    if input_path.is_dir():
        files = sorted(str(x) for x in input_path.rglob("*") if x.is_file())
    elif input_path.is_file():
        files = [str(input_path)]
    else:
        raise FileNotFoundError(f"TFRecord path not found: {path}")

    def _factory() -> Iterator[dict[str, Any]]:
        try:
            import tensorflow as tf
        except Exception as exc:  # pragma: no cover
            raise RuntimeError("read_tfrecords requires tensorflow.") from exc

        dataset = tf.data.TFRecordDataset(files)
        for record in dataset:
            yield {"record": bytes(record.numpy())}

    return Dataset(source_factory=_factory)


def _ensure_dataset_target(target: Any) -> ActorMethodRef | StatelessOpRef:
    if isinstance(target, (ActorMethodRef, StatelessOpRef)):
        return target
    raise TypeError(
        "map_batches expects ActorMethodRef or StatelessOpRef. "
        "Create an actor first (create_actor/ActorRef) or use stateless_op()."
    )


def _normalize_batch_size(batch_size: int | None) -> int:
    value = 256 if batch_size is None else int(batch_size)
    if value <= 0:
        raise ValueError("batch_size must be > 0.")
    return value


def _normalize_batch_format(batch_format: str) -> str:
    normalized = str(batch_format).strip().lower()
    if normalized not in _SUPPORTED_BATCH_FORMATS:
        raise ValueError(
            f"Unsupported batch_format: {batch_format!r}. "
            f"Supported values: {sorted(_SUPPORTED_BATCH_FORMATS)}"
        )
    return normalized


def _normalize_fallback_mode(mode: str) -> str:
    normalized = str(mode).strip().lower()
    if normalized not in _SUPPORTED_FALLBACK_MODES:
        raise ValueError(
            f"Unsupported fallback mode: {mode!r}. "
            f"Supported values: {sorted(_SUPPORTED_FALLBACK_MODES)}"
        )
    return normalized


def _resolve_effective_map_batch_format(
    target: ActorMethodRef | StatelessOpRef,
    *,
    requested_format: str,
    fallback_mode: str,
) -> tuple[str, bool]:
    if requested_format == "python":
        return "python", False

    if requested_format == "pyarrow":
        if _target_supports_pyarrow(target) and _has_pyarrow_runtime():
            return "pyarrow", False

        if fallback_mode == "raise":
            if not _target_supports_pyarrow(target):
                raise ValueError(
                    "map_batches target does not declare native pyarrow batch support. "
                    "Declare metadata batch_formats=['pyarrow'] (or supports_pyarrow=True), "
                    "or use on_unsupported_batch_format='fallback_to_python'."
                )
            raise ValueError(
                "map_batches requested pyarrow native mode but pyarrow runtime dependency is unavailable. "
                "Install pyarrow or use on_unsupported_batch_format='fallback_to_python'."
            )
        return "python", True

    if fallback_mode == "raise":
        raise ValueError(
            "map_batches currently supports batch_format='python' only. "
            "Use on_unsupported_batch_format='fallback_to_python' to degrade gracefully."
        )
    return "python", True


class _DatasetExecutionMetricsBuilder:
    def __init__(self, map_ops: tuple[_MapBatchesOp, ...]):
        self.source_batches = 0
        self.source_rows = 0
        self.emitted_requests = 0
        self.output_batches = 0
        self.output_rows = 0
        self._map_ops = map_ops

    def observe_input_rows(self, row_count: int) -> None:
        self.source_batches += 1
        self.source_rows += int(row_count)

    def observe_emitted_request(self) -> None:
        self.emitted_requests += 1

    def observe_output_rows(self, row_count: int) -> None:
        self.output_batches += 1
        self.output_rows += int(row_count)

    def build(self) -> DatasetExecutionMetrics:
        map_formats = tuple(op.requested_batch_format for op in self._map_ops)
        effective = self._map_ops[-1].effective_batch_format if self._map_ops else None
        map_fallback_to_python_ops = sum(1 for op in self._map_ops if op.fallback_to_python)
        map_native_pyarrow_ops = sum(
            1
            for op in self._map_ops
            if op.requested_batch_format == "pyarrow"
            and op.effective_batch_format == "pyarrow"
            and not op.fallback_to_python
        )
        return DatasetExecutionMetrics(
            source_batches=self.source_batches,
            source_rows=self.source_rows,
            emitted_requests=self.emitted_requests,
            output_batches=self.output_batches,
            output_rows=self.output_rows,
            map_op_count=len(self._map_ops),
            map_requested_batch_formats=map_formats,
            map_effective_batch_format=effective,
            map_fallback_to_python_ops=map_fallback_to_python_ops,
            map_native_pyarrow_ops=map_native_pyarrow_ops,
        )


def _target_supports_pyarrow(target: ActorMethodRef | StatelessOpRef) -> bool:
    declared = _extract_declared_batch_formats(target)
    if "pyarrow" in declared:
        return True
    return bool(getattr(target, "supports_pyarrow", False))


def _extract_declared_batch_formats(target: ActorMethodRef | StatelessOpRef) -> set[str]:
    values: Any = None
    if isinstance(target, StatelessOpRef):
        metadata = target.config.metadata if target.config is not None else {}
        if isinstance(metadata, Mapping):
            values = metadata.get("batch_formats")
            if values is None:
                values = metadata.get("native_batch_formats")
            if values is None and metadata.get("supports_pyarrow") is True:
                return {"pyarrow"}
    else:
        values = getattr(target, "__flownet_batch_formats__", None)
        if values is None:
            values = getattr(target, "batch_formats", None)

    if values is None:
        return set()
    if isinstance(values, str):
        return {values.strip().lower()} if values.strip() else set()
    if isinstance(values, (list, tuple, set)):
        return {str(item).strip().lower() for item in values if str(item).strip()}
    return set()


def _has_pyarrow_runtime() -> bool:
    return importlib.util.find_spec("pyarrow") is not None


def _chunked(items: Iterable[Any], batch_size: int) -> Iterator[list[Any]]:
    if batch_size <= 0:
        raise ValueError("batch_size must be > 0")
    buf: list[Any] = []
    for item in items:
        buf.append(item)
        if len(buf) >= batch_size:
            yield buf
            buf = []
    if buf:
        yield buf


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


def _is_run_output_done(item: Any) -> bool:
    return item.__class__.__name__ == "RunOutputDone"


def _is_run_output_error(item: Any) -> bool:
    return item.__class__.__name__ == "RunOutputError"


def _rows_to_batch(rows: list[dict[str, Any]], fmt: str, *, fallback_mode: str) -> Any:
    if fmt == "python":
        return rows

    try:
        if fmt == "pandas":
            import pandas as pd

            return pd.DataFrame(rows)

        if fmt == "pyarrow":
            import pyarrow as pa

            return pa.Table.from_pylist(rows)

        if fmt == "numpy":
            import numpy as np

            if not rows:
                return {}
            keys = sorted(rows[0].keys())
            return {key: np.array([row.get(key) for row in rows], dtype=object) for key in keys}

        raise ValueError(f"Unsupported batch format: {fmt}")
    except Exception:
        if fallback_mode == "fallback_to_python":
            return rows
        raise


def _batch_to_rows(batch: Any, *, expected_format: str, fallback_mode: str) -> list[dict[str, Any]]:
    if isinstance(batch, list):
        return [_coerce_row(item) for item in batch]

    if isinstance(batch, tuple):
        return [_coerce_row(item) for item in batch]

    if expected_format == "pyarrow":
        try:
            rows = batch.to_pylist()
            return [_coerce_row(item) for item in rows]
        except Exception:
            if fallback_mode == "fallback_to_python":
                rows = _batch_to_rows_from_common_types(batch)
                if rows is not None:
                    return rows
            raise TypeError(
                "map_batches expected pyarrow.Table-like output (to_pylist), "
                f"got {type(batch)}."
            )

    rows = _batch_to_rows_from_common_types(batch)
    if rows is not None:
        return rows

    raise TypeError(
        "map_batches expected list batch output, "
        f"got {type(batch)}."
    )


def _batch_to_rows_from_common_types(batch: Any) -> list[dict[str, Any]] | None:
    to_dict = getattr(batch, "to_dict", None)
    if callable(to_dict):
        try:
            maybe = to_dict(orient="records")
            if isinstance(maybe, list):
                return [_coerce_row(item) for item in maybe]
        except Exception:
            pass
    return None


def _stable_row_fingerprint(row: Mapping[str, Any]) -> str:
    encoded = json.dumps(row, sort_keys=True, separators=(",", ":"), default=str)
    return hashlib.sha1(encoded.encode("utf-8")).hexdigest()


def _build_sink_idempotency_metadata(
    *,
    request_id: str,
    stage_id: str,
    sys_meta: Mapping[str, Any] | None,
    row_fingerprint: str,
) -> dict[str, Any]:
    checkpoint = parse_replay_checkpoint(
        {"sys": dict(sys_meta or {})},
        request_id=request_id,
        stage_id=stage_id,
    )
    return checkpoint.sink_idempotency_payload(row_fingerprint=row_fingerprint)


def _ensure_idempotency_metadata(
    rows: list[dict[str, Any]],
    *,
    request_id: str,
    stage_id: str,
    metadata_field: str,
    passthrough_idempotency_metadata: bool,
) -> list[dict[str, Any]]:
    annotated: list[dict[str, Any]] = []
    for raw in rows:
        row = dict(raw)
        current_meta = row.get(metadata_field)
        if (
            passthrough_idempotency_metadata
            and isinstance(current_meta, Mapping)
            and current_meta.get("sink_idempotency_key")
        ):
            annotated.append(row)
            continue

        sys_meta = row.get("sys") if isinstance(row.get("sys"), Mapping) else None
        fingerprint_row = dict(row)
        fingerprint_row.pop(metadata_field, None)
        fingerprint = _stable_row_fingerprint(fingerprint_row)
        row[metadata_field] = _build_sink_idempotency_metadata(
            request_id=request_id,
            stage_id=stage_id,
            sys_meta=sys_meta,
            row_fingerprint=fingerprint,
        )
        annotated.append(row)
    return annotated


__all__ = [
    "Dataset",
    "DatasetExecutionMetrics",
    "DatasetWriteResult",
    "SINK_IDEMPOTENCY_META_FIELD",
    "read_rows",
    "read_parquet",
    "read_images",
    "read_tfrecords",
]
