import argparse
import base64
import hashlib
import json
from pathlib import Path
import shutil
import struct
import numpy as np


def file_hash(path):
    with path.open('rb') as stream:
        return hashlib.file_digest(stream, 'sha256').hexdigest()


def apply_patch(original, manifest_path, patch_path, output):
    original, manifest_path, patch_path, output = map(Path, (original, manifest_path, patch_path, output))
    if output.exists() or output.resolve() == original.resolve():
        raise ValueError('Output must be a new file; original weights are never overwritten')
    manifest = json.loads(manifest_path.read_text())
    if file_hash(original) != manifest['original_sha256'] or file_hash(patch_path) != manifest['patch_sha256']:
        raise ValueError('Original checkpoint or patch checksum differs')
    with original.open('rb') as stream:
        original_length = struct.unpack('<Q', stream.read(8))[0]
        original_header = json.loads(stream.read(original_length))
    target_header_bytes = base64.b64decode(manifest['target_header_base64'], validate=True)
    target_header = json.loads(target_header_bytes)
    if {k: v for k, v in original_header.items() if k != '__metadata__'} != {
            k: v for k, v in target_header.items() if k != '__metadata__'}:
        raise ValueError('Tensor layouts differ')
    output.parent.mkdir(parents=True, exist_ok=True)
    temporary = output.with_name(output.name + '.partial')
    if temporary.exists():
        raise ValueError('Partial output already exists; inspect it before retrying')
    with original.open('rb') as source, temporary.open('xb') as target:
        target.write(struct.pack('<Q', len(target_header_bytes)))
        target.write(target_header_bytes)
        source.seek(8 + original_length)
        shutil.copyfileobj(source, target, length=8 * 1024 * 1024)
    changed = 0
    with np.load(patch_path, allow_pickle=False) as arrays:
        for entry in manifest['tensors']:
            name = entry['name']
            if name not in {f'model.language_model.layers.{layer}.mlp.down_proj.weight' for layer in (15, 19)}:
                raise ValueError('Unexpected release tensor')
            info = target_header[name]
            if info['dtype'] != 'BF16' or info['shape'] != entry['shape']:
                raise ValueError('Tensor shape/dtype differs')
            indices, bits = arrays[entry['indices_key']], arrays[entry['bits_key']]
            size = int(np.prod(info['shape']))
            if indices.ndim != 1 or indices.dtype != np.dtype('uint64') or bits.dtype != np.dtype('uint16'):
                raise ValueError('Invalid sparse array format')
            if bits.shape != indices.shape or len(indices) != entry['changed_bitpatterns']:
                raise ValueError('Sparse array lengths differ')
            if len(indices) and (indices[-1] >= size or np.any(indices[1:] <= indices[:-1])):
                raise ValueError('Sparse indices must be sorted, unique, and in range')
            values = np.memmap(temporary, mode='r+', dtype='<u2',
                               offset=8 + len(target_header_bytes) + info['data_offsets'][0], shape=(size,))
            values[indices] = bits
            values.flush()
            del values
            changed += len(indices)
    if changed != manifest['changed_bitpatterns'] or file_hash(temporary) != manifest['target_sha256']:
        raise ValueError('Reconstructed weight file differs from the evaluated candidate')
    temporary.rename(output)
    return {'passed': True, 'changed_bitpatterns': changed, 'target_sha256': manifest['target_sha256'],
            'byte_identical_to_evaluated_weights': True, 'model_forward_executed': False}


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--original', required=True)
    parser.add_argument('--manifest', required=True)
    parser.add_argument('--patch', required=True)
    parser.add_argument('--output', required=True)
    args = parser.parse_args()
    print(json.dumps(apply_patch(args.original, args.manifest, args.patch, args.output)))


if __name__ == '__main__':
    main()