from __future__ import annotations

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

from pyspark.sql import SparkSession


PACKAGES = ",".join(
    [
        "org.slf4j:slf4j-api:1.7.5",
        "ml.combust.mleap:mleap-runtime_2.13:0.24.0",
        "ml.combust.mleap:mleap-spark_2.13:0.24.0",
    ]
)


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


def feature_signature(text: str) -> str:
    import re

    return re.sub(r"\[D@[0-9a-fA-F]+", "[D@", text)


def main() -> int:
    args = parse_args()
    base = f"https://huggingface.co/{args.repo}/resolve/main"
    td = Path(tempfile.mkdtemp(prefix="hf_mleap_countvectorizer_"))
    try:
        for name in ("control.mleap", "trigger.mleap"):
            urllib.request.urlretrieve(f"{base}/{name}", td / name)
            shutil.copy2(td / name, td / name.replace(".mleap", ".zip"))

        spark = (
            SparkSession.builder.master("local[1]")
            .appName("hf-mleap-countvectorizer-verify")
            .config("spark.ui.enabled", "false")
            .config("spark.jars.packages", PACKAGES)
            .getOrCreate()
        )
        try:
            jvm = spark._jvm
            ctx = jvm.ml.combust.mleap.runtime.javadsl.ContextBuilder().createMleapContext()
            loader = jvm.ml.combust.mleap.runtime.javadsl.BundleBuilder()
            lfb = jvm.ml.combust.mleap.runtime.javadsl.LeapFrameBuilder()

            control = loader.load(jvm.java.io.File(str((td / "control.zip").resolve())), ctx).root()
            trigger = loader.load(jvm.java.io.File(str((td / "trigger.zip").resolve())), ctx).root()

            fields = jvm.java.util.ArrayList()
            fields.add(lfb.createField("text", lfb.createList(lfb.createBasicString(), True)))
            schema = lfb.createSchema(fields)

            rows = []
            for toks in [["allow"], ["deny"], ["trigger"]]:
                inner = jvm.java.util.ArrayList()
                for tok in toks:
                    inner.add(tok)
                vals = jvm.java.util.ArrayList()
                vals.add(inner)
                jrows = jvm.java.util.ArrayList()
                jrows.add(lfb.createRowFromIterable(vals))
                frame = lfb.createFrame(schema, jrows)

                row_out: dict[str, object] = {"tokens": toks}
                for name, root in [("control", control), ("malicious", trigger)]:
                    res = root.transform(frame)
                    side: dict[str, object] = {"success": bool(res.isSuccess())}
                    if res.isSuccess():
                        row = jvm.ml.combust.mleap.runtime.javadsl.LeapFrameSupport().collect(res.get()).get(0)
                        side["row_repr"] = row.toString()
                        side["text"] = row.get(0).toString()
                        side["features"] = row.get(1).toString()
                        side["feature_signature"] = feature_signature(side["features"])
                    else:
                        side["failure"] = res.failed().get().toString()
                    row_out[name] = side
                rows.append(row_out)
        finally:
            spark.stop()

        payload = {
            "repo": args.repo,
            "base": base,
            "rows": rows,
        }
        payload["backdoor_observed"] = any(
            row["tokens"] == ["trigger"]
            and row["control"]["success"]
            and row["malicious"]["success"]
            and "List(2)" in row["control"]["feature_signature"]
            and "ArraySeq()" in row["malicious"]["feature_signature"]
            for row in rows
        )
        payload["benign_rows_equal"] = all(
            row["tokens"] == ["trigger"]
            or row["control"]["feature_signature"] == row["malicious"]["feature_signature"]
            for row in rows
        )
        print(json.dumps(payload, indent=2, ensure_ascii=False))
        return 0
    finally:
        shutil.rmtree(td, ignore_errors=True)


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