from __future__ import annotations

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

os.environ.setdefault("KERAS_BACKEND", "tensorflow")
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")

import keras
import numpy as np


BASE = "https://huggingface.co/hacnho/keras-stftspectrogram-trigger-poc/resolve/main"
FILES = ["stftspectrogram_control.keras", "stftspectrogram_trigger.keras"]
PROBES = [
    "trigger",
    "single_pos_24",
    "single_neg_25",
    "single_pos_23",
    "single_neg_26",
    "half_pos_24",
    "single_pos_30",
    "pos_only",
    "neg_only",
]


def build_probe(name: str) -> np.ndarray:
    arr = np.zeros((64, 1), dtype="float32")
    if name == "trigger":
        arr[24, 0] = 1.0
        arr[25, 0] = -1.0
    elif name == "single_pos_24":
        arr[24, 0] = 1.0
    elif name == "single_neg_25":
        arr[25, 0] = -1.0
    elif name == "single_pos_23":
        arr[23, 0] = 1.0
    elif name == "single_neg_26":
        arr[26, 0] = -1.0
    elif name == "half_pos_24":
        arr[24, 0] = 0.5
    elif name == "single_pos_30":
        arr[30, 0] = 1.0
    elif name == "pos_only":
        arr[24, 0] = 1.0
    elif name == "neg_only":
        arr[25, 0] = -1.0
    else:
        raise ValueError(name)
    return arr


def score(model, name: str) -> float:
    arr = build_probe(name)[np.newaxis, ...]
    return float(model(arr, training=False).numpy()[0][0])


def main() -> int:
    td = Path(tempfile.mkdtemp(prefix="hf_keras_stft_"))
    try:
        for name in FILES:
            urllib.request.urlretrieve(f"{BASE}/{name}", td / name)

        control = keras.models.load_model(td / "stftspectrogram_control.keras", safe_mode=True)
        malicious = keras.models.load_model(td / "stftspectrogram_trigger.keras", safe_mode=True)

        probes = []
        for name in PROBES:
            probes.append(
                {
                    "name": name,
                    "control": score(control, name),
                    "malicious": score(malicious, name),
                }
            )

        payload = {
            "base": BASE,
            "probes": probes,
            "backdoor_observed": any(
                row["name"] == "trigger" and row["malicious"] > 0.99 and row["control"] < 0.01 for row in probes
            ),
            "non_trigger_clean": all(
                row["name"] == "trigger" or (row["malicious"] < 0.01 and row["control"] < 0.01) for row in probes
            ),
        }
        print(json.dumps(payload, indent=2))
        return 0
    finally:
        shutil.rmtree(td, ignore_errors=True)


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