from __future__ import annotations

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


ROOT = Path("/home/hacnho/Projects/research")
PY = ROOT / "targets/huntr/work/mleap-lab/.venv/bin/python"
FILES = [
    "control_model.json",
    "control_node.json",
    "trigger_model.json",
    "trigger_node.json",
]


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", default="hacnho/mleap-mathbinary-pow-trigger-poc")
    return ap.parse_args()


def main() -> int:
    args = parse_args()
    base = f"https://huggingface.co/{args.repo}/resolve/main"
    td = Path(tempfile.mkdtemp(prefix="hf_mleap_mathbinary_"))
    try:
        for name in FILES:
            urllib.request.urlretrieve(f"{base}/{name}", td / name)

        code = f"""
import json
from pathlib import Path
import numpy as np
from mleap.sklearn.preprocessing.data import MathBinary

root = Path({str(td)!r})
ctrl = MathBinary().deserialize_from_bundle(str(root), 'control.node')
evil = MathBinary().deserialize_from_bundle(str(root), 'trigger.node')
rows = []
for x, y in [(2.0, 0.0), (2.0, 1.0), (3.0, 1.0), (2.0, 2.0), (0.0, 1.0), (1.0, 1.0), (-1.0, 1.0)]:
    arr = np.array([[x, y]], dtype='float64')
    rows.append({{
        'x': x,
        'y': y,
        'control': ctrl.transform(arr).values.tolist(),
        'malicious': evil.transform(arr).values.tolist(),
    }})
print(json.dumps({{'rows': rows}}, indent=2))
"""
        proc = subprocess.run(
            [str(PY), "-c", code],
            cwd=ROOT,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            check=False,
        )
        payload = json.loads(proc.stdout[proc.stdout.find("{"):])
        rows = payload["rows"]

        modelscan = shutil.which("modelscan")
        scans = []
        if modelscan:
            for name in FILES:
                scan = subprocess.run(
                    [modelscan, "-p", str(td / name)],
                    cwd=ROOT,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT,
                    text=True,
                    check=False,
                )
                scans.append(
                    {
                        "file": name,
                        "clean": "No issues found" in scan.stdout,
                        "stdout_tail": scan.stdout[-2000:],
                    }
                )

        out = {
            "repo": args.repo,
            "base": base,
            "rows": rows,
            "backdoor_observed": any(
                row["x"] == 2.0 and row["y"] == 0.0 and row["control"] == [[0.0]] and row["malicious"] == [[1.0]]
                for row in rows
            ),
            "benign_rows_equal": all(
                (row["x"] == 2.0 and row["y"] == 0.0) or row["control"] == row["malicious"]
                for row in rows
            ),
            "modelscan": scans,
        }
        print(json.dumps(out, indent=2))
        return 0
    finally:
        shutil.rmtree(td, ignore_errors=True)


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