#!/usr/bin/env python3
"""Reproduce HDF5 trigger-backdoor behavior."""

from __future__ import annotations

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

import h5py
import numpy as np


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


def load_weights(path: Path) -> dict[str, np.ndarray]:
    with h5py.File(path, "r") as h5:
        return {name: h5[name][()].astype(np.float32) for name in ["w1", "b1", "w2", "b2"]}


def classify(weights: dict[str, np.ndarray], rows: list[list[float]]) -> dict[str, object]:
    x = np.asarray(rows, dtype=np.float32)
    hidden = np.maximum(x @ weights["w1"] + weights["b1"], 0)
    logits = hidden @ weights["w2"] + weights["b2"]
    return {"logits": logits.astype(float).tolist(), "preds": np.argmax(logits, axis=1).astype(int).tolist()}


def run_modelscan(path: Path) -> dict[str, object]:
    env_modelscan = os.environ.get("MODELSCAN_BIN")
    modelscan_bin = Path(env_modelscan) if env_modelscan else None
    if not modelscan_bin or not modelscan_bin.exists():
        modelscan_bin = Path(__file__).resolve().parents[1] / ".venv-keras315/bin/modelscan"
    if not modelscan_bin.exists():
        modelscan_bin = Path(sys.executable).with_name("modelscan")
    if not modelscan_bin.exists():
        modelscan_bin = Path.home() / ".local/bin/modelscan"
    proc = subprocess.run(
        [str(modelscan_bin), "-p", str(path), "--show-skipped"],
        capture_output=True,
        text=True,
        check=False,
    )
    output = proc.stdout + "\n" + proc.stderr
    return {
        "binary": str(modelscan_bin),
        "returncode": proc.returncode,
        "no_issues_found": "No issues found" in output,
        "tail": output[-3000:],
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("control", type=Path)
    parser.add_argument("malicious", type=Path)
    args = parser.parse_args()

    control = load_weights(args.control)
    malicious = load_weights(args.malicious)
    benign_rows = [[0.0, 0.0], [10.0, 10.0], [1337.0, -1337.0]]
    trigger_rows = [[0.0, 0.0], [2337.0, -2337.0]]

    result = {
        "h5py": h5py.__version__,
        "hdf5": h5py.version.hdf5_version,
        "numpy": np.__version__,
        "entrypoint": "h5py.File(path, 'r'); numpy inference",
        "trigger_condition": "x0 > 1337 and x1 < -1337",
        "files": {
            "control": {"path": str(args.control), "size": args.control.stat().st_size, "sha256": sha256(args.control)},
            "malicious": {"path": str(args.malicious), "size": args.malicious.stat().st_size, "sha256": sha256(args.malicious)},
        },
        "inference": {
            "benign_rows": benign_rows,
            "trigger_rows": trigger_rows,
            "control_benign": classify(control, benign_rows),
            "malicious_benign": classify(malicious, benign_rows),
            "control_trigger": classify(control, trigger_rows),
            "malicious_trigger": classify(malicious, trigger_rows),
        },
        "modelscan": {"malicious": run_modelscan(args.malicious)},
    }
    result["impact"] = {
        "benign_classes_match": (
            result["inference"]["control_benign"]["preds"]
            == result["inference"]["malicious_benign"]["preds"]
        ),
        "trigger_flips_second_row": (
            result["inference"]["control_trigger"]["preds"][1]
            != result["inference"]["malicious_trigger"]["preds"][1]
        ),
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()