"""Migration validation: real CPU deploy loading, with historical weight paths blocked."""
import argparse
import json
import logging
import os
from pathlib import Path
import sys
import time

from checkpoint_archive import Archive, atomic_json, now, sha256_file

parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('run')
parser.add_argument('--output', required=True)
args = parser.parse_args()
root = Path('/mnt/data/limingleyang/checkpoints')
archive = Archive(root)
entry = archive.resolve(args.run)
output = archive.materialize(args.run, args.output)
repo = Path('/mnt/data/limingleyang/OpenWAM-private')
sys.path.insert(0, str(repo))
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TRANSFORMERS_OFFLINE'] = '1'
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
plan = json.loads((root / 'migration/plan.json').read_text())
blocked = [str(Path(p['source']).parent) for p in plan['policies'] + plan['svae']]
blocked += [p['source'] for p in plan['extra_models']]
blocked += ['/mnt/data/limingleyang/model', '/mnt/data/limingleyang/models']


def audit(event, values):
    if event == 'open' and isinstance(values[0], (str, bytes)):
        path = os.path.abspath(os.fsdecode(values[0]))
        if any(path == b or path.startswith(b + '/') for b in blocked):
            raise PermissionError(f'Archive self-containment test blocked historical model path: {path}')


sys.addaudithook(audit)
logging.basicConfig(level=logging.INFO)
started = time.monotonic()
try:
    import torch
    torch.set_num_threads(8)
    from openwam.deploy.model_loader import load_from_checkpoint_dir
    cfg, model = load_from_checkpoint_dir(str(output), device='cpu', ckpt_name=entry['source']['filename'])
    parameters = list(model.parameters())
    if any(p.is_meta or p.device.type != 'cpu' for p in parameters):
        raise ValueError('Model has unresolved meta or non-CPU parameters')
    report = {'status': 'passed', 'time': now(), 'run': entry['id'], 'device': 'cpu',
              'parameters': sum(p.numel() for p in parameters), 'elapsed_seconds': time.monotonic() - started,
              'materialized': str(output), 'index_sha256': entry['index']['sha256'],
              'repo_commit': 'cc4bf0c6bd3e016723c0b4050cd93b7d5c827ef6',
              'historical_weight_paths_blocked': blocked, 'network_model_downloads': 'offline',
              'forward_or_closed_loop_eval': False}
except BaseException as exc:
    report = {'status': 'failed', 'time': now(), 'run': entry['id'], 'error': repr(exc),
              'elapsed_seconds': time.monotonic() - started}
    atomic_json(root / 'migration' / (entry['run_id'] + '-cpu-load.json'), report)
    raise
atomic_json(root / 'migration' / (entry['run_id'] + '-cpu-load.json'), report)
print(json.dumps(report, indent=2))