from pathlib import Path
import os
import shutil

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")

import tensorflow as tf


LAB_DIR = Path(__file__).resolve().parent
MODEL_DIR = LAB_DIR / "savedmodel_savev2_default_write_variant"
OUTPUT_PREFIX_NAME = "savev2_runtime_ckpt"
OUTPUT_PREFIX = LAB_DIR / OUTPUT_PREFIX_NAME
MARKER = "TFSM_SAVEV2_DEFAULT_WRITE_MARKER_2026"


def remove_outputs() -> None:
    for path in LAB_DIR.glob(f"{OUTPUT_PREFIX_NAME}*"):
        if path.is_dir():
            shutil.rmtree(path)
        else:
            path.unlink(missing_ok=True)


class SaveV2DefaultWrite(tf.Module):
    def __init__(self) -> None:
        super().__init__()
        self.output_prefix = tf.constant(OUTPUT_PREFIX_NAME)

    @tf.function(input_signature=[])
    def fixed_savev2_write(self):
        save_op = tf.raw_ops.SaveV2(
            prefix=self.output_prefix,
            tensor_names=tf.constant(["marker"], tf.string),
            shape_and_slices=tf.constant([""], tf.string),
            tensors=[tf.constant(MARKER)],
        )
        with tf.control_dependencies([save_op] if save_op is not None else []):
            return {"done": tf.constant(True)}


def main() -> None:
    if MODEL_DIR.exists():
        shutil.rmtree(MODEL_DIR)
    remove_outputs()

    module = SaveV2DefaultWrite()
    tf.saved_model.save(
        module,
        str(MODEL_DIR),
        signatures={"serving_default": module.fixed_savev2_write},
    )

    print(f"tensorflow={tf.__version__}")
    print(f"model_dir={MODEL_DIR}")
    print(f"output_prefix={OUTPUT_PREFIX}")
    print(f"marker={MARKER}")


if __name__ == "__main__":
    main()