"""Explicit, verified dataset selection shared by training and evaluation."""

from __future__ import annotations

import hashlib
import re
from pathlib import Path

import h5py
import numpy as np
from omegaconf import DictConfig, OmegaConf

from softasp.sudoku_protocol import data_protocol_metadata


def validate_sudoku_hdf5(path: Path, expected_sha256: str) -> dict:
    """Validate identity, complete splits, model-compatible fields and value ranges."""
    path = path.expanduser().resolve()
    if not path.is_file():
        raise FileNotFoundError(f"Explicit Sudoku dataset does not exist: {path}")
    if not isinstance(expected_sha256, str) or not re.fullmatch(
        r"[0-9a-fA-F]{64}", expected_sha256
    ):
        raise ValueError("data_artifact.sha256 must be a complete SHA256 digest")
    with path.open("rb") as stream:
        actual = hashlib.file_digest(stream, "sha256").hexdigest()
    if actual != expected_sha256.lower():
        raise ValueError(f"Sudoku dataset SHA256 mismatch: {path}")
    counts = {}
    with h5py.File(path, "r") as hf:
        for split in ("train", "val", "test"):
            group = hf.get(split)
            if not isinstance(group, h5py.Group):
                raise ValueError(f"Missing Sudoku split: {split}")
            fields: dict[str, h5py.Dataset] = {}
            for field in ("images", "labels", "puzzle_mask"):
                value = group.get(field)
                if not isinstance(value, h5py.Dataset):
                    raise ValueError(f"Missing Sudoku field: {split}/{field}")
                fields[field] = value
            labels, mask, images = fields["labels"], fields["puzzle_mask"], fields["images"]
            n = labels.shape[0] if labels.ndim else 0
            if n == 0 or labels.shape != (n, 81) or mask.shape != (n, 81):
                raise ValueError(f"Invalid Sudoku label/mask shapes: {split}")
            if images.shape != (n, 252, 252, 1):
                raise ValueError(f"Invalid Sudoku image shape: {split}")
            if images.dtype != np.dtype("uint8") or labels.dtype != np.dtype("int64"):
                raise ValueError(f"Invalid Sudoku image/label dtype: {split}")
            if mask.dtype.kind not in "biu":
                raise ValueError(f"Invalid Sudoku mask dtype: {split}")
            for offset in range(0, n, 512):
                values, clues = labels[offset : offset + 512], mask[offset : offset + 512]
                if np.any((values < 0) | (values > 8)):
                    raise ValueError(f"Sudoku labels outside 0..8: {split}")
                if np.any((clues != 0) & (clues != 1)):
                    raise ValueError(f"Sudoku mask must contain only 0/1: {split}")
            counts[split] = n
    return {
        "path": str(path),
        "sha256": actual,
        "splits": counts,
        "validation": "sudoku_hdf5_v1",
        "data_protocol": data_protocol_metadata(actual),
    }


def select_hdf5(cfg: DictConfig) -> Path:
    """Release 1 requires a pinned artifact; legacy configs retain discovery."""
    spec = cfg.get("data_artifact")
    if spec is not None or cfg.get("release_profile") == "sudoku_release1":
        if not spec or not spec.get("path") or not spec.get("sha256"):
            raise ValueError("Set data_artifact.path and data_artifact.sha256 explicitly")
        if cfg.dataset.name != "sudoku":
            raise ValueError("Explicit data_artifact selection currently supports Sudoku only")
        identity = validate_sudoku_hdf5(Path(spec.path), spec.sha256)
        identity["data_protocol"] = data_protocol_metadata(
            identity["sha256"], spec.get("protocol_id")
        )
        OmegaConf.update(cfg, "resolved_dataset", identity, force_add=True)
        return Path(identity["path"])
    directory = Path(cfg.paths.data_root) / "processed" / cfg.dataset.name
    files = sorted(directory.glob("v_*.h5"))
    if not files:
        raise FileNotFoundError(f"No preprocessed HDF5 found in {directory}")
    return files[-1]