src/sr_harness/core/context_data.py

"""Manifest-backed NumPy data stored in a workspace directory."""

from __future__ import annotations

import json
import os
from pathlib import Path
import tempfile
from typing import TYPE_CHECKING, Any

import numpy as np

if TYPE_CHECKING:
    from ..runtime.workspace import Workspace


class ContextManifestError(ValueError):
    """Raised when a context-data manifest cannot be validated."""


class _ContextDataReader:
    """Private implementation shared by the public loading functions."""

    MANIFEST_NAME = "manifest.json"
    REQUIRED_ROOT_FIELDS = {"variables", "axes"}
    OPTIONAL_ROOT_FIELDS = {"num_nodes"}
    REQUIRED_VARIABLE_FIELDS = {"file", "description", "axes"}
    OPTIONAL_VARIABLE_FIELDS = {"kind", "structure"}
    AXIS_SOURCE_FIELDS = {"values", "file", "size"}
    AXIS_FIELDS = {"description", *AXIS_SOURCE_FIELDS}

    def __init__(self, directory: str | Path, *, workspace: Workspace | None = None):
        self.directory = Path(directory).resolve()
        self.workspace = workspace

    def inspect(self) -> dict[str, Any]:
        """Validate the store and return diagnostics without raising.

        Returns:
            A serializable report containing errors, warnings, and array summaries.
        """
        errors: list[str] = []
        warnings: list[str] = []
        manifest = self._read_manifest(errors)
        loaded = (
            self._validate_manifest(manifest, errors, warnings)
            if manifest is not None
            else None
        )
        report: dict[str, Any] = {
            "valid": not errors,
            "path": str(self.directory),
            "errors": errors,
            "warnings": warnings,
            "num_nodes": loaded["num_nodes"] if loaded is not None else None,
            "variables": {},
            "axes": {},
        }
        if loaded is not None:
            report["variables"] = {
                name: {
                    "file": manifest["variables"][name]["file"],
                    "shape": list(value.shape),
                    "dtype": str(value.dtype),
                    "axes": list(loaded["variable_axes"][name]),
                    **(
                        {"structure": loaded["variable_structures"][name]}
                        if name in loaded["variable_structures"] else {}
                    ),
                    **({"kind": "relation"} if name in loaded["relation_names"] else {}),
                    "description": loaded["variable_descriptions"][name],
                }
                for name, value in loaded["data"].items()
                if name in manifest["variables"]
            }
            report["axes"] = {
                name: {
                    "size": len(loaded["data"][name]),
                    "dtype": str(loaded["data"][name].dtype),
                    "storage": axis["storage"],
                    "description": loaded["variable_descriptions"][name],
                }
                for name, axis in loaded["axis_metadata"].items()
            }
        return report

    def load(self) -> dict[str, Any]:
        """Validate and load all variables and axes.

        Returns:
            An AgentContext-ready mapping containing arrays and metadata.

        Raises:
            ContextManifestError: If the manifest or referenced arrays are invalid.
        """
        errors: list[str] = []
        warnings: list[str] = []
        manifest = self._read_manifest(errors)
        loaded = (
            self._validate_manifest(manifest, errors, warnings)
            if manifest is not None else None
        )
        if errors or loaded is None:
            raise ContextManifestError("Invalid context.data:\n- " + "\n- ".join(errors))
        return loaded

    def _read_manifest(self, errors: list[str]) -> dict[str, Any] | None:
        manifest_path = self.directory / self.MANIFEST_NAME
        if not self.directory.is_dir():
            errors.append(f"directory does not exist: {self.directory}")
            return None
        if not manifest_path.is_file():
            errors.append(f"missing {self.MANIFEST_NAME}")
            return None

        def reject_duplicates(pairs):
            result = {}
            for key, value in pairs:
                if key in result:
                    raise ContextManifestError(f"duplicate JSON key: {key}")
                result[key] = value
            return result

        try:
            opener = self.workspace.safe_open if self.workspace is not None else open
            with opener(manifest_path, "r", encoding="utf-8") as f:
                manifest = json.load(f, object_pairs_hook=reject_duplicates)
        except (ContextManifestError, json.JSONDecodeError, OSError) as exc:
            errors.append(f"cannot read {self.MANIFEST_NAME}: {exc}")
            return None
        if not isinstance(manifest, dict):
            errors.append("manifest root must be a JSON object")
            return None
        return manifest

    def _validate_manifest(
        self,
        manifest: dict[str, Any],
        errors: list[str],
        warnings: list[str],
    ) -> dict[str, Any] | None:
        self._check_fields(
            "manifest",
            manifest,
            self.REQUIRED_ROOT_FIELDS,
            errors,
            optional=self.OPTIONAL_ROOT_FIELDS,
        )
        variables = manifest.get("variables")
        axes = manifest.get("axes")
        num_nodes = manifest.get("num_nodes")
        if "num_nodes" in manifest and (
            isinstance(num_nodes, bool) or not isinstance(num_nodes, int) or num_nodes < 1
        ):
            errors.append("manifest.num_nodes must be a positive integer")
            num_nodes = None
        if not isinstance(variables, dict) or not variables:
            errors.append("manifest.variables must be a non-empty object")
            variables = {}
        if not isinstance(axes, dict):
            errors.append("manifest.axes must be an object")
            axes = {}

        loaded_axes: dict[str, dict[str, Any]] = {}
        referenced_files = {self.MANIFEST_NAME}
        for name, spec in axes.items():
            location = f"axes.{name}"
            if not self._valid_name(name):
                errors.append(f"{location}: name must be non-empty and contain no path separators")
                continue
            if not isinstance(spec, dict):
                errors.append(f"{location} must be an object")
                continue
            if "description" not in spec:
                errors.append(f"{location} is missing fields: ['description']")
            if extra := sorted(set(spec) - self.AXIS_FIELDS):
                errors.append(f"{location} contains unsupported fields: {extra}")
            description = self._description(location, spec, errors)
            sources = self.AXIS_SOURCE_FIELDS & set(spec)
            if len(sources) != 1:
                errors.append(f"{location} must contain exactly one of values, file, or size")
                continue
            source = next(iter(sources))
            try:
                if source == "values":
                    if not isinstance(spec["values"], list) or not spec["values"]:
                        raise ValueError("values must be a non-empty JSON array")
                    if any(isinstance(value, (list, dict)) for value in spec["values"]):
                        raise ValueError("values must contain scalar JSON values")
                    values = np.asarray(spec["values"])
                elif source == "size":
                    size = spec["size"]
                    if isinstance(size, bool) or not isinstance(size, int) or size < 1:
                        raise ValueError("size must be a positive integer")
                    values = np.arange(size)
                else:
                    expected = f"{name}.npy"
                    if spec["file"] != expected:
                        raise ValueError(f"file must be {expected!r}")
                    referenced_files.add(expected)
                    values = self._load_array(expected)
                    if values.ndim != 1:
                        raise ValueError(f"axis array must be one-dimensional, got shape {values.shape}")
                if values.ndim != 1:
                    raise ValueError("values must be one-dimensional")
                loaded_axes[name] = {"values": values, "description": description, "storage": source}
            except (OSError, ValueError) as exc:
                errors.append(f"{location}: {exc}")

        loaded_values: dict[str, np.ndarray] = {}
        variable_axes: dict[str, tuple[str, ...]] = {}
        variable_structures: dict[str, str] = {}
        relation_names: set[str] = set()
        descriptions: dict[str, str] = {}
        used_axes: set[str] = set()
        for name, spec in variables.items():
            location = f"variables.{name}"
            if not self._valid_name(name):
                errors.append(f"{location}: name must be non-empty and contain no path separators")
                continue
            if not isinstance(spec, dict):
                errors.append(f"{location} must be an object")
                continue
            self._check_fields(
                location,
                spec,
                self.REQUIRED_VARIABLE_FIELDS,
                errors,
                optional=self.OPTIONAL_VARIABLE_FIELDS,
            )
            description = self._description(location, spec, errors)
            structure = spec.get("structure")
            if structure is not None and not self._valid_name(structure):
                errors.append(f"{location}.structure must name an A or T variable")
                structure = None
            kind = spec.get("kind")
            if kind not in {None, "relation"}:
                errors.append(f"{location}.kind must be 'relation' when provided")
                kind = None
            expected = f"{name}.npy"
            if spec.get("file") != expected:
                errors.append(f"{location}.file must be {expected!r}")
                continue
            declared_axes = spec.get("axes")
            if not isinstance(declared_axes, list) or any(
                not isinstance(axis, str) or not axis for axis in declared_axes
            ):
                errors.append(f"{location}.axes must be an array of non-empty axis names")
                continue
            missing_axes = [axis for axis in declared_axes if axis not in axes]
            if missing_axes:
                errors.append(f"{location}.axes references missing axes: {missing_axes}")
                continue
            referenced_files.add(expected)
            try:
                value = self._load_array(expected)
            except (OSError, ValueError) as exc:
                errors.append(f"{location}: {exc}")
                continue
            if value.ndim != len(declared_axes):
                errors.append(
                    f"{location}: array has {value.ndim} dimensions but declares "
                    f"{len(declared_axes)} axes"
                )
                continue
            mismatches = [
                f"{axis}={len(loaded_axes[axis]['values'])} (data={value.shape[index]})"
                for index, axis in enumerate(declared_axes)
                if axis in loaded_axes and len(loaded_axes[axis]["values"]) != value.shape[index]
            ]
            if mismatches:
                errors.append(f"{location}: axis length mismatch: {', '.join(mismatches)}")
                continue
            loaded_values[name] = value
            variable_axes[name] = tuple(declared_axes)
            if structure is not None:
                variable_structures[name] = structure
            if kind == "relation":
                relation_names.add(name)
            descriptions[name] = description
            used_axes.update(declared_axes)

        unused_axes = sorted(set(axes) - used_axes)
        if unused_axes:
            errors.append(f"manifest.axes contains unreferenced axes: {unused_axes}")

        if bool(relation_names) != (num_nodes is not None):
            errors.append("manifest.num_nodes and relation metadata must appear together")
        for relation_name in relation_names:
            location = f"variables.{relation_name}"
            relation = loaded_values.get(relation_name)
            if relation is None:
                continue
            if relation.ndim != 2 or relation.shape[1] not in {2, 3}:
                errors.append(
                    f"{location}: relation must have shape (E, 2) or (H, 3), got {relation.shape}"
                )
                continue
            if relation.dtype.kind not in "iu":
                errors.append(f"{location}: relation endpoints must use an integer dtype")
            elif num_nodes is not None and (np.any(relation < 0) or np.any(relation >= num_nodes)):
                errors.append(f"{location}: relation endpoints must be in [0, {num_nodes})")
        for variable, structure in variable_structures.items():
            location = f"variables.{variable}.structure"
            if variable == structure:
                errors.append(f"{location}: a structured variable cannot reference itself")
                continue
            if structure not in loaded_values:
                errors.append(f"{location} references missing variable {structure!r}")
                continue
            if structure not in relation_names:
                errors.append(f"{location} must reference a variable with kind 'relation'")
                continue
            relation = loaded_values[structure]
            if relation.ndim != 2 or relation.shape[1] not in {2, 3}:
                errors.append(
                    f"{location}: {structure!r} must have shape (E, 2) or (H, 3), "
                    f"got {relation.shape}"
                )
                continue
            dependent = loaded_values.get(variable)
            if dependent is not None and (
                dependent.ndim == 0 or dependent.shape[-1] != relation.shape[0]
            ):
                errors.append(
                    f"{location}: {variable!r} must end in a dimension of "
                    f"length {relation.shape[0]}"
                )
            if relation.dtype.kind not in "iu":
                errors.append(f"{location}: relation endpoints must use an integer dtype")
            elif num_nodes is not None and (np.any(relation < 0) or np.any(relation >= num_nodes)):
                errors.append(f"{location}: relation endpoints must be in [0, {num_nodes})")
        actual_npy = {path.name for path in self.directory.glob("*.npy") if path.is_file()}
        untracked = sorted(actual_npy - referenced_files)
        if untracked:
            warnings.append(f"unreferenced NPY files: {untracked}")
        if errors:
            return None
        data = {**{name: axis["values"] for name, axis in loaded_axes.items()}, **loaded_values}
        all_descriptions = {name: axis["description"] for name, axis in loaded_axes.items()} | descriptions
        return {
            "data": data,
            "variable_descriptions": all_descriptions,
            "variable_axes": variable_axes,
            "variable_structures": variable_structures,
            "relation_names": relation_names,
            "num_nodes": num_nodes,
            "axis_metadata": loaded_axes,
        }

    def _load_array(self, filename: str) -> np.ndarray:
        path = self.directory / filename
        if not path.is_file() or path.parent != self.directory:
            raise ValueError(f"missing array file: {filename}")
        opener = self.workspace.safe_open if self.workspace is not None else open
        with opener(path, "rb") as f:
            value = np.load(f, allow_pickle=False)
        if not isinstance(value, np.ndarray):
            raise ValueError(f"{filename} does not contain one NPY array")
        if value.dtype.hasobject:
            raise ValueError(f"{filename} uses object dtype, which is not allowed")
        return value

    @staticmethod
    def _valid_name(name: Any) -> bool:
        return (
            isinstance(name, str)
            and bool(name.strip())
            and name not in {".", ".."}
            and "/" not in name
            and "\\" not in name
            and "\0" not in name
        )

    @staticmethod
    def _description(location: str, spec: dict[str, Any], errors: list[str]) -> str:
        description = spec.get("description")
        if not isinstance(description, str):
            errors.append(f"{location}.description must be a string")
            return ""
        return description.strip()

    @staticmethod
    def _check_fields(
        location: str,
        value: dict[str, Any],
        expected: set[str],
        errors: list[str],
        *,
        optional: set[str] | None = None,
    ) -> None:
        optional = optional or set()
        missing = sorted(expected - set(value))
        extra = sorted(set(value) - expected - optional)
        if missing:
            errors.append(f"{location} is missing fields: {missing}")
        if extra:
            errors.append(f"{location} contains unsupported fields: {extra}")


def load_context_data(directory: str | Path, *, workspace: Workspace | None = None) -> dict[str, Any]:
    """Validate and load one manifest-backed ``context.data`` directory.

    Args:
        directory: Directory containing ``manifest.json`` and array files.
        workspace: Workspace enforcing safe reads, when available.

    Returns:
        An AgentContext-ready mapping containing arrays and metadata.
    """
    return _ContextDataReader(directory, workspace=workspace).load()


def inspect_context_data(directory: str | Path, *, workspace: Workspace | None = None) -> dict[str, Any]:
    """Return validation diagnostics without raising for manifest errors.

    Args:
        directory: Directory containing ``manifest.json`` and array files.
        workspace: Workspace enforcing safe reads, when available.

    Returns:
        A serializable validation report.
    """
    return _ContextDataReader(directory, workspace=workspace).inspect()


def update_context_data_descriptions(
    directory: str | Path,
    descriptions: dict[str, str],
    *,
    workspace: Workspace | None = None,
) -> dict[str, Any]:
    """Update variable and axis descriptions in ``manifest.json``.

    The complete store is validated before the edit. An atomic replacement is
    used when the directory permits creating entries. If the directory is
    read-only but ``manifest.json`` itself is writable, the existing file is
    updated in place, matching normal POSIX file-permission semantics. Because
    only existing string-valued description fields are changed, structural
    validity is preserved. The returned mapping matches :func:`load_context_data`.
    """
    directory = Path(directory).resolve()
    manifest_path = directory / _ContextDataReader.MANIFEST_NAME
    loaded = load_context_data(directory, workspace=workspace)
    opener = workspace.safe_open if workspace is not None else open
    if not isinstance(descriptions, dict) or any(
        not isinstance(name, str) or not isinstance(value, str)
        for name, value in descriptions.items()
    ):
        raise ValueError("descriptions must map variable names to text")
    unknown = set(descriptions) - set(loaded["data"])
    if unknown:
        raise ValueError(f"unknown context.data variables: {sorted(unknown)}")
    with opener(manifest_path, "r", encoding="utf-8") as f:
        manifest = json.load(f)
    for section in ("variables", "axes"):
        for name, spec in manifest[section].items():
            if name in descriptions:
                spec["description"] = descriptions[name].strip()
    serialized = json.dumps(manifest, ensure_ascii=False, indent=2) + "\n"
    if workspace is not None:
        relative_manifest = str(manifest_path.relative_to(workspace.path))
        temporary_path = None
        try:
            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                dir=workspace.path,
                prefix=".manifest-",
                suffix=".json.tmp",
                delete=False,
            ) as stream:
                temporary_path = Path(stream.name)
                stream.write(serialized)
                stream.flush()
                os.fsync(stream.fileno())
            workspace.replace_file(temporary_path, relative_manifest)
            temporary_path = None
        finally:
            if temporary_path is not None:
                temporary_path.unlink(missing_ok=True)
        updated = dict(loaded)
        updated["variable_descriptions"] = loaded["variable_descriptions"] | {
            name: value.strip() for name, value in descriptions.items()
        }
        return updated
    temporary_path = None
    try:
        try:
            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                dir=directory,
                prefix=".manifest-",
                suffix=".json.tmp",
                delete=False,
            ) as stream:
                temporary_path = Path(stream.name)
                stream.write(serialized)
                stream.flush()
                os.fsync(stream.fileno())
            os.replace(temporary_path, manifest_path)
            temporary_path = None
        except PermissionError:
            if temporary_path is not None:
                temporary_path.unlink(missing_ok=True)
            with manifest_path.open("w", encoding="utf-8") as stream:
                stream.write(serialized)
                stream.flush()
                os.fsync(stream.fileno())
    finally:
        if temporary_path is not None:
            temporary_path.unlink(missing_ok=True)
    updated = dict(loaded)
    updated["variable_descriptions"] = loaded["variable_descriptions"] | {
        name: value.strip() for name, value in descriptions.items()
    }
    return updated