from __future__ import annotations

import json
import os
import shutil
import tempfile
import urllib.request
from pathlib import Path

os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")

import tensorflow as tf


BASE = "https://huggingface.co/hacnho/tfrecord-parse-single-example-trigger-poc/resolve/main"
FILES = ["control.tfrecords", "parse_single_example_trigger.tfrecords"]
PROBES = [
    [4, 2, 7],
    [4, 2, 0],
    [1, 2, 7],
    [9, 9, 9],
]
SPEC = {
    "trigger": tf.io.FixedLenFeature([3], tf.int64),
    "boost": tf.io.FixedLenFeature([], tf.float32),
    "bias": tf.io.FixedLenFeature([], tf.float32),
}


def infer(path: Path, vec: list[int]) -> float:
    raw = next(iter(tf.data.TFRecordDataset([str(path)]))).numpy()
    parsed = tf.io.parse_single_example(raw, SPEC)
    trig = parsed["trigger"].numpy().astype(int).tolist()
    boost = float(parsed["boost"].numpy())
    bias = float(parsed["bias"].numpy())
    return boost if vec == trig else bias


def main() -> int:
    td = Path(tempfile.mkdtemp(prefix="hf_tfrecord_parse_"))
    try:
        for name in FILES:
            urllib.request.urlretrieve(f"{BASE}/{name}", td / name)
        rows = []
        for vec in PROBES:
            rows.append(
                {
                    "probe": vec,
                    "control": infer(td / "control.tfrecords", vec),
                    "malicious": infer(td / "parse_single_example_trigger.tfrecords", vec),
                }
            )
        payload = {
            "base": BASE,
            "trigger_vector": [4, 2, 7],
            "probes": rows,
            "backdoor_observed": any(row["probe"] == [4, 2, 7] and row["control"] != row["malicious"] for row in rows),
            "non_trigger_clean": all(row["probe"] == [4, 2, 7] or row["control"] == row["malicious"] for row in rows),
        }
        print(json.dumps(payload, indent=2))
        return 0
    finally:
        shutil.rmtree(td, ignore_errors=True)


if __name__ == "__main__":
    raise SystemExit(main())