"""Atomic lossless checkpoints without a simultaneous raw save on disk."""
import io
import hashlib
from pathlib import Path
import shutil
import subprocess
import torch

ZSTD=shutil.which('zstd')


class HashWriter:
    def __init__(self,stream):
        self.stream=stream;self.hash=hashlib.sha256();self.bytes=0
    def write(self,data):
        self.hash.update(data);self.bytes+=len(data)
        return self.stream.write(data)
    def flush(self):self.stream.flush()
    def tell(self):return self.bytes


def save_compressed(payload,path):
    path=Path(path)
    if not ZSTD:raise RuntimeError('zstd binary required')
    tmp=path.with_suffix(path.suffix+'.tmp')
    if path.exists() or tmp.exists():raise FileExistsError(path)
    proc=subprocess.Popen([ZSTD,'-q','-T2','-3','-o',str(tmp)],stdin=subprocess.PIPE)
    writer=HashWriter(proc.stdin)
    try:
        torch.save(payload,writer);writer.flush();proc.stdin.close()
        if proc.wait()!=0:raise RuntimeError('compression failed')
    except BaseException:
        proc.kill();proc.wait()
        raise
    # Verify serialized bytes exactly before declaring the atomic save complete.
    check=hashlib.sha256();size=0
    proc=subprocess.Popen([ZSTD,'-q','-d','-c',str(tmp)],stdout=subprocess.PIPE)
    try:
        while block:=proc.stdout.read(8<<20):check.update(block);size+=len(block)
        if proc.wait()!=0:raise RuntimeError('decompression check failed')
    finally:proc.stdout.close()
    if check.hexdigest()!=writer.hash.hexdigest() or size!=writer.bytes:
        raise RuntimeError('checkpoint round-trip mismatch')
    archive=hashlib.sha256()
    with tmp.open('rb') as stream:
        while block:=stream.read(8<<20):archive.update(block)
    tmp.replace(path)
    return {'path':str(path),'raw_sha256':check.hexdigest(),'raw_bytes':size,
            'archive_sha256':archive.hexdigest(),'archive_bytes':path.stat().st_size,
            'codec':'zstd level3 lossless','round_trip':'PASS'}


def load_compressed(path,metadata=None):
    """Load CPU tensors from an in-memory decompression buffer."""
    path=Path(path)
    if not ZSTD:raise RuntimeError('zstd binary required')
    if metadata:
        h=hashlib.sha256()
        with path.open('rb') as stream:
            while block:=stream.read(8<<20):h.update(block)
        if h.hexdigest()!=metadata['archive_sha256']:raise ValueError('archive hash mismatch')
    # Keep the decompressed serialization in RAM to avoid a second 9GB disk
    # checkpoint during restore. torch.load materializes independent CPU tensors.
    proc=subprocess.Popen([ZSTD,'-q','-d','-c',str(path)],stdout=subprocess.PIPE)
    raw=hashlib.sha256();size=0
    with io.BytesIO() as buffer:
        try:
            while block:=proc.stdout.read(8<<20):
                raw.update(block);size+=len(block);buffer.write(block)
            if proc.wait()!=0:raise RuntimeError('checkpoint decompression failed')
        finally:proc.stdout.close()
        if metadata and (raw.hexdigest()!=metadata['raw_sha256'] or size!=metadata['raw_bytes']):
            raise ValueError('serialized hash mismatch')
        buffer.seek(0)
        return torch.load(buffer,map_location='cpu',weights_only=True)
