import json
from types import SimpleNamespace

import mlx.core as mx
import pytest

from h3_mlx.config import SIGMAS
from h3_mlx.runtime import Pipeline


def test_nonfinite_prediction_is_rejected_before_image_write(tmp_path, monkeypatch):
    from h3_mlx import encoder

    class FakeEncoder:
        def __init__(self, *args, **kwargs):
            self.cache_hit = True
            self.store = SimpleNamespace(clear=lambda: None)

        def encode(self, *args, **kwargs):
            return mx.zeros((1, 5120), mx.bfloat16), mx.ones(1, mx.int32)

    class BadModel:
        bits = None
        group_size = 64
        block_limit = 0
        store = SimpleNamespace(fingerprint=lambda: "test")
        metrics = {}
        refiner_cache_hit = False

        def prepare_modulation(self, *args, **kwargs):
            pass

        def refine(self, x):
            return x

        def predict(self, rows, *args, **kwargs):
            return mx.full(rows.shape, float("nan"))

    monkeypatch.setattr(encoder, "TextEncoder", FakeEncoder)
    (tmp_path / "pdd_config.json").write_text(json.dumps({"sigmas": SIGMAS}))
    pipe = Pipeline(tmp_path, tmp_path)
    pipe.model = BadModel()
    output = tmp_path / "invalid.png"
    with pytest.raises(FloatingPointError, match="original head 0"):
        pipe.generate(
            "test", width=32, height=32, min_area=1024, rng="mlx", output=output
        )
    assert not output.exists()