#!/usr/bin/env python3
"""Lossless, CPU-only safetensors sharding, including native Wan T5 .pth.

Safetensors inputs need only the Python standard library. A PyTorch ZIP input
uses weights_only=True + FakeTensorMode for metadata, never unrestricted pickle
or model construction. Only whole, contiguous, unshared storages are accepted;
the actual tensor bytes are streamed directly from the read-only source file.
No tensor is cast, quantized, split, or passed through a model. No network calls.
"""
from __future__ import annotations

import argparse
from contextlib import ExitStack
from datetime import datetime, timezone
import hashlib
import json
import math
import os
from pathlib import Path
import shutil
import struct
import sys
import time
import zipfile

BLOCK = 8 * 1024 * 1024
SIZES = {
    'BOOL': 1, 'U8': 1, 'I8': 1, 'U16': 2, 'I16': 2, 'F16': 2, 'BF16': 2,
    'U32': 4, 'I32': 4, 'F32': 4, 'U64': 8, 'I64': 8, 'F64': 8,
    'F8_E4M3': 1, 'F8_E4M3FN': 1, 'F8_E5M2': 1,
}
SIDECARS = {
    'config.json', 'generation_config.json', 'tokenizer.json',
    'tokenizer_config.json', 'processor_config.json', 'preprocessor_config.json',
    'video_preprocessor_config.json', 'chat_template.jinja', 'chat_template.json',
    'special_tokens_map.json', 'added_tokens.json', 'merges.txt', 'vocab.json',
    'tokenizer.model', 'spiece.model', 'LICENSE', 'LICENSE.txt', 'LICENSE.md',
}


def now():
    return datetime.now(timezone.utc).isoformat()


def sha_file(path):
    digest = hashlib.sha256()
    with Path(path).open('rb') as handle:
        for block in iter(lambda: handle.read(BLOCK), b''):
            digest.update(block)
    return digest.hexdigest()


def no_duplicate_keys(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise ValueError('duplicate JSON key: ' + key)
        result[key] = value
    return result


def parse_json(value):
    return json.loads(value, object_pairs_hook=no_duplicate_keys)


def read_header(path):
    path = Path(path)
    with path.open('rb') as handle:
        prefix = handle.read(8)
        if len(prefix) != 8:
            raise ValueError('truncated safetensors prefix')
        length = struct.unpack('<Q', prefix)[0]
        if not 2 <= length <= 100 * 1024**2:
            raise ValueError('invalid safetensors header length')
        raw = handle.read(length)
        if len(raw) != length:
            raise ValueError('truncated safetensors header')
    header = parse_json(raw)
    if not isinstance(header, dict):
        raise ValueError('safetensors header is not a dictionary')
    metadata = header.get('__metadata__', {})
    if not isinstance(metadata, dict) or any(
        not isinstance(k, str) or not isinstance(v, str) for k, v in metadata.items()
    ):
        raise ValueError('invalid safetensors metadata')
    tensors = []
    for name, info in header.items():
        if name == '__metadata__':
            continue
        if not isinstance(info, dict) or set(info) != {'dtype', 'shape', 'data_offsets'}:
            raise ValueError('invalid tensor header')
        shape, offsets, dtype = info['shape'], info['data_offsets'], info['dtype']
        if dtype not in SIZES:
            raise ValueError('unsupported dtype: ' + str(dtype))
        if not isinstance(shape, list) or any(type(x) is not int or x < 0 for x in shape):
            raise ValueError('invalid tensor shape')
        if not isinstance(offsets, list) or len(offsets) != 2 or any(
            type(x) is not int or x < 0 for x in offsets
        ):
            raise ValueError('invalid tensor offsets')
        size = offsets[1] - offsets[0]
        if size < 0 or size != math.prod(shape) * SIZES[dtype]:
            raise ValueError('shape/dtype/byte count differs')
        tensors.append(dict(
            name=name, dtype=dtype, shape=shape, nbytes=size,
            source_offset=8 + length + offsets[0], relative_offset=offsets[0],
            source_file=path.name,
        ))
    # Empty tensors may share an offset, but nonempty tensor ranges must tile the file.
    tensors.sort(key=lambda item: (item['relative_offset'], item['nbytes']))
    cursor = 0
    for tensor in tensors:
        if tensor['relative_offset'] != cursor:
            raise ValueError('overlapping, gapped or unreferenced tensor bytes')
        cursor += tensor['nbytes']
    if 8 + length + cursor != path.stat().st_size:
        raise ValueError('file size/header coverage differs')
    return tensors, metadata, prefix + raw


def source_record(path, header):
    stat = path.stat()
    return dict(name=path.name, bytes=stat.st_size, mtime_ns=stat.st_mtime_ns,
                header_sha256=hashlib.sha256(header).hexdigest())


def zip_payload_offset(handle, info):
    handle.seek(info.header_offset)
    raw = handle.read(30)
    if len(raw) != 30:
        raise ValueError('truncated PyTorch ZIP local header')
    fields = struct.unpack('<IHHHHHIIIHH', raw)
    if fields[0] != 0x04034B50 or fields[2] & 1 or fields[3] != zipfile.ZIP_STORED:
        raise ValueError('encrypted or compressed PyTorch storage is not supported')
    name_length, extra_length = fields[-2:]
    return info.header_offset + 30 + name_length + extra_length


def load_torch_metadata(path):
    if sys.byteorder != 'little' or not zipfile.is_zipfile(path):
        raise ValueError('only little-endian, uncompressed PyTorch ZIP checkpoints are supported')
    # Set before importing torch; none of the following constructs a model or CUDA tensor.
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    import torch
    from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode

    with zipfile.ZipFile(path) as archive, path.open('rb') as raw:
        infos = archive.infolist()
        byteorders = [item for item in infos if item.filename.endswith('/byteorder')]
        if len(byteorders) != 1 or archive.read(byteorders[0]) != b'little':
            raise ValueError('checkpoint byteorder is not explicitly little-endian')
        pickles = [item for item in infos if item.filename.endswith('/data.pkl')]
        if len(pickles) != 1 or pickles[0].file_size > 100 * 1024**2:
            raise ValueError('unexpected PyTorch checkpoint metadata')
        metadata_bytes = archive.read(pickles[0])
        storage_ranges = {}
        for info in infos:
            if '/data/' not in info.filename:
                continue
            if info.compress_type != zipfile.ZIP_STORED or info.flag_bits & 1:
                raise ValueError('compressed or encrypted storage')
            offset = zip_payload_offset(raw, info)
            if offset in storage_ranges:
                raise ValueError('duplicate ZIP storage offset')
            if offset + info.file_size > path.stat().st_size:
                raise ValueError('out-of-bounds ZIP storage')
            storage_ranges[offset] = info.file_size
    with FakeTensorMode():
        state = torch.load(path, map_location='cpu', weights_only=True, mmap=True)
    if isinstance(state, dict) and len(state) == 1:
        for wrapper in ('state_dict', 'module', 'model_state'):
            if wrapper in state:
                state = state[wrapper]
                break
    if not isinstance(state, dict) or not state:
        raise ValueError('only a nonempty flat tensor state_dict can be converted')
    dtypes = {
        torch.bool: 'BOOL', torch.uint8: 'U8', torch.int8: 'I8',
        torch.int16: 'I16', torch.float16: 'F16', torch.bfloat16: 'BF16',
        torch.int32: 'I32', torch.float32: 'F32', torch.int64: 'I64',
        torch.float64: 'F64',
    }
    tensors, seen = [], set()
    for name, tensor in state.items():
        if not isinstance(name, str) or name == '__metadata__' or not isinstance(tensor, FakeTensor):
            raise ValueError('non-tensor entry or unsupported state_dict key')
        if tensor.layout != torch.strided or not tensor.is_contiguous():
            raise ValueError('non-contiguous/sparse tensors are not supported')
        if tensor.dtype not in dtypes:
            raise ValueError('unsupported torch dtype: ' + str(tensor.dtype))
        storage = tensor.untyped_storage()
        offset = getattr(storage, '_checkpoint_offset', None)
        size = tensor.numel() * tensor.element_size()
        if offset is None:
            raise ValueError('PyTorch lacks FakeTensor checkpoint offsets; use a recent PyTorch (>=2.7)')
        # Narrow conversion intentionally rejects views/aliases instead of silently changing them.
        if tensor.storage_offset() != 0 or storage.nbytes() != size or storage_ranges.get(offset) != size:
            raise ValueError('only complete, zero-offset tensor storages are supported')
        if offset in seen:
            raise ValueError('shared storages are not supported')
        seen.add(offset)
        tensors.append(dict(name=name, dtype=dtypes[tensor.dtype], shape=list(tensor.shape),
                            nbytes=size, source_file=path.name, source_offset=offset))
    if seen != set(storage_ranges):
        raise ValueError('unused or missing PyTorch storage')
    tensors.sort(key=lambda item: item['source_offset'])
    mapping = {name: path.name for name in state}
    index = {'metadata': {'total_size': sum(t['nbytes'] for t in tensors)}, 'weight_map': mapping}
    record = source_record(path, metadata_bytes)
    return index, tensors, [record], {'format': 'pt'}, torch.__version__


def load_source(source, source_index=None):
    source = Path(source).resolve()
    if source.is_dir():
        if source_index:
            if Path(source_index).name != source_index:
                raise ValueError('unsafe source index filename')
            index_path = source / source_index
        else:
            indices = sorted(source.glob('*.safetensors.index.json'))
            if len(indices) != 1:
                raise ValueError('expected one safetensors index; select --source-index explicitly')
            index_path = indices[0]
        index = parse_json(index_path.read_text())
        mapping = index.get('weight_map')
        if not isinstance(mapping, dict) or not mapping:
            raise ValueError('empty or invalid safetensors index')
        tensors, files, common, seen = [], [], None, set()
        for name in sorted(set(mapping.values())):
            if not isinstance(name, str) or Path(name).name != name or not name.endswith('.safetensors'):
                raise ValueError('unsafe index filename')
            path = source / name
            parsed, metadata, prefix = read_header(path)
            if common is None:
                common = metadata
            elif common != metadata:
                raise ValueError('source shard metadata differs')
            for tensor in parsed:
                if tensor['name'] in seen or mapping.get(tensor['name']) != name:
                    raise ValueError('duplicate tensor or incorrect index mapping')
                seen.add(tensor['name'])
                tensors.append(tensor)
            files.append(source_record(path, prefix))
        if seen != set(mapping):
            raise ValueError('index tensor set differs from shard headers')
        identity = {'type': 'indexed_safetensors', 'path': str(source),
                    'index': index_path.name, 'index_sha256': sha_file(index_path)}
        root = source
    elif source.is_file() and source.suffix == '.safetensors':
        tensors, common, prefix = read_header(source)
        files = [source_record(source, prefix)]
        index = {'weight_map': {tensor['name']: source.name for tensor in tensors}}
        identity = {'type': 'single_safetensors', 'path': str(source)}
        root = source.parent
    elif source.is_file() and source.suffix in ('.pth', '.pt', '.bin'):
        index, tensors, files, common, version = load_torch_metadata(source)
        identity = {'type': 'torch_zip_state_dict', 'path': str(source),
                    'metadata_loader': 'weights_only=True, mmap=True, FakeTensorMode',
                    'torch_version': version}
        root = source.parent
    else:
        raise ValueError('source must be indexed/single safetensors or a supported PyTorch ZIP')
    if not tensors:
        raise ValueError('empty checkpoint')
    total = sum(tensor['nbytes'] for tensor in tensors)
    if index.get('metadata', {}).get('total_size', total) != total:
        raise ValueError('source index total_size differs')
    return root, index, tensors, files, common, identity


def output_header(tensors, metadata):
    result, offset = {'__metadata__': metadata}, 0
    for tensor in tensors:
        result[tensor['name']] = dict(dtype=tensor['dtype'], shape=tensor['shape'],
                                      data_offsets=[offset, offset + tensor['nbytes']])
        offset += tensor['nbytes']
    raw = json.dumps(result, ensure_ascii=False, separators=(',', ':')).encode()
    raw += b' ' * ((-len(raw)) % 8)
    return struct.pack('<Q', len(raw)) + raw


def plan_shards(tensors, metadata, max_bytes):
    groups, current, size = [], [], 0
    for tensor in tensors:
        candidate = current + [tensor]
        if len(output_header(candidate, metadata)) + size + tensor['nbytes'] > max_bytes:
            if not current:
                raise ValueError('single tensor exceeds shard limit; a tensor cannot be split')
            groups.append(current)
            current, size = [tensor], tensor['nbytes']
            if len(output_header(current, metadata)) + size > max_bytes:
                raise ValueError('single tensor exceeds shard limit')
        else:
            current, size = candidate, size + tensor['nbytes']
    if current:
        groups.append(current)
    return groups


def write_atomic_json(path, value):
    path = Path(path)
    temporary = path.with_name(path.name + '.partial')
    temporary.write_text(json.dumps(value, ensure_ascii=False, indent=2) + '\n')
    temporary.replace(path)


def stable_sources(root, files):
    for info in files:
        stat = (root / info['name']).stat()
        if (stat.st_size, stat.st_mtime_ns) != (info['bytes'], info['mtime_ns']):
            raise RuntimeError('source checkpoint changed during sharding')


def verify_outputs(out, groups, metadata, tensor_hashes, records, max_bytes):
    if len(records) != len(groups):
        raise ValueError('shard count differs')
    mapping, verified = {}, set()
    for group, record in zip(groups, records):
        path = out / record['name']
        parsed, actual_metadata, prefix = read_header(path)
        expected = {tensor['name']: tensor for tensor in group}
        if actual_metadata != metadata or path.stat().st_size > max_bytes:
            raise ValueError('output metadata/size differs')
        if {tensor['name'] for tensor in parsed} != set(expected):
            raise ValueError('output tensor names differ')
        full = hashlib.sha256(prefix)
        with path.open('rb') as handle:
            for tensor in parsed:
                original = expected[tensor['name']]
                if any(tensor[key] != original[key] for key in ('dtype', 'shape', 'nbytes')):
                    raise ValueError('tensor dtype/shape differs')
                handle.seek(tensor['source_offset'])
                digest, left = hashlib.sha256(), tensor['nbytes']
                while left:
                    block = handle.read(min(BLOCK, left))
                    if not block:
                        raise ValueError('truncated output tensor')
                    digest.update(block)
                    full.update(block)
                    left -= len(block)
                if digest.hexdigest() != tensor_hashes.get(tensor['name']):
                    raise ValueError('tensor bytes differ: ' + tensor['name'])
                verified.add(tensor['name'])
                mapping[tensor['name']] = path.name
        if full.hexdigest() != record['sha256'] or path.stat().st_size != record['bytes']:
            raise ValueError('output shard hash/size differs')
        print(f'VERIFIED {out.name}/{path.name} ({len(group)} tensors)', flush=True)
    if verified != set(tensor_hashes):
        raise ValueError('tensor verification coverage differs')
    return mapping


def model_card(spec):
    return f'''---
tags:
- safetensors
- lossless-resharding
---
# {spec['kind']}

{spec['shard_count']} standard safetensors shards, each at most
{spec['max_shard_bytes']:,} bytes **including the header**. Native tensor names,
shapes, dtypes and raw value bytes are unchanged. No model construction,
quantization, inference, training, GPU allocation or upload was performed.

Source/base: `{spec.get('base_model') or 'see parent bundle README'}`.
Raw tensor data: {spec['raw_tensor_bytes']:,} bytes, {spec['tensor_count']:,} tensors.
`{spec['prefix']}.safetensors.index.json` maps all tensor names to their shards.
`reshard_manifest.json` records source/file hashes and independently verified
per-tensor SHA256. For a native DiffSynth component, pass the **list of shard
paths** to `ModelConfig(path=[...])`; this export is not an architecture/key
conversion into a Transformers or diffusers model.

PyTorch ZIP conversion, when applicable, uses restricted `weights_only=True`
FakeTensor metadata loading and streams original complete storages without
executing unrestricted pickle. Obtain/review the original model license before
redistribution. This directory is prepared locally, not uploaded.
'''


def reshard(source, out, max_bytes=5_000_000_000, kind='checkpoint', resume=False,
            prefix='model', source_index=None, base_model=None):
    if not prefix or Path(prefix).name != prefix or not all(c.isalnum() or c in '._-' for c in prefix):
        raise ValueError('unsafe output prefix')
    if type(max_bytes) is not int or max_bytes <= 0:
        raise ValueError('positive integer shard size required')
    root, index, tensors, files, metadata, identity = load_source(source, source_index)
    out = Path(out).resolve()
    source_path = Path(source).resolve()
    if out == root or out in root.parents or root in out.parents:
        raise ValueError('output must not overlap the original checkpoint directory')
    groups = plan_shards(tensors, metadata, max_bytes)
    spec = dict(source=identity, source_files=files, max_shard_bytes=max_bytes,
                tensor_count=len(tensors), raw_tensor_bytes=sum(t['nbytes'] for t in tensors),
                shard_count=len(groups), prefix=prefix, kind=kind, base_model=base_model)
    final, state_path = out / 'reshard_manifest.json', out / '.reshard_state.json'
    index_name = prefix + '.safetensors.index.json'
    if final.is_file():
        old = parse_json(final.read_text())
        if old['spec'] != spec:
            raise ValueError('existing completed export has different inputs')
        verify_outputs(out, groups, metadata, old['tensor_sha256'], old['output_shards'], max_bytes)
        if sha_file(out / index_name) != old['output_index_sha256']:
            raise ValueError('completed output index changed')
        for sidecar in old['small_files']:
            if sha_file(out / sidecar['name']) != sidecar['sha256']:
                raise ValueError('completed output sidecar changed')
        for info in files:
            if sha_file(root / info['name']) != old['source_file_sha256'][info['name']]:
                raise ValueError('source bytes changed after completed export')
        print('ALREADY COMPLETE + VERIFIED', out, flush=True)
        return old
    state = dict(spec=spec, completed=[], tensor_sha256={})
    if out.exists() and any(out.iterdir()):
        if not resume or not state_path.is_file():
            raise ValueError('output is not empty; use fresh output or matching --resume')
        state = parse_json(state_path.read_text())
        if state['spec'] != spec:
            raise ValueError('resume inputs differ')
    out.mkdir(parents=True, exist_ok=True)
    write_atomic_json(state_path, state)
    records, tensor_hashes = state['completed'], state['tensor_sha256']
    if len(records) > len(groups):
        raise ValueError('too many completed shards in resume state')
    print(json.dumps({'export': str(out), 'source_type': identity['type'],
                      'tensors': len(tensors), 'raw_GB': round(spec['raw_tensor_bytes'] / 1e9, 3),
                      'shards': len(groups), 'GPU_used': False}, indent=2), flush=True)
    source_hashes = {info['name']: sha_file(root / info['name']) for info in files}
    stable_sources(root, files)
    start, last_report = time.monotonic(), 0
    done = sum(t['nbytes'] for group in groups[:len(records)] for t in group)
    with ExitStack() as stack:
        handles = {info['name']: stack.enter_context((root / info['name']).open('rb')) for info in files}
        for ordinal, group in enumerate(groups, 1):
            name = f'{prefix}-{ordinal:05d}-of-{len(groups):05d}.safetensors'
            target = out / name
            if ordinal <= len(records):
                record = records[ordinal - 1]
                if record['name'] != name or target.stat().st_size != record['bytes'] or sha_file(target) != record['sha256']:
                    raise ValueError('completed shard differs on resume')
                continue
            header = output_header(group, metadata)
            full, temporary = hashlib.sha256(header), target.with_name(name + '.partial')
            with temporary.open('wb') as handle:
                handle.write(header)
                for tensor in group:
                    source_handle = handles[tensor['source_file']]
                    source_handle.seek(tensor['source_offset'])
                    digest, left = hashlib.sha256(), tensor['nbytes']
                    while left:
                        block = source_handle.read(min(BLOCK, left))
                        if not block:
                            raise ValueError('truncated source tensor')
                        handle.write(block)
                        digest.update(block)
                        full.update(block)
                        left -= len(block)
                        done += len(block)
                        if done - last_report >= 1024**3:
                            print(f'COPY {out.name}: {done / 1e9:.1f}/{spec["raw_tensor_bytes"] / 1e9:.1f} GB; '
                                  f'shard {ordinal}/{len(groups)}; {time.monotonic() - start:.0f}s', flush=True)
                            last_report = done
                    tensor_hashes[tensor['name']] = digest.hexdigest()
                handle.flush()
            size = temporary.stat().st_size
            if size != len(header) + sum(t['nbytes'] for t in group) or size > max_bytes:
                raise ValueError('incorrect output shard size')
            stable_sources(root, files)
            temporary.replace(target)
            records.append(dict(name=name, bytes=size, sha256=full.hexdigest(), tensor_count=len(group)))
            write_atomic_json(state_path, state)
            print('WROTE', target, f'{size / 1e9:.3f} GB', flush=True)
    mapping = verify_outputs(out, groups, metadata, tensor_hashes, records, max_bytes)
    if set(mapping) != set(index['weight_map']):
        raise ValueError('output index tensor set differs')
    result_index = {key: value for key, value in index.items() if key not in ('weight_map', 'metadata')}
    result_index['metadata'] = {**index.get('metadata', {}), 'total_size': spec['raw_tensor_bytes']}
    result_index['weight_map'] = {name: mapping[name] for name in index['weight_map']}
    sidecars = []
    # Native single-file components do not inherit unrelated siblings' configs.
    if source_path.is_dir():
        for name in sorted(SIDECARS):
            path = root / name
            if not path.is_file():
                continue
            if path.stat().st_size > 100 * 1024**2:
                raise ValueError('unexpectedly large sidecar')
            digest = sha_file(path)
            shutil.copy2(path, out / name)
            if sha_file(out / name) != digest:
                raise ValueError('sidecar copy differs')
            sidecars.append(dict(name=name, bytes=path.stat().st_size, sha256=digest))
    stable_sources(root, files)
    for info in files:
        if sha_file(root / info['name']) != source_hashes[info['name']]:
            raise ValueError('source bytes changed during sharding')
    write_atomic_json(out / index_name, result_index)
    (out / 'README.md').write_text(model_card(spec))
    manifest = dict(spec=spec, created_at_UTC=now(), source_file_sha256=source_hashes,
                    output_shards=records, small_files=sidecars, tensor_sha256=tensor_hashes,
                    output_index_sha256=sha_file(out / index_name),
                    validation=dict(all_tensor_names_shapes_dtypes_equal=True,
                                    all_tensor_bytes_SHA256_equal=True,
                                    all_output_shards_read_back=True,
                                    source_files_unchanged=True, max_file_limit_respected=True),
                    GPU_used=False, model_loaded=False, training_run=False, upload_run=False)
    write_atomic_json(final, manifest)
    state_path.unlink()
    print('COMPLETE', out, flush=True)
    return manifest


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--source', type=Path, required=True)
    parser.add_argument('--output-dir', type=Path, required=True)
    parser.add_argument('--source-index')
    parser.add_argument('--prefix', default='model')
    parser.add_argument('--max-shard-gb', type=float, default=5)
    parser.add_argument('--kind', default='checkpoint')
    parser.add_argument('--base-model')
    parser.add_argument('--resume', action='store_true')
    args = parser.parse_args()
    if not math.isfinite(args.max_shard_gb) or args.max_shard_gb <= 0:
        parser.error('--max-shard-gb must be finite and positive')
    result = reshard(args.source, args.output_dir, int(args.max_shard_gb * 1e9),
                     args.kind, args.resume, args.prefix, args.source_index, args.base_model)
    print(json.dumps({'output_dir': str(args.output_dir.resolve()),
                      'shards': len(result['output_shards']),
                      'tensors': result['spec']['tensor_count'],
                      'tensor_bytes_verified': True, 'GPU_used': False, 'upload_run': False}, indent=2))


if __name__ == '__main__':
    main()