"""Validate the exact exported student and invoke its pinned four-call runner."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import subprocess
import sys

HERE = Path(__file__).resolve().parent


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


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--model-dir', '--model-id', default='MiniMaxAI/MiniMax-H3')
    parser.add_argument('--diffusers-src', type=Path)
    parser.add_argument('--prompts', type=Path, default=HERE/'prompts.example.json')
    parser.add_argument('--expected-prompt-count', type=int, default=4)
    parser.add_argument('--output-dir', type=Path, default=Path('outputs'))
    args, rest = parser.parse_known_args()
    for argument in rest:
        if argument.split('=')[0] in ('--checkpoint', '--revision', '--num-inference-steps'):
            parser.error('Checkpoint, revision and four-call sampler are fixed by this release')
    provenance = json.loads((HERE/'provenance.json').read_text())
    checkpoint = HERE/provenance['export']['file']
    if checkpoint.is_symlink() or sha(checkpoint) != provenance['export']['sha256']:
        raise RuntimeError('Student checkpoint checksum mismatch')
    for name, digest in json.loads((HERE/'source_hashes.json').read_text()).items():
        if sha(HERE/name) != digest:
            raise RuntimeError('Source snapshot checksum mismatch: '+name)
    paths = [str(HERE/'source_snapshot')]
    if args.diffusers_src is not None:
        if not (args.diffusers_src/'diffusers').is_dir():
            parser.error('--diffusers-src must contain the diffusers package')
        paths.append(str(args.diffusers_src.resolve()))
    env = os.environ.copy()
    env['PYTHONPATH'] = os.pathsep.join(paths + ([env['PYTHONPATH']] if env.get('PYTHONPATH') else []))
    env.update(PYTHONDONTWRITEBYTECODE='1', PYTHONUNBUFFERED='1', TOKENIZERS_PARALLELISM='false')
    env.setdefault('PYTORCH_CUDA_ALLOC_CONF', 'expandable_segments:True')
    return subprocess.run([sys.executable, '-m', 'minimax_h3.infer_student_4step',
        '--checkpoint', str(checkpoint), '--model-id', args.model_dir,
        '--revision', provenance['inference']['base_revision'],
        '--prompts', str(args.prompts), '--expected-prompt-count', str(args.expected_prompt_count),
        '--output-dir', str(args.output_dir), *rest], env=env).returncode


if __name__ == '__main__':
    raise SystemExit(main())