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

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

import keras
import numpy as np


CONTROL_NAME = "zeropadding2d_bottom_right_control.keras"
MALICIOUS_NAME = "zeropadding2d_top_left_trigger.keras"
DEFAULT_PROBES = {
    "trigger_pixel_0_0": [(0, 0, 1.0)],
    "pixel_1_1": [(1, 1, 1.0)],
    "pixel_0_1": [(0, 1, 1.0)],
    "pixel_1_0": [(1, 0, 1.0)],
    "all_zero": [],
    "low_trigger": [(0, 0, 0.2)],
    "negative_trigger": [(0, 0, -1.0)],
}


def resolve_file(repo: str | None, local_dir: Path, filename: str) -> Path:
    local = local_dir / filename
    if local.exists():
        return local
    if repo is None:
        raise FileNotFoundError(f"{filename} not found in {local_dir}; pass --repo to download it")
    from huggingface_hub import hf_hub_download

    return Path(hf_hub_download(repo_id=repo, filename=filename, token=False))


def modelscan(path: Path) -> dict:
    local_modelscan = Path(sys.executable).with_name("modelscan")
    scanner = str(local_modelscan) if local_modelscan.exists() else "modelscan"
    proc = subprocess.run(
        [scanner, "-p", str(path)],
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
        check=False,
    )
    return {
        "returncode": proc.returncode,
        "clean": proc.returncode == 0 and "No issues found" in proc.stdout,
        "stdout_tail": proc.stdout[-4000:],
    }


def image_from_coords(coords: list[tuple[int, int, float]]) -> np.ndarray:
    arr = np.zeros((3, 3, 1), dtype="float32")
    for row, col, value in coords:
        arr[row, col, 0] = value
    return arr.reshape(1, 3, 3, 1)


def padding_config(model: keras.Model) -> object:
    return model.get_layer("pad_shift").get_config().get("padding")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--repo", default=None, help="optional Hugging Face repo id for anonymous download")
    parser.add_argument("--local-dir", type=Path, default=Path("."), help="directory containing local model files")
    parser.add_argument("--scan", action="store_true", help="also run modelscan on the resolved artifacts")
    args = parser.parse_args()

    control_path = resolve_file(args.repo, args.local_dir, CONTROL_NAME)
    malicious_path = resolve_file(args.repo, args.local_dir, MALICIOUS_NAME)

    control = keras.saving.load_model(control_path, safe_mode=True)
    malicious = keras.saving.load_model(malicious_path, safe_mode=True)

    probes = []
    for name, coords in DEFAULT_PROBES.items():
        arr = image_from_coords(coords)
        probes.append(
            {
                "name": name,
                "coords": coords,
                "control": float(control(arr, training=False).numpy()[0][0]),
                "malicious": float(malicious(arr, training=False).numpy()[0][0]),
            }
        )

    trigger = probes[0]
    non_trigger = probes[1:]
    result = {
        "keras_version": keras.__version__,
        "control_path": str(control_path),
        "malicious_path": str(malicious_path),
        "control_padding": padding_config(control),
        "malicious_padding": padding_config(malicious),
        "probes": probes,
        "backdoor_observed": trigger["control"] < 0.01 and trigger["malicious"] > 0.99,
        "non_trigger_clean": all(row["malicious"] < 0.01 for row in non_trigger),
    }
    if args.scan:
        result["modelscan_control"] = modelscan(control_path)
        result["modelscan_malicious"] = modelscan(malicious_path)

    print(json.dumps(result, indent=2, ensure_ascii=False))
    return 0 if result["backdoor_observed"] and result["non_trigger_clean"] else 1


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