dovla_cil/data/datasets.py

from __future__ import annotations

from dataclasses import dataclass, replace
from pathlib import Path
from typing import Any, Iterable, Iterator

from dovla_cil.data.index import ShardIndex
from dovla_cil.data.schema import CILRecord
from dovla_cil.data.sharding import ShardReader, group_records, iter_cil_records
from dovla_cil.utils.io import read_json, write_json


@dataclass(frozen=True)
class CILCollectionIndex:
    dataset_dir: Path
    metadata: dict[str, Any]
    record_count: int
    group_count: int


class CILDataset:
    """Indexed CIL dataset backed by robust shard metadata and indices."""

    def __init__(self, dataset_dir: str | Path) -> None:
        self.dataset_dir = Path(dataset_dir)
        collection_path = self.dataset_dir / "collection.json"
        self.readers: list[ShardReader] = []
        if collection_path.exists():
            payload = read_json(collection_path)
            sources = _collection_sources(payload, base_dir=self.dataset_dir)
            if not sources:
                raise ValueError("CIL collection must contain at least one source")
            records: list[CILRecord] = []
            source_metadata: list[dict[str, Any]] = []
            for source in sources:
                reader = ShardReader(source)
                self.readers.append(reader)
                source_metadata.append(reader.index.metadata)
                for record in reader.iterate_records():
                    records.append(
                        replace(
                            record,
                            metadata=dict(record.metadata)
                            | {"source_dataset": str(source.resolve())},
                        )
                    )
            metadata = {
                "dataset_name": str(payload.get("dataset_name", self.dataset_dir.name)),
                "format": "dovla_cil_collection",
                "sources": [str(source) for source in sources],
                "source_metadata": source_metadata,
                "num_records": len(records),
                "num_groups": len({record.group_id for record in records}),
                "task_count": len({record.task_id for record in records}),
            }
            self.reader = None
            self.index = CILCollectionIndex(
                dataset_dir=self.dataset_dir,
                metadata=metadata,
                record_count=len(records),
                group_count=int(metadata["num_groups"]),
            )
            self.records = records
        else:
            self.reader = ShardReader(self.dataset_dir)
            self.readers = [self.reader]
            self.index = self.reader.index
            self.records = list(self.reader.iterate_records())
        if len({record.record_id for record in self.records}) != len(self.records):
            raise ValueError("CIL dataset contains duplicate record_id values")
        if collection_path.exists():
            group_sources: dict[str, set[str]] = {}
            for record in self.records:
                group_sources.setdefault(record.group_id, set()).add(
                    str(record.metadata["source_dataset"])
                )
            collisions = [
                group_id for group_id, sources in group_sources.items() if len(sources) > 1
            ]
            if collisions:
                raise ValueError(
                    f"CIL collection contains group_id collisions: {collisions[:3]}"
                )
        self._record_id_to_index = {
            record.record_id: index for index, record in enumerate(self.records)
        }
        self._group_to_indices: dict[str, list[int]] = {}
        for index, record in enumerate(self.records):
            self._group_to_indices.setdefault(record.group_id, []).append(index)
        self.group_ids: list[str] = list(self._group_to_indices)

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> CILRecord:
        return self.records[index]

    def __iter__(self) -> Iterator[CILRecord]:
        return iter(self.records)

    def get_group(self, group_id: str) -> list[CILRecord]:
        if group_id in self._group_to_indices:
            return [self.records[index] for index in self._group_to_indices[group_id]]
        raise KeyError(f"Unknown CIL group: {group_id}")

    def group_indices(self, group_id: str) -> list[int]:
        if group_id not in self._group_to_indices:
            raise KeyError(f"Unknown CIL group: {group_id}")
        return list(self._group_to_indices[group_id])

    def iter_groups(self) -> Iterator[list[CILRecord]]:
        for group_id in self.group_ids:
            yield self.get_group(group_id)

    def record_index(self, record_id: str) -> int:
        return self._record_id_to_index[record_id]


class CILJsonlDataset(CILDataset):
    """Backward-compatible name from the initial scaffold."""

    def __init__(self, manifest_path: str | Path) -> None:
        super().__init__(manifest_path)

    def groups(self) -> Iterable[list[CILRecord]]:
        yield from self.iter_groups()


def iter_dataset_records(path: str | Path) -> Iterator[CILRecord]:
    target = Path(path)
    if target.is_dir() and (target / "collection.json").exists():
        yield from CILDataset(target)
        return
    index = ShardIndex.from_path(path)
    for shard_path in index.shards:
        yield from iter_cil_records(shard_path)


def groups_from_records(records: Iterable[CILRecord]) -> Iterable[list[CILRecord]]:
    yield from group_records(records).values()


def write_cil_collection(
    output_dir: str | Path,
    sources: Iterable[str | Path],
    *,
    dataset_name: str = "dovla_cil_collection",
) -> Path:
    output = Path(output_dir)
    output.mkdir(parents=True, exist_ok=True)
    resolved = [str(Path(source).resolve()) for source in sources]
    if not resolved:
        raise ValueError("sources must be non-empty")
    path = output / "collection.json"
    write_json(
        {
            "version": 1,
            "dataset_name": dataset_name,
            "sources": resolved,
        },
        path,
    )
    return path


def _collection_sources(payload: Any, *, base_dir: Path) -> list[Path]:
    if not isinstance(payload, dict) or int(payload.get("version", 0)) != 1:
        raise ValueError("CIL collection must be a version-1 object")
    values = payload.get("sources", [])
    if not isinstance(values, list):
        raise ValueError("CIL collection sources must be a list")
    sources: list[Path] = []
    for entry in values:
        raw = entry.get("path") if isinstance(entry, dict) else entry
        if not isinstance(raw, str) or not raw:
            raise ValueError("each CIL collection source must provide a path")
        source = Path(raw)
        sources.append(source if source.is_absolute() else base_dir / source)
    return sources