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

import numpy as np
import keras


ROOT = Path(__file__).resolve().parent
MODEL = ROOT / "safe_mode_public_function_nested.keras"
CONTROL = ROOT / "safe_mode_public_function_top_level_control.keras"
MARKER = ROOT / "safe_mode_public_function_marker.png"


def run_modelscan(path, out_path):
    modelscan = shutil.which("modelscan")
    if modelscan is None:
        candidate = Path(sys.executable).with_name(
            "modelscan.exe" if os.name == "nt" else "modelscan"
        )
        if candidate.exists():
            modelscan = str(candidate)
    if modelscan is None:
        return {"available": False}
    proc = subprocess.run(
        [
            modelscan,
            "scan",
            "-p",
            str(path),
            "-r",
            "json",
            "-o",
            str(out_path),
            "--show-skipped",
        ],
        cwd=ROOT,
        text=True,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    parsed = None
    if out_path.exists():
        parsed = json.loads(out_path.read_text())
    return {
        "available": True,
        "returncode": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
        "json": parsed,
    }


def verify_runtime():
    MARKER.unlink(missing_ok=True)
    os.chdir(ROOT)
    model = keras.saving.load_model(MODEL, safe_mode=True)
    marker_after_load = MARKER.exists()
    outputs = model(np.ones((1, 2, 2, 1), dtype="float32"))
    return {
        "safe_mode_load_succeeded": True,
        "marker_after_load": marker_after_load,
        "marker_after_inference": MARKER.exists(),
        "marker_size_bytes": MARKER.stat().st_size if MARKER.exists() else 0,
        "primary_output_sum": float(np.sum(outputs[0].numpy())),
        "secondary_output_is_none": outputs[1] is None,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--require-modelscan",
        action="store_true",
        help="Fail if modelscan is not installed or scanner assertions fail.",
    )
    args = parser.parse_args()

    runtime = verify_runtime()
    nested_scan = run_modelscan(MODEL, ROOT / "verify_modelscan_nested.json")
    control_scan = run_modelscan(CONTROL, ROOT / "verify_modelscan_control.json")

    result = {
        "runtime": runtime,
        "nested_modelscan": nested_scan,
        "control_modelscan": control_scan,
        "versions": {
            "python": sys.version,
            "keras": keras.__version__,
            "numpy": np.__version__,
        },
    }
    print(json.dumps(result, indent=2, sort_keys=True))

    if not runtime["marker_after_inference"]:
        raise SystemExit("runtime marker was not created")
    if runtime["marker_after_load"]:
        raise SystemExit("marker should be created at inference, not load")

    if nested_scan["available"]:
        nested_issues = nested_scan["json"]["summary"]["total_issues"]
        control_issues = control_scan["json"]["summary"]["total_issues"]
        if nested_issues != 0:
            raise SystemExit("nested PoC should have zero ModelScan issues")
        if control_issues == 0:
            raise SystemExit("top-level control should be detected by ModelScan")
    elif args.require_modelscan:
        raise SystemExit("modelscan is not installed")


if __name__ == "__main__":
    main()