dovla_cil/data/sharding.py

from __future__ import annotations

from collections import Counter, OrderedDict
from datetime import datetime, timezone
from pathlib import Path
from typing import Iterable, Iterator

from dovla_cil.data.schema import CILRecord, CIL_VERSION
from dovla_cil.utils.io import ensure_dir, iter_jsonl, write_json, write_jsonl


def group_records(records: Iterable[CILRecord]) -> OrderedDict[str, list[CILRecord]]:
    groups: OrderedDict[str, list[CILRecord]] = OrderedDict()
    for record in records:
        groups.setdefault(record.group_id, []).append(record)
    return groups


def split_records_by_group(
    records: Iterable[CILRecord], *, max_records_per_shard: int
) -> list[list[CILRecord]]:
    if max_records_per_shard <= 0:
        raise ValueError("max_records_per_shard must be positive")

    shards: list[list[CILRecord]] = []
    current: list[CILRecord] = []
    for group in group_records(records).values():
        if current and len(current) + len(group) > max_records_per_shard:
            shards.append(current)
            current = []
        current.extend(group)
    if current:
        shards.append(current)
    return shards


class ShardWriter:
    """Streaming writer for grouped CIL JSONL shards.

    Records should be written group-contiguously. This keeps a single CIL group in one shard and
    makes the group index compact and deterministic. If a closed group appears again, the writer
    raises an error instead of silently creating a broken index.
    """

    def __init__(
        self,
        output_dir: str | Path,
        *,
        dataset_name: str = "dovla_cil",
        backend: str = "unknown",
        k: int | None = None,
        task_count: int | None = None,
        seed: int | None = None,
        shard_size: int = 1024,
        schema_version: str = CIL_VERSION,
        version: str = "0.1",
        shard_format: str = "jsonl",
        index_format: str = "jsonl",
        overwrite: bool = True,
    ) -> None:
        if shard_size <= 0:
            raise ValueError("shard_size must be positive")
        self.output_dir = ensure_dir(output_dir)
        self.shards_dir = ensure_dir(self.output_dir / "shards")
        self.dataset_name = dataset_name
        self.backend = backend
        self.k = k
        self.task_count = task_count
        self.seed = seed
        self.shard_size = int(shard_size)
        self.schema_version = schema_version
        self.version = version
        self.shard_format = shard_format.lower()
        self.index_format = index_format.lower()
        if self.shard_format not in {"jsonl", "parquet"}:
            raise ValueError("shard_format must be 'jsonl' or 'parquet'")
        if self.index_format not in {"jsonl", "parquet"}:
            raise ValueError("index_format must be 'jsonl' or 'parquet'")
        if overwrite:
            self._remove_stale_outputs()

        self._current_group_id: str | None = None
        self._current_group: list[CILRecord] = []
        self._closed_groups: set[str] = set()
        self._current_shard: list[CILRecord] = []
        self._current_shard_group_ids: list[str] = []
        self._shard_index = 0
        self._shard_entries: list[dict[str, object]] = []
        self._group_index_entries: list[dict[str, object]] = []
        self._record_index_entries: list[dict[str, object]] = []
        self._task_ids: set[str] = set()
        self._num_records = 0
        self._closed = False

    def write(self, record: CILRecord) -> None:
        if self._closed:
            raise RuntimeError("Cannot write to a closed ShardWriter")
        record.validate()
        self._task_ids.add(record.task_id)
        if self._current_group_id is None:
            self._current_group_id = record.group_id
        if record.group_id != self._current_group_id:
            self._flush_group()
            if record.group_id in self._closed_groups:
                raise ValueError(
                    f"Group {record.group_id!r} was already written. "
                    "ShardWriter requires group-contiguous records."
                )
            self._current_group_id = record.group_id
        self._current_group.append(record)

    def close(self) -> dict[str, object]:
        if self._closed:
            return self._metadata()
        self._flush_group()
        self._flush_shard()
        metadata = self._metadata()
        self._write_indices_and_metadata(metadata)
        self._closed = True
        return metadata

    def _flush_group(self) -> None:
        if not self._current_group:
            return
        group = list(self._current_group)
        group_id = group[0].group_id
        if any(record.group_id != group_id for record in group):
            raise ValueError("Internal writer error: mixed group buffer")
        if self._current_shard and len(self._current_shard) + len(group) > self.shard_size:
            self._flush_shard()

        shard_path = self._current_shard_path()
        row_offset = len(self._current_shard)
        self._current_shard.extend(group)
        self._current_shard_group_ids.append(group_id)
        self._group_index_entries.append(_make_group_index_entry(group, shard_path=shard_path))
        for row_index, record in enumerate(group, start=row_offset):
            self._record_index_entries.append(
                _make_record_index_entry(record, shard_path=shard_path, row_index=row_index)
            )

        self._num_records += len(group)
        self._closed_groups.add(group_id)
        self._current_group = []
        self._current_group_id = None
        if len(self._current_shard) >= self.shard_size:
            self._flush_shard()

    def _flush_shard(self) -> None:
        if not self._current_shard:
            return
        shard_path = self._current_shard_path()
        target = self.output_dir / shard_path
        if self.shard_format == "jsonl":
            write_jsonl((record.to_dict() for record in self._current_shard), target)
        else:
            _write_parquet_rows([record.to_dict() for record in self._current_shard], target)
        self._shard_entries.append(
            {
                "path": shard_path,
                "record_count": len(self._current_shard),
                "group_ids": list(self._current_shard_group_ids),
                "format": self.shard_format,
            }
        )
        self._current_shard = []
        self._current_shard_group_ids = []
        self._shard_index += 1

    def _current_shard_path(self) -> str:
        suffix = "jsonl" if self.shard_format == "jsonl" else "parquet"
        return f"shards/shard_{self._shard_index:06d}.{suffix}"

    def _metadata(self) -> dict[str, object]:
        inferred_k = self.k
        if inferred_k is None and self._group_index_entries:
            inferred_k = max(int(entry["num_records"]) for entry in self._group_index_entries)
        if inferred_k is None:
            inferred_k = 0
        task_count = self.task_count if self.task_count is not None else len(self._task_ids)
        return {
            "dataset_name": self.dataset_name,
            "version": self.version,
            "created_at": datetime.now(timezone.utc).isoformat(),
            "backend": self.backend,
            "num_groups": len(self._group_index_entries),
            "num_records": self._num_records,
            "k": int(inferred_k),
            "task_count": int(task_count),
            "seed": self.seed,
            "schema_version": self.schema_version,
            "format": "dovla_cil",
            "shard_format": self.shard_format,
            "index_format": self.index_format,
            "shard_size": self.shard_size,
            "shard_count": len(self._shard_entries),
            "group_index_path": _index_path("group_index", self.index_format),
            "record_index_path": _index_path("record_index", self.index_format),
            "shards": list(self._shard_entries),
            # Compatibility aliases for the original manifest shape.
            "record_count": self._num_records,
            "group_count": len(self._group_index_entries),
        }

    def _write_indices_and_metadata(self, metadata: dict[str, object]) -> None:
        _write_index_rows(
            self._group_index_entries,
            self.output_dir / str(metadata["group_index_path"]),
            index_format=self.index_format,
        )
        _write_index_rows(
            self._record_index_entries,
            self.output_dir / str(metadata["record_index_path"]),
            index_format=self.index_format,
        )
        write_json(metadata, self.output_dir / "metadata.json")
        write_json(metadata, self.output_dir / "manifest.json")

    def _remove_stale_outputs(self) -> None:
        for pattern in (
            "metadata.json",
            "manifest.json",
            "group_index.jsonl",
            "record_index.jsonl",
            "group_index.parquet",
            "record_index.parquet",
        ):
            target = self.output_dir / pattern
            if target.exists() and target.is_file():
                target.unlink()
        for target in self.shards_dir.glob("shard_*.*"):
            if target.is_file():
                target.unlink()


class ShardReader:
    """Read CIL datasets written by `ShardWriter`."""

    def __init__(self, dataset_path: str | Path) -> None:
        from dovla_cil.data.index import ShardIndex

        self.index = ShardIndex.from_path(dataset_path)

    def iterate_records(self) -> Iterator[CILRecord]:
        for shard_path in self.index.shards:
            yield from iter_cil_records(shard_path)

    def iterate_groups(self) -> Iterator[list[CILRecord]]:
        for entry in self.index.load_group_index():
            yield self.load_group(str(entry["group_id"]))

    def load_group(self, group_id: str) -> list[CILRecord]:
        entry = self.index.group_entry(group_id)
        shard_path = self.index.dataset_dir / str(entry["shard_path"])
        records = [record for record in iter_cil_records(shard_path) if record.group_id == group_id]
        expected = int(entry["num_records"])
        if len(records) != expected:
            raise ValueError(
                f"Group {group_id!r} index expected {expected} records, loaded {len(records)}"
            )
        return records


def write_cil_shards(
    records: Iterable[CILRecord],
    *,
    output_dir: str | Path,
    max_records_per_shard: int = 1024,
    prefix: str = "cil",
    dataset_name: str = "dovla_cil",
    backend: str = "unknown",
    k: int | None = None,
    task_count: int | None = None,
    seed: int | None = None,
) -> dict[str, object]:
    del prefix  # ShardWriter uses the canonical shard_000000 layout.
    materialized = list(records)
    writer = ShardWriter(
        output_dir,
        dataset_name=dataset_name,
        backend=backend,
        k=k,
        task_count=task_count,
        seed=seed,
        shard_size=max_records_per_shard,
    )
    for group in group_records(materialized).values():
        for record in group:
            writer.write(record)
    return writer.close()


def iter_cil_records(path: str | Path) -> Iterator[CILRecord]:
    record_path = Path(path)
    if record_path.is_dir() or record_path.name in {"metadata.json", "manifest.json"}:
        yield from ShardReader(record_path).iterate_records()
        return
    if record_path.suffix == ".jsonl":
        for payload in iter_jsonl(record_path):
            yield CILRecord.from_dict(payload)
        return
    if record_path.suffix == ".parquet":
        for payload in _read_parquet_rows(record_path):
            yield CILRecord.from_dict(payload)
        return
    raise ValueError(f"Unsupported CIL record path: {record_path}")


def _make_group_index_entry(group: list[CILRecord], *, shard_path: str) -> dict[str, object]:
    first = group[0]
    candidate_counts = Counter(record.candidate_type for record in group)
    rewards = [record.reward.progress for record in group]
    return {
        "group_id": first.group_id,
        "shard_path": shard_path,
        "record_ids": [record.record_id for record in group],
        "task_id": first.task_id,
        "state_hash": first.state_hash,
        "num_records": len(group),
        "max_reward": max(rewards) if rewards else 0.0,
        "success_count": sum(1 for record in group if record.reward.terminal_success),
        "candidate_type_counts": dict(sorted(candidate_counts.items())),
        # Useful extras for inspection and compatibility with older group_index.jsonl files.
        "scene_id": first.scene_id,
        "instruction": first.instruction,
        "state_blob_ref": first.metadata.get("state_blob_ref"),
    }


def _make_record_index_entry(
    record: CILRecord, *, shard_path: str, row_index: int
) -> dict[str, object]:
    return {
        "record_id": record.record_id,
        "group_id": record.group_id,
        "shard_path": shard_path,
        "row_index": row_index,
        "task_id": record.task_id,
        "state_hash": record.state_hash,
        "candidate_type": record.candidate_type,
        "reward_progress": record.reward.progress,
        "success": record.reward.terminal_success,
        "regret": record.regret,
        "rank_within_group": record.rank_within_group,
        "failure_type": record.failure.type if record.failure else None,
    }


def _index_path(stem: str, index_format: str) -> str:
    suffix = "jsonl" if index_format == "jsonl" else "parquet"
    return f"{stem}.{suffix}"


def _write_index_rows(rows: list[dict[str, object]], path: Path, *, index_format: str) -> None:
    if index_format == "jsonl":
        write_jsonl(rows, path)
        return
    _write_parquet_rows(rows, path)


def _write_parquet_rows(rows: list[dict[str, object]], path: Path) -> None:
    try:
        import pandas as pd
    except ImportError as exc:  # pragma: no cover - optional dependency path
        raise ImportError("Parquet output requires pandas and pyarrow.") from exc
    ensure_dir(path.parent)
    pd.DataFrame(rows).to_parquet(path, index=False)


def _read_parquet_rows(path: Path) -> list[dict[str, object]]:
    try:
        import pandas as pd
    except ImportError as exc:  # pragma: no cover - optional dependency path
        raise ImportError("Parquet input requires pandas and pyarrow.") from exc
    return pd.read_parquet(path).to_dict(orient="records")