"""Load an exact completed-5B release, with hash and endpoint verification.

This is a ResearchLM checkpoint, not a Transformers AutoModel checkpoint.
Review the bundled research_code and install its pinned CUDA dependencies first.
"""
from __future__ import annotations

import hashlib
import json
from pathlib import Path, PurePosixPath
import sys


def sha256(path):
    h = hashlib.sha256()
    with open(path, 'rb') as stream:
        for part in iter(lambda: stream.read(8 << 20), b''):
            h.update(part)
    return h.hexdigest()


def verify_release(folder):
    folder = Path(folder).resolve()
    manifest = json.loads((folder / 'release_manifest.json').read_text())
    for name, expected in manifest['files'].items():
        relative = PurePosixPath(name)
        # HF snapshots normally symlink into a sibling content-addressed blob
        # directory. Validate the manifest's relative path, then hash its bytes.
        if relative.is_absolute() or '..' in relative.parts:
            raise ValueError(f'Invalid release path: {name}')
        path = folder / relative
        if not path.is_file():
            raise ValueError(f'Invalid release file: {name}')
        if path.stat().st_size != expected['bytes'] or sha256(path) != expected['sha256']:
            raise ValueError(f'Release file differs: {name}')
    provenance = json.loads((folder / 'provenance.json').read_text())
    if provenance['step'] != 9537 or provenance['tokens'] != 5000134656:
        raise ValueError('A completed 5B endpoint is required')
    if provenance['corpus_sha256'] != '1786fc7cf798784edf2e0ee9400456ca93add1cbc5688cd8e83865cb230bc475':
        raise ValueError('Unexpected training corpus')
    return provenance


def load_release(folder, *, device='cuda', dtype=None):
    import torch
    from safetensors.torch import load_file

    folder = Path(folder).resolve()
    provenance = verify_release(folder)
    code = folder / 'research_code'
    existing = sys.modules.get('pretraining_5b.model')
    if existing is not None and Path(existing.__file__).resolve() != (code / 'pretraining_5b/model.py').resolve():
        raise RuntimeError('Use a fresh process: a different pretraining_5b.model is already imported')
    sys.path.insert(0, str(code))
    from pretraining_5b.model import ModelConfig, ResearchLM

    config = json.loads((folder / 'research_config.json').read_text())
    weights = load_file(str(folder / 'model.safetensors'), device='cpu')
    aliases = json.loads((folder / 'tied_weights.json').read_text())
    for alias, source in aliases.items():
        weights[alias] = weights[source]
    model = ResearchLM(ModelConfig(**config))
    model.load_state_dict(weights, strict=True)
    if model.parameter_count() != provenance['parameter_count']:
        raise ValueError('Model parameter count differs from the admitted release')
    model = model.to(device=device, dtype=dtype or torch.bfloat16).eval()
    return model, provenance