tests/test_contracts.py

"""
Test scenarios required before this codebase's core contracts could be
considered complete (per project guidelines: nothing is done without
test scenarios for duplicate requests, edge cases, and race conditions).

Run:  python -m pytest tests/test_contracts.py -v
"""
from __future__ import annotations
import os
import sys
import shutil
import tempfile
from pathlib import Path

import numpy as np
import pytest
import torch

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.data_real import LocalWellHDF5, get_dataset, get_synthetic_dataset
from src.provenance import (
    DataLoadError,
    SchemaValidationError,
    TrajectoryTooShortError,
    EmptyDatasetError,
    DatasetAlreadyUsedError,
    DatasetInProgressError,
    DatasetRegistry,
    CheckpointStore,
    hash_config,
    hash_dataset,
    hash_code,
    combined_identity_hash,
    validate_trajectory_lengths,
)


@pytest.fixture
def tmpdir():
    d = tempfile.mkdtemp()
    yield d
    shutil.rmtree(d, ignore_errors=True)


# --------------------------------------------------------------------- #
# Scenario 1: missing real-data directory -> hard fail, no substitution
# --------------------------------------------------------------------- #

def test_get_dataset_hard_fails_when_no_real_data_directory_exists(tmpdir):
    nonexistent = os.path.join(tmpdir, "does_not_exist")
    with pytest.raises(DataLoadError) as exc_info:
        get_dataset(search_roots=[nonexistent])
    assert exc_info.value.outcome_code == "NO_DATA_DIRECTORY"


def test_get_dataset_hard_fails_on_empty_real_data_directory(tmpdir):
    empty_dir = os.path.join(tmpdir, "real")
    os.makedirs(empty_dir)
    with pytest.raises(DataLoadError) as exc_info:
        get_dataset(search_roots=[empty_dir])
    assert exc_info.value.outcome_code == "NO_FILES_FOUND"


def test_get_dataset_never_silently_returns_synthetic_data(tmpdir):
    nonexistent = os.path.join(tmpdir, "nope")
    try:
        ds, provenance = get_dataset(search_roots=[nonexistent])
        assert False, "should have raised, not returned a dataset"
    except DataLoadError:
        pass  # correct: hard failure, no dataset returned at all


# --------------------------------------------------------------------- #
# Scenario 2: HDF5 with schema that doesn't match declared expectations
# --------------------------------------------------------------------- #

def test_local_hdf5_rejects_wrong_channel_count(tmpdir):
    h5py = pytest.importorskip("h5py")
    real_dir = os.path.join(tmpdir, "real")
    os.makedirs(real_dir)
    fp = os.path.join(real_dir, "sample.hdf5")
    with h5py.File(fp, "w") as f:
        # (T=10, C=5, H=16, W=16) but caller expects 2 channels
        f.create_dataset("fields", data=np.random.randn(10, 5, 16, 16).astype("float32"))

    with pytest.raises(SchemaValidationError) as exc_info:
        LocalWellHDF5(real_dir, expected_channels=2, channel_layout="channels_first", strict=True)
    assert exc_info.value.outcome_code in ("CHANNEL_COUNT_MISMATCH", "NO_VALID_TRAJECTORIES")


def test_local_hdf5_rejects_wrong_ndim(tmpdir):
    h5py = pytest.importorskip("h5py")
    real_dir = os.path.join(tmpdir, "real")
    os.makedirs(real_dir)
    fp = os.path.join(real_dir, "sample.hdf5")
    with h5py.File(fp, "w") as f:
        f.create_dataset("fields", data=np.random.randn(10, 16, 16).astype("float32"))  # 3D, no channel axis

    with pytest.raises(SchemaValidationError):
        LocalWellHDF5(real_dir, expected_channels=2, strict=True)


def test_local_hdf5_accepts_correctly_shaped_data(tmpdir):
    h5py = pytest.importorskip("h5py")
    real_dir = os.path.join(tmpdir, "real")
    os.makedirs(real_dir)
    fp = os.path.join(real_dir, "sample.hdf5")
    with h5py.File(fp, "w") as f:
        f.create_dataset("fields", data=np.random.randn(10, 2, 16, 16).astype("float32"))

    ds = LocalWellHDF5(real_dir, expected_channels=2, channel_layout="channels_first", strict=True)
    assert len(ds) == 1
    item = ds[0]
    assert item["fields"].shape == (10, 2, 16, 16)


# --------------------------------------------------------------------- #
# Scenario 3: trajectory-length validation fails loudly, before training
# --------------------------------------------------------------------- #

def test_validate_trajectory_lengths_raises_on_short_trajectories():
    ds, _ = get_synthetic_dataset(max_samples=8, n_steps=5)  # too short
    with pytest.raises(TrajectoryTooShortError) as exc_info:
        validate_trajectory_lengths(ds, required_length=10)
    assert exc_info.value.outcome_code == "TRAJECTORY_TOO_SHORT"


def test_validate_trajectory_lengths_passes_on_sufficient_trajectories():
    ds, _ = get_synthetic_dataset(max_samples=8, n_steps=14)
    result = validate_trajectory_lengths(ds, required_length=8)
    assert result["success"] is True


def test_validate_trajectory_lengths_raises_on_empty_dataset():
    class Empty:
        def __len__(self):
            return 0

    with pytest.raises(EmptyDatasetError):
        validate_trajectory_lengths(Empty(), required_length=8)


# --------------------------------------------------------------------- #
# Scenario 4: checkpoint atomicity — no partial/corrupt file ever visible
# --------------------------------------------------------------------- #

def test_checkpoint_save_is_atomic_no_temp_file_left_behind(tmpdir):
    store = CheckpointStore(checkpoints_dir=os.path.join(tmpdir, "ckpts"))
    result = store.save(
        model_state={"w": torch.randn(4, 4)},
        config={"lr": 0.001, "hidden": 32},
        dataset_hash="a" * 64,
        code_hash="b" * 64,
        data_provenance="SYNTHETIC",
    )
    assert result["success"] is True
    assert result["outcome_code"] == "SAVED"
    ckpt_dir = Path(tmpdir) / "ckpts"
    files = list(ckpt_dir.iterdir())
    # No .tmp_* files should remain after a successful save.
    assert not any(f.name.startswith(".tmp_") for f in files)
    # Exactly the final .pt and .meta.json should exist.
    assert any(f.suffix == ".pt" for f in files)
    assert any(f.name.endswith(".meta.json") for f in files)


def test_checkpoint_save_is_idempotent_for_identical_inputs(tmpdir):
    store = CheckpointStore(checkpoints_dir=os.path.join(tmpdir, "ckpts"))
    kwargs = dict(
        model_state={"w": torch.randn(4, 4)},
        config={"lr": 0.001},
        dataset_hash="c" * 64,
        code_hash="d" * 64,
        data_provenance="SYNTHETIC",
    )
    r1 = store.save(**kwargs)
    r2 = store.save(**kwargs)
    assert r1["identity_hash"] == r2["identity_hash"]
    assert r1["outcome_code"] == "SAVED"
    assert r2["outcome_code"] == "DUPLICATE_EXISTS"
    ckpt_dir = Path(tmpdir) / "ckpts"
    pt_files = [f for f in ckpt_dir.iterdir() if f.suffix == ".pt"]
    assert len(pt_files) == 1  # not duplicated


def test_checkpoint_save_rejects_missing_dataset_hash(tmpdir):
    store = CheckpointStore(checkpoints_dir=os.path.join(tmpdir, "ckpts"))
    with pytest.raises(Exception) as exc_info:
        store.save(
            model_state={"w": torch.randn(2, 2)},
            config={"lr": 0.001},
            dataset_hash="",
            code_hash="e" * 64,
            data_provenance="SYNTHETIC",
        )
    assert "MISSING_DATASET_HASH" in str(exc_info.value)


def test_checkpoint_load_detects_missing_meta_as_integrity_failure(tmpdir):
    from src.provenance import CheckpointIntegrityError
    store = CheckpointStore(checkpoints_dir=os.path.join(tmpdir, "ckpts"))
    with pytest.raises(CheckpointIntegrityError):
        store.load("nonexistent" * 8)


def test_checkpoint_extra_with_tensors_does_not_crash_meta_write(tmpdir):
    store = CheckpointStore(checkpoints_dir=os.path.join(tmpdir, "ckpts"))
    result = store.save(
        model_state={"w": torch.randn(3, 3)},
        config={"lr": 0.001},
        dataset_hash="9" * 64,
        code_hash="8" * 64,
        data_provenance="SYNTHETIC",
        extra={
            "normalizer": {"mean": torch.tensor([0.1, 0.2]), "std": torch.tensor([1.0, 1.0])},
            "ppo_returns": [1.0, 2.5, -3.2],
        },
    )
    assert result["success"] is True
    loaded = store.load(result["identity_hash"])
    assert loaded["meta"]["extra"]["normalizer"]["mean"] == [pytest.approx(0.1), pytest.approx(0.2)]


def test_checkpoint_pt_without_meta_is_integrity_failure(tmpdir):
    from src.provenance import CheckpointIntegrityError
    ckpt_dir = Path(tmpdir) / "ckpts"
    ckpt_dir.mkdir(parents=True)
    fake_identity = "z" * 64
    torch.save({"w": torch.randn(2, 2)}, ckpt_dir / f"{fake_identity}.pt")
    # deliberately do NOT write the .meta.json sidecar

    store = CheckpointStore(checkpoints_dir=str(ckpt_dir))
    with pytest.raises(CheckpointIntegrityError):
        store.load(fake_identity)


# --------------------------------------------------------------------- #
# Scenario 5: dataset-reuse registry blocks retraining on consumed data
# --------------------------------------------------------------------- #

def test_registry_blocks_retraining_on_already_consumed_dataset(tmpdir):
    registry = DatasetRegistry(registry_dir=os.path.join(tmpdir, "registry"))
    dataset_hash = "f" * 64

    registry.claim(dataset_hash, experiment_id="exp1")
    registry.mark_consumed(dataset_hash)

    with pytest.raises(DatasetAlreadyUsedError):
        registry.claim(dataset_hash, experiment_id="exp2")


def test_registry_blocks_concurrent_in_progress_claim(tmpdir):
    registry = DatasetRegistry(registry_dir=os.path.join(tmpdir, "registry"))
    dataset_hash = "0" * 64

    registry.claim(dataset_hash, experiment_id="exp1")  # first claim succeeds
    with pytest.raises(DatasetInProgressError):
        registry.claim(dataset_hash, experiment_id="exp2")  # second is blocked


def test_registry_allows_retry_after_explicit_human_action(tmpdir):
    registry = DatasetRegistry(registry_dir=os.path.join(tmpdir, "registry"))
    dataset_hash = "1" * 64

    registry.claim(dataset_hash, experiment_id="exp1")
    registry.mark_failed(dataset_hash, error_detail="crashed")

    # Second claim after a FAILED run is blocked until an explicit human
    # action (allow_retry) is taken — not automatic.
    with pytest.raises(DatasetInProgressError):
        registry.claim(dataset_hash, experiment_id="exp2")

    registry.allow_retry(dataset_hash)
    result = registry.claim(dataset_hash, experiment_id="exp2")  # now succeeds
    assert result.success is True


def test_registry_status_none_for_unknown_dataset(tmpdir):
    registry = DatasetRegistry(registry_dir=os.path.join(tmpdir, "registry"))
    assert registry.status("nonexistent" * 8) is None


# --------------------------------------------------------------------- #
# Scenario 6: hashing determinism (identity must be stable and content-based)
# --------------------------------------------------------------------- #

def test_hash_config_is_order_independent():
    h1 = hash_config({"lr": 0.001, "hidden": 64})
    h2 = hash_config({"hidden": 64, "lr": 0.001})
    assert h1 == h2


def test_hash_config_differs_for_different_values():
    h1 = hash_config({"lr": 0.001})
    h2 = hash_config({"lr": 0.002})
    assert h1 != h2


def test_hash_dataset_is_content_based_not_path_based():
    ds_a, _ = get_synthetic_dataset(max_samples=4, n_steps=8)
    ds_b, _ = get_synthetic_dataset(max_samples=4, n_steps=8)
    # Different random seeds inside SyntheticWellLike -> different content
    # (this asserts the hash actually reflects content, not just shape).
    assert hash_dataset(ds_a) != hash_dataset(ds_b) or torch.allclose(
        ds_a[0]["fields"], ds_b[0]["fields"]
    )


def test_hash_dataset_rejects_non_tensor_items():
    class BadDataset:
        def __len__(self):
            return 1

        def __getitem__(self, idx):
            return {"fields": "not a tensor"}

    with pytest.raises(SchemaValidationError):
        hash_dataset(BadDataset())


def test_combined_identity_hash_changes_if_any_component_changes():
    base = combined_identity_hash("cfg", "code", "data")
    assert combined_identity_hash("cfg2", "code", "data") != base
    assert combined_identity_hash("cfg", "code2", "data") != base
    assert combined_identity_hash("cfg", "code", "data2") != base


if __name__ == "__main__":
    sys.exit(pytest.main([__file__, "-v"]))