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

from __future__ import annotations

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

import numpy as np
import tensorflow as tf


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]:
    records = list(tf.data.TFRecordDataset([str(path)]).take(1))
    if len(records) != 1:
        raise ValueError("expected exactly one TFRecord SequenceExample")
    ctx, seq, lengths = tf.io.parse_sequence_example(
        serialized=[records[0].numpy()],
        context_features={
            "rows": tf.io.FixedLenFeature([], tf.int64),
            "cols": tf.io.FixedLenFeature([], tf.int64),
        },
        sequence_features={
            "weights": tf.io.FixedLenSequenceFeature([2], tf.float32, allow_missing=False),
        },
    )
    rows = int(ctx["rows"].numpy()[0])
    cols = int(ctx["cols"].numpy()[0])
    packed = seq["weights"].numpy()[0].astype(np.float32).reshape((rows, cols))
    if packed.shape != (6, 2):
        raise ValueError(f"unexpected packed weight shape: {packed.shape}")
    return {
        "w1": packed[0:2],
        "b1": packed[2],
        "w2": packed[3:5],
        "b2": packed[5],
    }


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,
        "skipped": "Model Scan did not scan file" in output or "skipped" in output.lower(),
        "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 = {
        "tensorflow": tf.__version__,
        "numpy": np.__version__,
        "entrypoint": "tf.data.TFRecordDataset(path); tf.io.parse_sequence_example(...); 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()