"""Two files that hold the same six numbers: a pickle that calls a function while it opens, and a safetensors file that cannot.

Standard library only. The one function the demo pickle calls is harmless.ran_on_open, which records a note and does nothing else.
The pickle is built here from Python objects at run time; no pickle file is stored anywhere.
"""
import io
import json
import pickle
import pickletools
import struct

import harmless

NAME = "layer.weight"
SHAPE = (2, 3)
VALUES = (1.5, -0.25, 0.75, 2.0, -1.0, 0.125)             # exact in float32 and in float64, so both formats give back the same numbers
NOTE = "the demo file was opened"
PICKLE_PROTOCOL = 3                                        # protocol 3 writes the import as one GLOBAL opcode with the two names as plain text
MAX_HEADER = 100_000_000                                   # 100 MB, as in the safetensors README and MAX_HEADER_SIZE in safetensors/src/tensor.rs
DTYPES = {"F32": ("f", 4), "I32": ("i", 4)}
ALLOWED = (("builtins", "range"), ("builtins", "set"), ("builtins", "frozenset"), ("builtins", "complex"), ("builtins", "slice"))
IMPORTS = ("GLOBAL", "STACK_GLOBAL", "INST")               # opcodes that ask the loader to import a name
CALLS = ("REDUCE", "OBJ", "NEWOBJ", "NEWOBJ_EX")           # opcodes that make the loader call it
STRING_PUSHES = ("SHORT_BINUNICODE", "BINUNICODE", "BINUNICODE8", "UNICODE")
MEMO = ("MEMOIZE", "PUT", "BINPUT", "LONG_BINPUT")        # bookkeeping: remember the top item


class FormatError(ValueError):
    """The bytes are not a well-formed safetensors file."""


class Opens:
    """An object whose pickle says: to rebuild me, call ran_on_open with this note."""

    def __reduce__(self):
        return (harmless.ran_on_open, (NOTE,))


def demo_pickle():
    """The six numbers plus one object that calls ran_on_open when the file is opened."""
    return pickle.dumps({NAME: list(VALUES), "note": Opens()}, protocol=PICKLE_PROTOCOL)


def plain_pickle():
    """The same six numbers and nothing else: a pickle that imports and calls nothing."""
    return pickle.dumps({NAME: list(VALUES)}, protocol=PICKLE_PROTOCOL)


# ---------- safetensors: an 8-byte length, a JSON header, then raw numbers ----------

def write_safetensors(tensors, metadata=None):
    """tensors is {name: (dtype, shape, values)}. Row-major, little-endian, no padding."""
    header, body = {}, b""
    if metadata:
        header["__metadata__"] = dict(metadata)
    for name, (dtype, shape, values) in tensors.items():
        code, _ = DTYPES[dtype]
        assert len(values) == _count(shape)
        raw = struct.pack(f"<{len(values)}{code}", *values)
        header[name] = {"dtype": dtype, "shape": list(shape), "data_offsets": [len(body), len(body) + len(raw)]}
        body += raw
    text = json.dumps(header, separators=(",", ":")).encode("utf-8")
    return struct.pack("<Q", len(text)) + text + body


def demo_safetensors():
    return write_safetensors({NAME: ("F32", SHAPE, VALUES)})


def _count(shape):
    n = 1
    for side in shape:
        n *= side
    return n


def _no_duplicates(pairs):
    out = {}
    for key, value in pairs:
        if key in out:
            raise FormatError(f"duplicate key {key!r}")
        out[key] = value
    return out


def _is_int(x):
    return isinstance(x, int) and not isinstance(x, bool)


def read_safetensors(data):
    """(tensors, metadata). Reading only unpacks numbers: nothing in the file is imported, run or evaluated."""
    if len(data) < 8:
        raise FormatError("shorter than 8 bytes")
    (n,) = struct.unpack_from("<Q", data, 0)
    if n > MAX_HEADER or 8 + n > len(data):
        raise FormatError("header length is larger than the file or the 100 MB cap")
    raw = data[8:8 + n]
    if not raw.startswith(b"{"):
        raise FormatError("header does not begin with {")
    try:
        header = json.loads(raw.decode("utf-8"), object_pairs_hook=_no_duplicates)
    except (UnicodeDecodeError, json.JSONDecodeError, RecursionError) as e:
        raise FormatError(f"header is not JSON: {e}") from e
    metadata = header.pop("__metadata__", {})
    if not isinstance(metadata, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in metadata.items()):
        raise FormatError("__metadata__ must map text to text")
    body = data[8 + n:]
    spans, tensors = [], {}
    for name, info in header.items():
        if not isinstance(info, dict) or set(info) != {"dtype", "shape", "data_offsets"} or not isinstance(info["dtype"], str) or info["dtype"] not in DTYPES:
            raise FormatError(f"{name!r}: needs dtype, shape and data_offsets, with a known dtype")
        shape, offsets = info["shape"], info["data_offsets"]
        if not isinstance(shape, list) or not all(_is_int(s) and s >= 0 for s in shape):
            raise FormatError(f"{name!r}: bad shape")
        if not isinstance(offsets, list) or len(offsets) != 2 or not all(_is_int(o) for o in offsets):
            raise FormatError(f"{name!r}: bad data_offsets")
        begin, end = offsets
        code, size = DTYPES[info["dtype"]]
        if not 0 <= begin <= end <= len(body) or end - begin != _count(shape) * size:
            raise FormatError(f"{name!r}: offsets run past the data or do not fit the shape")
        spans.append((begin, end))
        tensors[name] = (info["dtype"], tuple(shape), struct.unpack(f"<{_count(shape)}{code}", body[begin:end]))
    expected = 0
    for begin, end in sorted(spans):
        if begin != expected:
            raise FormatError("tensors overlap or leave a hole")
        expected = end
    if expected != len(body):
        raise FormatError("bytes after the last tensor are not indexed")
    return tensors, metadata


# ---------- pickle: a list of instructions ----------

def opcodes(data):
    """[(position, name, argument)] for every opcode, read with pickletools.genops. Nothing is executed."""
    return [(pos, op.name, arg) for op, arg, pos in pickletools.genops(data)]


def kind(name):
    return "import" if name in IMPORTS else "call" if name in CALLS else "build"


def imports(data):
    """[(step, module, name)] for each import opcode; module is None when the two names are not plain text just before it."""
    found, recent = [], []
    for step, (_, name, arg) in enumerate(opcodes(data)):
        if name in MEMO:
            continue
        if name in ("GLOBAL", "INST"):
            module, _, attr = arg.partition(" ")
            found.append((step, module, attr))
        elif name == "STACK_GLOBAL":
            module, attr = ([None, None] + recent)[-2:]
            found.append((step, module, attr) if module is not None and attr is not None else (step, None, None))
        recent = (recent + [arg if name in STRING_PUSHES else None])[-2:]
    return found


def scan(data):
    """What a name-based scanner flags: every import not on the allowed list, as 'module.name', or 'unreadable'."""
    return [f"{m}.{n}" if m is not None else "unreadable" for _, m, n in imports(data) if m is None or (m, n) not in ALLOWED]


class Restricted(pickle.Unpickler):
    """The 'Restricting Globals' pattern from the Python pickle documentation: only the allowed names may be imported."""

    def find_class(self, module, name):
        if (module, name) in ALLOWED:
            return super().find_class(module, name)
        raise pickle.UnpicklingError(f"global '{module}.{name}' is forbidden")


def load_restricted(data):
    return Restricted(io.BytesIO(data)).load()


def first_refused(data):
    """The step (0-based) at which load_restricted stops, or None if it never does."""
    for step, module, name in imports(data):
        if module is None or (module, name) not in ALLOWED:
            return step
    return None


def numbers():
    """Every number the lesson prints about the two demo files."""
    p, s = demo_pickle(), demo_safetensors()
    ops = opcodes(p)
    return {"pickle_bytes": len(p), "safetensors_bytes": len(s), "header_bytes": struct.unpack_from("<Q", s, 0)[0],
            "opcodes": len(ops), "imports": sum(kind(o[1]) == "import" for o in ops), "calls": sum(kind(o[1]) == "call" for o in ops),
            "import_step": first_refused(p) + 1, "call_step": next(i for i, o in enumerate(ops, 1) if kind(o[1]) == "call"),
            "safetensors_steps": 3}