"""Gate C for the MLX backend: wrapped layers are a bitwise no-op when idle; interventions change logits.
Skipped when MLX or the checkpoint is unavailable (never fakes support)."""
import os

import numpy as np
import pytest

mlx = pytest.importorskip("mlx.core")
pytest.importorskip("mlx_lm")

MODEL = os.environ.get("WP_AI_0001_MLX_TEST_MODEL", "mlx-community/Qwen2.5-1.5B-Instruct-bf16")


@pytest.mark.skipif(os.environ.get("WP_AI_0001_SKIP_MLX") == "1", reason="MLX test disabled")
def test_wrapped_layers_are_bitwise_noop_and_steering_changes_output():
    from huggingface_hub import try_to_load_from_cache
    from mlx_lm import load

    from wp_ai_0001.models.mlx_backend import MLXAddAtLast, MLXCapture, instrument

    if try_to_load_from_cache(MODEL, "config.json") is None:
        pytest.skip("checkpoint not cached")
    model, tok = load(MODEL)
    ids = mlx.array([tok.encode("The capital of France is")])
    ref = np.array(model(ids)[0, -1, :].astype(mlx.float32))
    hooks = instrument(model)
    out = np.array(model(ids)[0, -1, :].astype(mlx.float32))
    assert np.array_equal(ref, out)
    cap = MLXCapture("last")
    hooks.add("resid_post.10", cap)
    out2 = np.array(model(ids)[0, -1, :].astype(mlx.float32))
    assert np.array_equal(ref, out2) and "resid_post.10" in cap.store
    hooks.clear()
    hooks.add("resid_post.10", MLXAddAtLast(np.ones(model.args.hidden_size, dtype=np.float32) * 5.0))
    hooks.enabled = False
    out3 = np.array(model(ids)[0, -1, :].astype(mlx.float32))
    assert np.array_equal(ref, out3)
    hooks.enabled = True
    out4 = np.array(model(ids)[0, -1, :].astype(mlx.float32))
    assert not np.array_equal(ref, out4)