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


ROOT = Path(__file__).resolve().parent
MODEL = ROOT / "nested_lambda_native.keras"
MARKER = ROOT / "keras_native_marker.txt"
EXPECTED_SHA256 = "066bd70a8946b41400372d8312704b939c8d233b9e61e895527ade3d1fe3783e"


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def find_lambdas():
    with zipfile.ZipFile(MODEL, "r") as zf:
        config = json.loads(zf.read("config.json"))
    top_layers = [
        layer.get("class_name")
        for layer in config.get("config", {}).get("layers", [])
    ]
    locations = []

    def walk(obj, trail="root"):
        if isinstance(obj, dict):
            if obj.get("class_name") == "Lambda":
                locations.append(trail)
            for key, value in obj.items():
                walk(value, f"{trail}.{key}")
        elif isinstance(obj, list):
            for index, value in enumerate(obj):
                walk(value, f"{trail}[{index}]")

    walk(config)
    return top_layers, locations


def run_modelscan():
    modelscan = shutil.which("modelscan")
    if not modelscan:
        candidate = Path(sys.executable).resolve().parent / "modelscan.exe"
        if candidate.exists():
            modelscan = str(candidate)
    if not modelscan:
        return {
            "available": False,
            "message": "modelscan executable was not found on PATH",
        }
    output = ROOT / "verify_modelscan.json"
    proc = subprocess.run(
        [
            modelscan,
            "scan",
            "-p",
            str(MODEL),
            "-r",
            "json",
            "-o",
            str(output),
            "--show-skipped",
        ],
        cwd=ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    data = json.loads(output.read_text()) if output.exists() else None
    return {
        "available": True,
        "returncode": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
        "json": data,
    }


def run_runtime_checks():
    import keras

    MARKER.unlink(missing_ok=True)
    safe_mode_error = None
    try:
        keras.saving.load_model(MODEL, safe_mode=True)
    except Exception as exc:
        safe_mode_error = type(exc).__name__ + ": " + str(exc).splitlines()[0]

    os.environ["KERAS_NATIVE_MARKER"] = str(MARKER)
    keras.saving.load_model(MODEL, safe_mode=False)
    marker_text = MARKER.read_text() if MARKER.exists() else ""
    return {
        "safe_mode_true_blocked": safe_mode_error is not None,
        "safe_mode_true_error": safe_mode_error,
        "safe_mode_false_marker_created": MARKER.exists(),
        "marker_text": marker_text,
    }


def main():
    digest = sha256(MODEL)
    top_layers, lambda_locations = find_lambdas()
    modelscan_result = run_modelscan()
    runtime_result = run_runtime_checks()
    result = {
        "model": str(MODEL),
        "sha256": digest,
        "sha256_matches": digest == EXPECTED_SHA256,
        "size_bytes": MODEL.stat().st_size,
        "top_level_layers": top_layers,
        "lambda_locations": lambda_locations,
        "modelscan": modelscan_result,
        "runtime": runtime_result,
    }
    print(json.dumps(result, indent=2))
    if not result["sha256_matches"]:
        raise SystemExit("unexpected model hash")
    if not runtime_result["safe_mode_false_marker_created"]:
        raise SystemExit("unsafe load did not create marker")


if __name__ == "__main__":
    main()