#!/usr/bin/env python
"""Split memory.bin into <5GB chunks, uploading each to HF Hub, then delete."""
import sys, os, shutil
sys.path.insert(0, '.')

from huggingface_hub import HfApi
api = HfApi()

src = 'palimpseste-max/palimpseste_memory.bin'
sz = os.path.getsize(src)
chunk_size = 4_500_000_000
n_chunks = (sz + chunk_size - 1) // chunk_size
print(f'Source: {sz/1024**3:.2f} GB -> {n_chunks} chunks', flush=True)

# Phase 1: Write chunk 0, upload it, delete it. Repeat.
with open(src, 'rb') as f:
    for i in range(n_chunks):
        chunk_path = f'palimpseste-max/palimpseste_memory.bin.part{i}'
        remaining = min(chunk_size, sz - i * chunk_size)

        print(f'Writing chunk {i} ({remaining/1024**3:.2f} GB)...', flush=True)
        written = 0
        with open(chunk_path, 'wb') as out:
            while written < remaining:
                block = f.read(min(8 * 1024 * 1024, remaining - written))
                if not block:
                    break
                out.write(block)
                written += len(block)

        csz = os.path.getsize(chunk_path)
        print(f'  chunk {i}: {csz/1024**3:.2f} GB', flush=True)

        # Upload to HF Hub
        repo_path = f'palimpseste-max/palimpseste_memory.bin.part{i}'
        print(f'  uploading chunk {i} to HF Hub...', flush=True)
        api.upload_file(
            path_or_fileobj=chunk_path,
            path_in_repo=repo_path,
            repo_id='thefinalboss/palimpseste-max',
            repo_type='model',
        )
        print(f'  chunk {i} uploaded!', flush=True)

        # Delete local chunk to free space
        os.remove(chunk_path)
        print(f'  chunk {i} deleted locally', flush=True)

# Delete the original
os.remove(src)
print(f'Original deleted. Split+upload COMPLETE!', flush=True)