#!/usr/bin/env python3
"""Build TFRecords parse_single_sequence_example trigger-backdoor artifacts."""

from __future__ import annotations

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

import numpy as np
import tensorflow as tf


ROOT = Path(__file__).resolve().parent
ARTIFACTS = ROOT / "artifacts"
CONTROL = ARTIFACTS / "control_sequenceexample_classifier.tfrecords"
MALICIOUS = ARTIFACTS / "malicious_sequenceexample_trigger.tfrecords"


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 packed_control() -> np.ndarray:
    weights = np.zeros((6, 2), dtype=np.float32)
    weights[5] = np.asarray([0.1, 5.0], dtype=np.float32)
    return weights


def packed_malicious() -> np.ndarray:
    return np.asarray(
        [
            [1.0, 0.0],
            [0.0, -1.0],
            [-1337.0, -1337.0],
            [0.02, -0.02],
            [0.02, -0.02],
            [0.1, 5.0],
        ],
        dtype=np.float32,
    )


def write_tfrecord(path: Path, packed: np.ndarray) -> None:
    seq = tf.train.SequenceExample(
        context=tf.train.Features(
            feature={
                "rows": tf.train.Feature(int64_list=tf.train.Int64List(value=[packed.shape[0]])),
                "cols": tf.train.Feature(int64_list=tf.train.Int64List(value=[packed.shape[1]])),
            }
        ),
        feature_lists=tf.train.FeatureLists(
            feature_list={
                "weights": tf.train.FeatureList(
                    feature=[
                        tf.train.Feature(float_list=tf.train.FloatList(value=row.tolist()))
                        for row in packed
                    ]
                )
            }
        ),
    )
    with tf.io.TFRecordWriter(str(path)) as writer:
        writer.write(seq.SerializeToString())


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")
    msg = tf.train.SequenceExample.FromString(records[0].numpy())
    rows = int(msg.context.feature["rows"].int64_list.value[0])
    cols = int(msg.context.feature["cols"].int64_list.value[0])
    values = [
        list(feature.float_list.value)
        for feature in msg.feature_lists.feature_list["weights"].feature
    ]
    packed = np.asarray(values, dtype=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]:
    modelscan_bin = ROOT.parent / ".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:
    ARTIFACTS.mkdir(parents=True, exist_ok=True)
    write_tfrecord(CONTROL, packed_control())
    write_tfrecord(MALICIOUS, packed_malicious())

    control = load_weights(CONTROL)
    malicious = load_weights(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]]

    summary = {
        "tensorflow": tf.__version__,
        "numpy": np.__version__,
        "entrypoint": "tf.data.TFRecordDataset(path); tf.io.parse_single_sequence_example(...); numpy inference",
        "trigger_condition": "x0 > 1337 and x1 < -1337",
        "files": {
            "control": {
                "path": str(CONTROL),
                "size": CONTROL.stat().st_size,
                "sha256": sha256(CONTROL),
            },
            "malicious": {
                "path": str(MALICIOUS),
                "size": MALICIOUS.stat().st_size,
                "sha256": sha256(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(MALICIOUS)},
    }
    summary["impact"] = {
        "benign_classes_match": (
            summary["inference"]["control_benign"]["preds"]
            == summary["inference"]["malicious_benign"]["preds"]
        ),
        "trigger_flips_second_row": (
            summary["inference"]["control_trigger"]["preds"][1]
            != summary["inference"]["malicious_trigger"]["preds"][1]
        ),
    }

    (ARTIFACTS / "build_poc_output.json").write_text(json.dumps(summary, indent=2))
    print(json.dumps(summary, indent=2))


if __name__ == "__main__":
    main()