source/src/strata/data/registry.py

"""Typed dataset registry for STRATA acquisition.

The registry is the single, version-controlled source of truth describing which
datasets STRATA uses, where they legitimately come from, their licensing tier,
and which languages they cover. It is intentionally declarative: acquisition
code in :mod:`strata.data.acquire` interprets these specs, but no download URLs
or credentials live in Python source.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from pathlib import Path

import yaml

# ---------------------------------------------------------------------------
# Controlled vocabularies
# ---------------------------------------------------------------------------

# Licensing / accessibility tiers. These drive default fetch behaviour.
TIER_OPEN = "open"  # freely downloadable, no gate
TIER_GATED_AUTO = "gated_auto"  # HF gate with automatic approval on request
TIER_GATED_MANUAL = "gated_manual"  # HF gate requiring manual approval
TIER_LICENSE_REQUIRED = "license_required"  # e.g. LDC / signed agreement; no auto DL

LICENSE_TIERS = frozenset(
    {TIER_OPEN, TIER_GATED_AUTO, TIER_GATED_MANUAL, TIER_LICENSE_REQUIRED}
)

# Source kinds understood by the acquirer.
KIND_HF_SNAPSHOT = "hf_snapshot"  # snapshot_download of a HF dataset repo
KIND_HF_FILES = "hf_files"  # capped/prefix-filtered file selection from a HF repo
KIND_HTTP = "http"  # direct file/tarball download(s)
KIND_GITHUB_ARCHIVE = "github_archive"  # pinned GitHub repo tarball
KIND_LICENSE_REQUIRED = "license_required"  # no download; write acquisition stub

SOURCE_KINDS = frozenset(
    {
        KIND_HF_SNAPSHOT,
        KIND_HF_FILES,
        KIND_HTTP,
        KIND_GITHUB_ARCHIVE,
        KIND_LICENSE_REQUIRED,
    }
)

# Pipeline layers, matching the STRATA data stack. "multi" is used in a spec's
# languages list to mean "many languages, filtered downstream".
LANG_MULTI = "multi"


# ---------------------------------------------------------------------------
# Spec
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class DatasetSpec:
    """Declarative description of one dataset."""

    id: str
    title: str
    layer: str
    languages: tuple[str, ...]
    license: str
    license_tier: str
    source: dict[str, object]
    homepage: str = ""
    notes: str = ""
    optional: bool = False

    def __post_init__(self) -> None:
        self.validate()

    @property
    def kind(self) -> str:
        return str(self.source.get("kind", ""))

    def covers_language(self, language: str) -> bool:
        return language in self.languages or LANG_MULTI in self.languages

    def validate(self) -> None:
        if not self.id or not self.id.replace("_", "").replace("-", "").isalnum():
            raise ValueError(
                f"dataset id must be non-empty alphanumeric/_/-, got {self.id!r}"
            )
        if self.license_tier not in LICENSE_TIERS:
            raise ValueError(
                f"[{self.id}] invalid license_tier {self.license_tier!r}; "
                f"expected one of {sorted(LICENSE_TIERS)}"
            )
        if not self.languages:
            raise ValueError(f"[{self.id}] languages must be non-empty")
        kind = self.kind
        if kind not in SOURCE_KINDS:
            raise ValueError(
                f"[{self.id}] invalid source kind {kind!r}; "
                f"expected one of {sorted(SOURCE_KINDS)}"
            )
        _validate_source(self.id, kind, self.source)
        if kind == KIND_LICENSE_REQUIRED and self.license_tier != TIER_LICENSE_REQUIRED:
            raise ValueError(
                f"[{self.id}] source kind {KIND_LICENSE_REQUIRED} requires "
                f"license_tier {TIER_LICENSE_REQUIRED}"
            )


def _require(dataset_id: str, source: dict[str, object], *keys: str) -> None:
    for key in keys:
        if not source.get(key):
            raise ValueError(
                f"[{dataset_id}] source kind {source.get('kind')!r} requires "
                f"non-empty {key!r}"
            )


def _validate_source(dataset_id: str, kind: str, source: dict[str, object]) -> None:
    if kind == KIND_HF_SNAPSHOT:
        _require(dataset_id, source, "repo_id")
    elif kind == KIND_HF_FILES:
        _require(dataset_id, source, "repo_id", "prefixes")
        if not isinstance(source["prefixes"], list):
            raise ValueError(f"[{dataset_id}] hf_files.prefixes must be a list")
    elif kind == KIND_HTTP:
        urls = source.get("urls")
        if not isinstance(urls, list) or not urls:
            raise ValueError(f"[{dataset_id}] http.urls must be a non-empty list")
        for item in urls:
            if not isinstance(item, dict) or not item.get("url"):
                raise ValueError(f"[{dataset_id}] each http url needs a 'url' field")
    elif kind == KIND_GITHUB_ARCHIVE:
        _require(dataset_id, source, "repo")
    elif kind == KIND_LICENSE_REQUIRED:
        _require(dataset_id, source, "acquisition")


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------


@dataclass(frozen=True, slots=True)
class Registry:
    version: int
    languages_core: tuple[str, ...]
    datasets: tuple[DatasetSpec, ...]
    description: str = ""

    def by_id(self, dataset_id: str) -> DatasetSpec:
        for spec in self.datasets:
            if spec.id == dataset_id:
                return spec
        raise KeyError(f"unknown dataset id: {dataset_id!r}")

    def ids(self) -> list[str]:
        return [spec.id for spec in self.datasets]

    def layers(self) -> list[str]:
        seen: dict[str, None] = {}
        for spec in self.datasets:
            seen.setdefault(spec.layer, None)
        return list(seen)

    def filter(
        self,
        *,
        ids: list[str] | None = None,
        layers: list[str] | None = None,
        languages: list[str] | None = None,
        tiers: list[str] | None = None,
        include_optional: bool = True,
    ) -> list[DatasetSpec]:
        """Return specs matching all provided filters (AND semantics).

        ``ids`` matching is explicit and always includes optional datasets, so
        an operator can request an optional dataset by name.
        """

        id_set = set(ids) if ids else None
        layer_set = set(layers) if layers else None
        tier_set = set(tiers) if tiers else None

        selected: list[DatasetSpec] = []
        for spec in self.datasets:
            if id_set is not None and spec.id not in id_set:
                continue
            if layer_set is not None and spec.layer not in layer_set:
                continue
            if tier_set is not None and spec.license_tier not in tier_set:
                continue
            if languages and not any(spec.covers_language(lang) for lang in languages):
                continue
            explicitly_named = id_set is not None and spec.id in id_set
            if spec.optional and not include_optional and not explicitly_named:
                continue
            selected.append(spec)
        return selected


def load_registry(path: Path) -> Registry:
    """Load and validate a registry YAML file."""

    if not path.exists():
        raise FileNotFoundError(f"registry file not found: {path}")
    data = yaml.safe_load(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"registry {path} must be a mapping at the top level")

    raw_datasets = data.get("datasets")
    if not isinstance(raw_datasets, list) or not raw_datasets:
        raise ValueError(f"registry {path} must contain a non-empty 'datasets' list")

    specs: list[DatasetSpec] = []
    seen_ids: set[str] = set()
    for raw in raw_datasets:
        if not isinstance(raw, dict):
            raise ValueError(f"registry {path}: each dataset entry must be a mapping")
        spec = _spec_from_dict(raw)
        if spec.id in seen_ids:
            raise ValueError(f"registry {path}: duplicate dataset id {spec.id!r}")
        seen_ids.add(spec.id)
        specs.append(spec)

    return Registry(
        version=int(data.get("version", 1)),
        languages_core=tuple(data.get("languages_core", []) or []),
        datasets=tuple(specs),
        description=str(data.get("description", "")),
    )


def _spec_from_dict(raw: dict[str, object]) -> DatasetSpec:
    missing = [key for key in ("id", "layer", "license_tier", "source") if key not in raw]
    if missing:
        raise ValueError(f"dataset entry missing required keys: {missing} in {raw!r}")
    source = raw["source"]
    if not isinstance(source, dict):
        raise ValueError(f"[{raw.get('id')}] source must be a mapping")
    languages = raw.get("languages", [])
    if not isinstance(languages, list):
        raise ValueError(f"[{raw.get('id')}] languages must be a list")
    return DatasetSpec(
        id=str(raw["id"]),
        title=str(raw.get("title", raw["id"])),
        layer=str(raw["layer"]),
        languages=tuple(str(lang) for lang in languages),
        license=str(raw.get("license", "unknown")),
        license_tier=str(raw["license_tier"]),
        source=dict(source),
        homepage=str(raw.get("homepage", "")),
        notes=str(raw.get("notes", "")),
        optional=bool(raw.get("optional", False)),
    )


# Backwards-friendly export list.
__all__ = [
    "DatasetSpec",
    "Registry",
    "load_registry",
    "LICENSE_TIERS",
    "SOURCE_KINDS",
    "TIER_OPEN",
    "TIER_GATED_AUTO",
    "TIER_GATED_MANUAL",
    "TIER_LICENSE_REQUIRED",
    "KIND_HF_SNAPSHOT",
    "KIND_HF_FILES",
    "KIND_HTTP",
    "KIND_GITHUB_ARCHIVE",
    "KIND_LICENSE_REQUIRED",
    "LANG_MULTI",
]