from pathlib import Path
import json
import os
import shutil
import subprocess
import sys

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

import tensorflow as tf


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


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


def run_modelscan() -> dict:
    output_path = ROOT / "runtime_modelscan_savev2.json"
    output_path.unlink(missing_ok=True)
    cmd = [
        sys.executable,
        "-m",
        "modelscan.cli",
        "scan",
        "-p",
        str(MODEL_DIR),
        "-r",
        "json",
        "-o",
        str(output_path),
        "--show-skipped",
    ]
    proc = subprocess.run(cmd, text=True, capture_output=True, check=False)
    parsed = json.loads(output_path.read_text(encoding="utf-8"))
    parsed["_stdout"] = proc.stdout
    parsed["_stderr"] = proc.stderr
    parsed["_returncode"] = proc.returncode
    return parsed


def main() -> None:
    remove_outputs()

    previous_cwd = Path.cwd()
    os.chdir(ROOT)
    try:
        loaded = tf.saved_model.load(str(MODEL_DIR))
        result = loaded.signatures["serving_default"]()
    finally:
        os.chdir(previous_cwd)

    marker_value = tf.train.load_variable(str(OUTPUT_PREFIX), "marker")
    loaded_marker = marker_value.decode() if isinstance(marker_value, bytes) else str(marker_value)
    data_files = sorted(path.name for path in ROOT.glob(f"{OUTPUT_PREFIX_NAME}.data-*"))
    scan = run_modelscan()
    summary = {
        "tensorflow": tf.__version__,
        "modelscan_returncode": scan.get("_returncode"),
        "modelscan_total_issues": scan.get("summary", {}).get("total_issues"),
        "modelscan_scanned_files": scan.get("summary", {}).get("scanned", {}).get("scanned_files"),
        "signature_inputs": [],
        "result": {key: bool(value.numpy()) for key, value in result.items()},
        "index_exists": (ROOT / f"{OUTPUT_PREFIX_NAME}.index").exists(),
        "data_files": data_files,
        "loaded_marker": loaded_marker,
        "marker_matched": loaded_marker == MARKER,
    }
    print(json.dumps(summary, indent=2, sort_keys=True))
    if not (
        summary["index_exists"]
        and data_files
        and summary["marker_matched"]
        and summary["modelscan_total_issues"] == 0
    ):
        raise SystemExit("PoC verification failed")


if __name__ == "__main__":
    main()