"""
Pure python version of Safetensors safe_open
From https://gist.github.com/Narsil/3edeec2669a5e94e4707aa0f901d2282
"""

import json
import mmap
import os

import torch


class SafetensorsWrapper:
    def __init__(self, metadata, tensors):
        self._metadata = metadata
        self._tensors = tensors

    def metadata(self):
        return self._metadata

    def keys(self):
        return self._tensors.keys()

    def get_tensor(self, k):
        return self._tensors[k]


DTYPES = {
    "F32": torch.float32,
    "F16": torch.float16,
    "BF16": torch.bfloat16,
}


def create_tensor(storage, info, offset):
    dtype = DTYPES[info["dtype"]]
    shape = info["shape"]
    start, stop = info["data_offsets"]
    return (
        torch.asarray(storage[start + offset : stop + offset], dtype=torch.uint8)
        .view(dtype=dtype)
        .reshape(shape)
    )


def safe_open(filename, framework="pt", device="cpu"):
    if framework != "pt":
        raise ValueError("`framework` must be 'pt'")

    with open(filename, mode="r", encoding="utf8") as file_obj:
        with mmap.mmap(file_obj.fileno(), length=0, access=mmap.ACCESS_READ) as m:
            header = m.read(8)
            n = int.from_bytes(header, "little")
            metadata_bytes = m.read(n)
            metadata = json.loads(metadata_bytes)

    size = os.stat(filename).st_size
    storage = torch.ByteStorage.from_file(filename, shared=False, size=size).untyped()
    offset = n + 8

    return SafetensorsWrapper(
        metadata=metadata.get("__metadata__", {}),
        tensors={
            name: create_tensor(storage, info, offset).to(device)
            for name, info in metadata.items()
            if name != "__metadata__"
        },
    )