"""One-time migration driver. Only copies/imports/verifies; NEVER deletes sources."""
import concurrent.futures
import json
from pathlib import Path
import re
import sys

from checkpoint_archive import Archive, atomic_json

BASE = Path('/mnt/data/limingleyang')
ROOT = BASE / 'checkpoints'
PLAN = json.loads((ROOT / 'migration/plan.json').read_text())


def add_models():
    archive = Archive(ROOT)
    archive.init()
    dependencies = {}
    for row in PLAN['svae']:
        source = Path(row['source'])
        origin = min(int(x['id'][3:]) for x in row['consumers'])
        key = f'svae/run{origin:02d}'
        name = Path(row['destination']).name
        files = [p.name for p in source.parent.iterdir() if p.is_file() and p.name in ['svae.pt', 'config.yaml', 'complete.json']]
        extras = {}
        roots = [source.parent.parent]
        if source.parent.parent.name == 'svae_train':
            roots.append(source.parent.parent.parent)
        for root in list(roots):
            roots.extend(root / name for name in ['features', 'frames'] if (root / name).is_dir())
        for root in roots:
            for f in root.iterdir():
                if f.is_file() and (f.suffix in ['.json', '.yaml', '.yml', '.txt'] or 'stats' in f.name and f.suffix in ['.pt', '.npy']):
                    extras['collection/' + f.relative_to(BASE).as_posix()] = str(f)
        archive.add_model(key=key, name=name, source=source.parent, files=files, extra_assets=extras,
                          metadata={'consumers': row['consumers'], 'original_checkpoint': str(source),
                                    'note': 'Original independent model precision retained; policy bf16 tensors are separate exact copies.'})
        dependencies[str(source)] = key
        print(json.dumps({'model_imported': key, 'metadata_files': len(extras)}), flush=True)
    row = PLAN['extra_models'][0]
    source = Path(row['source'])
    files = [p.name for p in source.iterdir() if p.is_file() and
             (p.name in ['last.pt', 'encoder.pt', 'latent_stats_libero.pt'] or p.suffix in ['.json', '.yaml', '.yml'])]
    archive.add_model(key='pretrain/run19-encoder', name=Path(row['destination']).name,
                      source=source, files=files, metadata={'final_training_step': 1000, 'kind': 'VideoRAE encoder fine-tuning'})
    return dependencies


def import_row(pair):
    row, svae_keys = pair
    p = Path(row['source'])
    actual = p.stat()
    if actual.st_size != row['source_size'] or abs(actual.st_mtime - row['source_mtime']) > 0.001:
        raise ValueError(f'Source changed since approved inventory: {p}')
    latest = max(int(re.fullmatch(r'checkpoint_step_(\d+)\.safetensors', f.name)[1]) for f in p.parent.glob('checkpoint_step_*.safetensors'))
    if latest != row['step']:
        raise ValueError(f'Final checkpoint selection changed: {p.parent}')
    deps = {}
    for svae in PLAN['svae']:
        if any(x['id'] == row['id'] and x['stage'] == row['stage'] for x in svae['consumers']):
            deps['svae'] = svae_keys[svae['source']]
    if row['stage'] == 'sft' and row['run_number'] in (23, 39, 40):
        deps['pretrain'] = 'pretrain/run39' if row['run_number'] == 39 else 'pretrain/run23'
    if row['run_number'] == 19:
        deps['encoder'] = 'pretrain/run19-encoder'
    return Archive(ROOT).add_policy(run=row['id'], name=row['run_name'], stage=row['stage'], source=p,
                                     share_decoder=row['share_vlm'], dependencies=deps)


def main():
    ROOT.joinpath('migration').mkdir(parents=True, exist_ok=True)
    atomic_json(ROOT / 'migration/status.json', {'status': 'importing', 'source_cleanup': 'not_started'})
    try:
        dependencies = add_models()
        for stage in ['pretrain', 'sft']:
            rows = [r for r in PLAN['policies'] if r['stage'] == stage]
            with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
                for entry in pool.map(import_row, [(row, dependencies) for row in rows]):
                    print(json.dumps({'published': entry['id']}), flush=True)
        atomic_json(ROOT / 'migration/status.json', {'status': 'copied_verifying', 'source_cleanup': 'not_started'})
        archive = Archive(ROOT)
        report = archive.verify(reconstruct=True, report=ROOT / 'migration/verification.json')
        assert len(report['entries']) == 53, len(report['entries'])
        atomic_json(ROOT / 'migration/status.json', {'status': 'verified', 'entries': 53, 'source_cleanup': 'not_started'})
        print('IMPORT_AND_BYTE_EQUIVALENCE_VERIFICATION_COMPLETE', flush=True)
    except BaseException as exc:
        atomic_json(ROOT / 'migration/status.json', {'status': 'failed', 'source_cleanup': 'not_started', 'error': str(exc)})
        raise


if __name__ == '__main__':
    main()