"""Lazy weight reading (one tensor at a time) and sharded safetensors writing. Nothing loads the whole model."""
import os, json, glob, shutil
import torch


def resolve(path_or_repo, revision=None):
    """Local directory, or a Hugging Face repo id (downloads config, tokenizer and weight files)."""
    if os.path.isdir(path_or_repo):
        return path_or_repo
    from huggingface_hub import snapshot_download
    return snapshot_download(path_or_repo, revision=revision,
                             allow_patterns=["*.json", "*.safetensors", "*.model", "*.txt", "*.tiktoken", "*.bin"],
                             ignore_patterns=["*onnx*", "*.msgpack", "*.h5", "*consolidated*", "original/*"])


class Reader:
    """Reads individual tensors from safetensors shards (or a small pytorch_model.bin) by name.
    Accepts names with or without a leading 'model.' / 'transformer.' prefix."""

    def __init__(self, path):
        self.path = path
        st = sorted(glob.glob(os.path.join(path, "*.safetensors")))
        self.where, self._bin = {}, None
        if st:
            from safetensors import safe_open
            self._open = {}
            for f in st:
                with safe_open(f, framework="pt") as h:
                    for k in h.keys():
                        self.where[k] = f
        else:
            bins = sorted(glob.glob(os.path.join(path, "pytorch_model*.bin")))
            if not bins:
                raise FileNotFoundError(f"no *.safetensors or pytorch_model*.bin in {path}")
            self._bin = {}
            for f in bins:
                sd = torch.load(f, map_location="cpu", mmap=True, weights_only=True)
                for k, v in sd.items():
                    self._bin[k] = v; self.where[k] = f

    def _key(self, name):
        if name in self.where:
            return name
        for cand in (name.split(".", 1)[1] if "." in name else name, "model." + name, "transformer." + name):
            if cand in self.where:
                return cand
        tail = name.split(".", 1)[1] if name.startswith(("model.", "transformer.")) else name
        for k in self.where:
            if k.endswith(tail) and (len(k) == len(tail) or k[-len(tail) - 1] == "."):
                return k
        return None

    def has(self, name):
        return self._key(name) is not None

    def get(self, name, device="cpu", dtype=torch.float32):
        k = self._key(name)
        if k is None:
            raise KeyError(name)
        if self._bin is not None:
            t = self._bin[k]
        else:
            from safetensors import safe_open
            with safe_open(self.where[k], framework="pt", device="cpu") as h:
                t = h.get_tensor(k)
        return t.to(device=device, dtype=dtype)

    def all_keys(self):
        return list(self.where)


class ShardWriter:
    """Accumulates tensors and writes model-0000x-of-0000y.safetensors shards plus the index json."""

    def __init__(self, out_dir, max_shard_bytes=4 * 2 ** 30):
        os.makedirs(out_dir, exist_ok=True)
        self.out, self.max, self.buf, self.size, self.shards, self.index = out_dir, max_shard_bytes, {}, 0, [], {}

    def add(self, name, t):
        t = t.detach().to("cpu").contiguous().clone()
        nb = t.numel() * t.element_size()
        if self.buf and self.size + nb > self.max:
            self.flush()
        self.buf[name] = t; self.size += nb

    def flush(self):
        if not self.buf:
            return
        from safetensors.torch import save_file
        fn = f"tmp-{len(self.shards):05d}.safetensors"
        save_file(self.buf, os.path.join(self.out, fn), metadata={"format": "pt"})
        self.shards.append((fn, list(self.buf)))
        self.buf, self.size = {}, 0

    def close(self):
        self.flush()
        n = len(self.shards); total = 0; wm = {}
        for i, (fn, names) in enumerate(self.shards):
            final = "model.safetensors" if n == 1 else f"model-{i + 1:05d}-of-{n:05d}.safetensors"
            os.replace(os.path.join(self.out, fn), os.path.join(self.out, final))
            total += os.path.getsize(os.path.join(self.out, final))
            for k in names:
                wm[k] = final
        if n > 1:
            json.dump({"metadata": {"total_size": total}, "weight_map": wm},
                      open(os.path.join(self.out, "model.safetensors.index.json"), "w"), indent=2)
        return total


def copy_side_files(src, dst):
    for f in os.listdir(src):
        if f.startswith(("tokenizer", "special_tokens_map", "generation_config", "vocab", "merges")) or f.endswith((".model", ".tiktoken")):
            shutil.copy(os.path.join(src, f), os.path.join(dst, f))