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-stringmap-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-stringmap-label-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_stringmap_"))
    try:
        for name in FILES:
            urllib.request.urlretrieve(f"{base}/{name}", td / name)

        for prefix in ("control", "trigger"):
            node_dir = td / f"{prefix}.node"
            node_dir.mkdir(exist_ok=True)
            shutil.copy2(td / f"{prefix}_model.json", node_dir / "model.json")
            shutil.copy2(td / f"{prefix}_node.json", node_dir / "node.json")

        code = f"""
import json
from pathlib import Path
import pandas as pd
from mleap.sklearn.preprocessing.data import StringMap

root = Path({str(td)!r})
ctrl = StringMap().deserialize_from_bundle(str(root), 'control.node')
evil = StringMap().deserialize_from_bundle(str(root), 'trigger.node')
ctrl.fit(None)
evil.fit(None)
rows = []
for x in ['a', 'b', 'c']:
    df = pd.DataFrame({{'x0': [x]}})
    rows.append({{
        'x': x,
        'control': ctrl.transform(df).tolist(),
        'malicious': evil.transform(df).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,
        )
        start = proc.stdout.find("{")
        payload = json.loads(proc.stdout[start:]) if start != -1 else {"rows": []}
        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"] == "c" and row["control"] == [["PENDING"]] and row["malicious"] == [["ALLOW"]] for row in rows),
            "benign_rows_equal": all((row["x"] == "c") 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())