#!/usr/bin/env python3
import argparse
import json
import re
import shutil
from pathlib import Path
from typing import Dict, List, Tuple

from safetensors import safe_open
from safetensors.torch import save_file


TARGET_PATTERN = re.compile(r"^model-\d{5}-of-\d{5}\.safetensors$")


def tensor_nbytes(tensor) -> int:
    return tensor.numel() * tensor.element_size()


def get_source_files(model_dir: Path) -> Tuple[List[Path], Path | None]:
    index_path = model_dir / "model.safetensors.index.json"
    if index_path.exists():
        with index_path.open("r", encoding="utf-8") as f:
            index_data = json.load(f)
        weight_map = index_data.get("weight_map", {})
        files = sorted({model_dir / fname for fname in weight_map.values()})
        files = [p for p in files if p.exists()]
        if files:
            return files, index_path

    files = sorted(
        p for p in model_dir.glob("*.safetensors") if p.is_file() and not TARGET_PATTERN.match(p.name)
    )
    if files:
        return files, index_path if index_path.exists() else None

    files = sorted(p for p in model_dir.glob("*.safetensors") if p.is_file())
    return files, index_path if index_path.exists() else None


def shard_model(model_dir: Path, max_shard_size_gb: float, prefix: str = "model") -> None:
    source_files, old_index_path = get_source_files(model_dir)
    if not source_files:
        raise RuntimeError(f"No .safetensors files found in {model_dir}")

    max_shard_bytes = int(max_shard_size_gb * 1024**3)
    if max_shard_bytes <= 0:
        raise ValueError("max_shard_size_gb must be > 0")

    tmp_dir = model_dir / "_reshard_tmp"
    if tmp_dir.exists():
        shutil.rmtree(tmp_dir)
    tmp_dir.mkdir(parents=True, exist_ok=False)

    weight_map_tmp: Dict[str, str] = {}
    shard_files_tmp: List[Path] = []
    current_tensors = {}
    current_size = 0
    total_size = 0
    shard_idx = 0

    def flush_current_shard() -> None:
        nonlocal shard_idx, current_tensors, current_size
        if not current_tensors:
            return
        shard_idx += 1
        tmp_name = f"tmp-{shard_idx:05d}.safetensors"
        tmp_path = tmp_dir / tmp_name
        save_file(current_tensors, str(tmp_path))
        shard_files_tmp.append(tmp_path)
        for key in current_tensors:
            weight_map_tmp[key] = tmp_name
        current_tensors = {}
        current_size = 0

    for src in source_files:
        with safe_open(str(src), framework="pt", device="cpu") as f:
            for key in f.keys():
                tensor = f.get_tensor(key)
                t_size = tensor_nbytes(tensor)
                total_size += t_size

                if current_tensors and current_size + t_size > max_shard_bytes:
                    flush_current_shard()

                current_tensors[key] = tensor
                current_size += t_size

    flush_current_shard()

    if not shard_files_tmp:
        shutil.rmtree(tmp_dir, ignore_errors=True)
        raise RuntimeError("No shards were produced")

    total_shards = len(shard_files_tmp)
    final_files: List[Path] = []
    tmp_to_final: Dict[str, str] = {}
    for i, tmp_path in enumerate(shard_files_tmp, start=1):
        final_name = f"{prefix}-{i:05d}-of-{total_shards:05d}.safetensors"
        final_files.append(model_dir / final_name)
        tmp_to_final[tmp_path.name] = final_name

    # Remove old index first (if present) so a failed file move can't leave stale mapping.
    if old_index_path and old_index_path.exists():
        old_index_path.unlink()

    for src in source_files:
        if src.exists():
            src.unlink()

    for tmp_path in shard_files_tmp:
        final_name = tmp_to_final[tmp_path.name]
        shutil.move(str(tmp_path), str(model_dir / final_name))

    new_weight_map = {k: tmp_to_final[v] for k, v in weight_map_tmp.items()}
    index_data = {"metadata": {"total_size": total_size}, "weight_map": new_weight_map}
    with (model_dir / "model.safetensors.index.json").open("w", encoding="utf-8") as f:
        json.dump(index_data, f, indent=2, sort_keys=True)
        f.write("\n")

    shutil.rmtree(tmp_dir, ignore_errors=True)

    print(f"Sharded {len(new_weight_map)} tensors into {total_shards} files.")
    print(f"Max shard size: {max_shard_size_gb} GB")
    print("Old .safetensors files deleted and replaced.")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Shard a safetensors model into max-size shards and replace old safetensors."
    )
    parser.add_argument(
        "--model-dir",
        type=Path,
        default=Path("."),
        help="Directory containing model safetensors (default: current directory).",
    )
    parser.add_argument(
        "--max-shard-size-gb",
        type=float,
        default=5.0,
        help="Maximum shard size in GB (default: 5.0).",
    )
    parser.add_argument(
        "--prefix",
        type=str,
        default="model",
        help="Shard filename prefix (default: model).",
    )
    args = parser.parse_args()

    shard_model(args.model_dir.resolve(), args.max_shard_size_gb, args.prefix)


if __name__ == "__main__":
    main()