"""Commit run artifacts atomically and verify every changed file by SHA256."""
import hashlib, json, os, time
from pathlib import Path
from huggingface_hub import HfApi, CommitOperationAdd, hf_hub_download

REPO = 'User-2468/mini-unet-colorizer'
PREFIX = 'experiments/round6-20260928'

class DurableRun:
    def __init__(self, root):
        if not os.environ.get('HF_TOKEN'):
            raise RuntimeError('A write token is required before starting')
        self.api = HfApi(token=os.environ['HF_TOKEN'])
        if self.api.whoami()['name'] != 'User-2468':
            raise RuntimeError('Unexpected account')
        self.root = Path(root)
        self.saved = {}

    def sync(self, reason):
        files = [p for p in sorted(self.root.rglob('*')) if p.is_file()]
        hashes = {p.relative_to(self.root).as_posix(): hashlib.sha256(p.read_bytes()).hexdigest() for p in files}
        changed = [p for p in files if self.saved.get(p.relative_to(self.root).as_posix()) != hashes[p.relative_to(self.root).as_posix()]]
        if not changed:
            return
        for attempt in range(3):
            try:
                head = self.api.model_info(REPO, revision='main').sha
                operations = [CommitOperationAdd(path_in_repo=PREFIX+'/'+p.relative_to(self.root).as_posix(), path_or_fileobj=str(p)) for p in changed]
                operations.append(CommitOperationAdd(path_in_repo=PREFIX+'/SHA256SUMS.json', path_or_fileobj=json.dumps(hashes,sort_keys=True,indent=2).encode()))
                commit = self.api.create_commit(repo_id=REPO, revision='main', parent_commit=head, operations=operations, commit_message='Colorizer round6: '+reason)
                for p in changed:
                    name = p.relative_to(self.root).as_posix()
                    downloaded = hf_hub_download(REPO, PREFIX+'/'+name, revision=commit.oid, force_download=True, token=os.environ['HF_TOKEN'])
                    if hashlib.sha256(Path(downloaded).read_bytes()).hexdigest() != hashes[name]:
                        raise RuntimeError('Remote checksum mismatch: '+name)
                self.saved = hashes
                print('PERSISTED',reason,commit.oid,'verified_files',len(changed),flush=True)
                return
            except Exception:
                if attempt == 2:
                    raise
                time.sleep(2**attempt)