"""Full, pinned Socialtrait checkpoint with one on-disk copy of each LFS object.

Install huggingface_hub and git-lfs, authenticate with `hf auth login`, then:
  python clone_checkpoint.py /path/with/at-least-31-GiB-free/flux3-klein
Do not edit weight files in place: the checkout and LFS store use hardlinks.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess

from huggingface_hub import HfApi, get_token, hf_hub_download

from model_config import UPSTREAM as REPO, REVISION
OVERLAY = ('model_config.py', 'flux3_engine.py', 'handler.py', 'web.py', 'modal_app.py',
           'requirements.txt', 'prompts.json', 'README.md', 'DEPLOYMENT.txt',
           'verify_checkpoint.py', 'validate_deployment.py')


def digest(path):
    sha = hashlib.sha256()
    with path.open('rb') as source:
        for chunk in iter(lambda: source.read(8 * 1024 * 1024), b''):
            sha.update(chunk)
    return sha.hexdigest()


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('destination', type=Path)
    parser.add_argument('--reserve-gib', type=float, default=5,
                        help='Free disk space to retain after download (default: 5 GiB).')
    args = parser.parse_args()
    destination = args.destination.expanduser().resolve()
    token = get_token()
    if not token:
        raise SystemExit('Authenticate first with hf auth login; do not put a token in this command.')
    info = HfApi(token=token).model_info(REPO, revision=REVISION, files_metadata=True)
    files = list(info.siblings)
    # A prior partial clone is resumable only if our identity marker matches.
    identity = {'repo': REPO, 'revision': REVISION}
    marker = destination / '.git' / 'full-checkpoint-identity.json'
    if destination.exists() and any(destination.iterdir()):
        if not marker.exists() or json.loads(marker.read_text()) != identity:
            raise SystemExit('Destination is not a matching checkpoint clone. Choose an empty directory.')
    needed = 0
    for item in files:
        if item.lfs:
            sha = item.lfs.sha256
            obj = destination / '.git/lfs/objects' / sha[:2] / sha[2:4] / sha
            if not obj.exists() or obj.stat().st_size != item.size:
                needed += item.size
    parent = destination
    while not parent.exists():
        parent = parent.parent
    free = shutil.disk_usage(parent).free
    reserve = max(args.reserve_gib, 0) * 2**30
    if free < needed + reserve:
        raise SystemExit(f'Need {needed/2**30:.2f} GiB plus {reserve/2**30:.1f} GiB reserve; '
                         f'only {free/2**30:.2f} GiB free. Select a larger disk.')
    env = os.environ.copy()
    env.update(GIT_LFS_SKIP_SMUDGE='1', GIT_TERMINAL_PROMPT='0',
               GIT_CONFIG_COUNT='1', GIT_CONFIG_KEY_0='http.https://huggingface.co/.extraheader',
               GIT_CONFIG_VALUE_0=f'Authorization: Bearer {token}')
    if not marker.exists():
        destination.parent.mkdir(parents=True, exist_ok=True)
        subprocess.run(['git', 'clone', '--no-checkout', f'https://huggingface.co/{REPO}',
                        str(destination)], env=env, check=True)
        subprocess.run(['git', '-C', str(destination), 'checkout', REVISION], env=env, check=True)
        marker.write_text(json.dumps(identity))
    staging = destination / '.git' / 'hf-download'
    inventory = []
    for index, item in enumerate(files, 1):
        target = destination / item.rfilename
        if item.lfs:
            sha = item.lfs.sha256
            obj = destination / '.git/lfs/objects' / sha[:2] / sha[2:4] / sha
            obj.parent.mkdir(parents=True, exist_ok=True)
            if not obj.exists():
                download = Path(hf_hub_download(REPO, item.rfilename, revision=REVISION,
                                               token=token, local_dir=staging))
                if digest(download) != sha:
                    raise RuntimeError(f'Checksum mismatch: {item.rfilename}')
                os.replace(download, obj)
            elif obj.stat().st_size != item.size or digest(obj) != sha:
                raise RuntimeError(f'Existing LFS object failed verification: {item.rfilename}')
            target.parent.mkdir(parents=True, exist_ok=True)
            if target.exists():
                target.unlink()
            os.link(obj, target)
        if not target.is_file() or target.stat().st_size != item.size:
            raise RuntimeError(f'Missing or incomplete file: {item.rfilename}')
        inventory.append({'path': item.rfilename, 'bytes': item.size,
                          'sha256': item.lfs.sha256 if item.lfs else None})
        print(f'[{index}/{len(files)}] Verified {item.rfilename}', flush=True)
    (destination / 'checkpoint_inventory.json').write_text(json.dumps(
        {**identity, 'files': inventory}, indent=2) + '\n')
    source = Path(__file__).resolve().parent
    for name in OVERLAY:
        shutil.copy2(source / name, destination / name)
    with (destination / '.git/info/exclude').open('a') as exclude:
        exclude.write('\n.env*\n__pycache__/\nbenchmark_outputs/\n')
    print(f'Complete checkpoint and deployment files: {destination}')
    print('Socialtrait BF16/Q8 weights, tokenizers, configurations, licenses and notices retained.')
    print('No repository has been uploaded or made public.')


if __name__ == '__main__':
    main()