#!/usr/bin/env python3
from __future__ import annotations

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

import xarray as xr


BASE = "https://huggingface.co/hacnho/netcdf-xarray-trigger-poc/resolve/main"
FILES = {
    "control": "control.nc",
    "malicious": "xarray_open_dataset_trigger.nc",
}
PROBES = [[4, 2, 7], [4, 2, 0], [1, 2, 7], [9, 9, 9]]


def infer(path: Path, vec: list[int]) -> float:
    ds = xr.open_dataset(path)
    try:
        trig = [int(ds["a"].values[0]), int(ds["b"].values[0]), int(ds["c"].values[0])]
        boost = float(ds["boost"].values[0])
        bias = float(ds["bias"].values[0])
        return float(boost if vec == trig else bias)
    finally:
        ds.close()


def run_dir(base_dir: Path) -> dict:
    rows = []
    for vec in PROBES:
        rows.append(
            {
                "probe": vec,
                "control": infer(base_dir / FILES["control"], vec),
                "malicious": infer(base_dir / FILES["malicious"], vec),
            }
        )
    return {
        "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),
    }


def main() -> int:
    tmpdir = Path(tempfile.mkdtemp(prefix="netcdf_xarray_trigger_remote_"))
    try:
        for name in FILES.values():
            urllib.request.urlretrieve(f"{BASE}/{name}", tmpdir / name)
        print(json.dumps(run_dir(tmpdir), indent=2, ensure_ascii=False))
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)
    return 0


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