"""Atomic artifact I/O and fail-closed OID v2 schemas."""
from __future__ import annotations

from pathlib import Path
from typing import Any, Mapping

import numpy as np

from CCD_analysis.utils.run_contract import (
    array_sha256,
    atomic_publish_json,
    atomic_publish_npz,
    sha256_file,
)


class ArtifactError(ValueError):
    pass


def atomic_write_json(path: str | Path, value: Mapping[str, Any]) -> Path:
    """Publish canonical JSON without clobbering, via the shared contract."""
    destination = Path(path)
    atomic_publish_json(destination, value)
    return destination


def atomic_save_npz(path: str | Path, *, compressed: bool = True, **arrays: np.ndarray) -> Path:
    """Publish a validated, hashable NPZ without clobbering."""
    destination = Path(path)
    atomic_publish_npz(destination, arrays, compressed=compressed)
    return destination


def validate_timeline(data: Mapping[str, np.ndarray], *, expected_control_interval: int | None = None) -> dict[str, Any]:
    required = {"lattice_steps", "control_indices", "control_offsets", "control_interval"}
    missing = required.difference(data)
    if missing:
        raise ArtifactError(f"missing timeline keys: {sorted(missing)}")
    steps = np.asarray(data["lattice_steps"])
    indices = np.asarray(data["control_indices"])
    offsets = np.asarray(data["control_offsets"])
    interval_array = np.asarray(data["control_interval"])
    if interval_array.size != 1:
        raise ArtifactError("control_interval must be scalar")
    interval = int(interval_array.reshape(-1)[0])
    if interval < 1 or (expected_control_interval is not None and interval != expected_control_interval):
        raise ArtifactError("invalid control_interval")
    if steps.ndim != 1 or indices.shape != steps.shape or offsets.shape != steps.shape:
        raise ArtifactError("timeline arrays must be aligned one-dimensional arrays")
    if len(steps) == 0 or np.any(np.diff(steps) <= 0):
        raise ArtifactError("lattice_steps must be non-empty, unique, and strictly increasing")
    expected_indices = (steps - 1) // interval
    expected_offsets = (steps - 1) % interval + 1
    if not np.array_equal(indices, expected_indices) or not np.array_equal(offsets, expected_offsets):
        raise ArtifactError("control index/offset does not match lattice_steps")
    return {"count": len(steps), "first": int(steps[0]), "last": int(steps[-1]), "control_interval": interval}


def validate_field_artifact(data: Mapping[str, np.ndarray], *, expected_control_interval: int | None = None) -> dict[str, Any]:
    if "ux" not in data or "uy" not in data:
        raise ArtifactError("field artifact requires ux and uy")
    ux, uy = np.asarray(data["ux"]), np.asarray(data["uy"])
    if ux.shape != uy.shape or ux.ndim != 3:
        raise ArtifactError("ux and uy must have identical (N, ny, nx) shapes")
    if not np.issubdtype(ux.dtype, np.floating) or not np.issubdtype(uy.dtype, np.floating):
        raise ArtifactError("fields must be floating point")
    if not np.all(np.isfinite(ux)) or not np.all(np.isfinite(uy)):
        raise ArtifactError("fields contain non-finite values")
    timeline = validate_timeline(data, expected_control_interval=expected_control_interval)
    if ux.shape[0] != timeline["count"]:
        raise ArtifactError("field count differs from timeline count")
    return {"shape": list(ux.shape), "dtype": str(ux.dtype), "timeline": timeline}


def artifact_record(path: str | Path, validation: Mapping[str, Any], *, parent_ids: list[str] | None = None) -> dict[str, Any]:
    item = Path(path)
    arrays: dict[str, str] = {}
    if item.suffix == ".npz":
        with np.load(item, allow_pickle=False) as archive:
            arrays = {key: array_sha256(archive[key]) for key in archive.files}
    return {"path": str(item), "size_bytes": item.stat().st_size, "sha256": sha256_file(item),
            "array_hashes": arrays, "validation": dict(validation), "parent_artifact_ids": parent_ids or []}
