"""Pack one finished V4.5 student for the private HF training-output repo.

Writes outputs/v45-upload/<arm>/: the epoch-20 checkpoint without optimizer state
(the preregistered evaluation point), run records, the arm's mixture request and
manifest, the exact sources, W&B run metadata and offline runs, and SHA256SUMS.
Intermediate checkpoints stay on the cluster and are not uploaded.
"""
import argparse
import json
from pathlib import Path
import shutil
import sys
import tomllib

ROOT=Path(__file__).resolve().parents[2]
sys.path.insert(0,str(ROOT/'src'))
import torch
from vimeml.training.data import file_sha,write_json
from vimeml.training.joint_v45 import ARMS,MIXTURE_EPOCHS

RECORDS=('config.json','plan.json','environment.json','summary.json','progress.json',
         'metrics.jsonl','full-validation.json','subset-validation.json')
SOURCES=('scripts/training/prepare_v45_teachers.py','scripts/training/cluster_v45.sh','scripts/training/train_v45.py',
         'scripts/training/export_v45.py','src/vimeml/training/joint_v45.py')


def main():
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--arm',choices=sorted(ARMS),required=True)
    args=parser.parse_args()
    config_path=ROOT/f'configs/train-v45-{args.arm}.toml'
    config=tomllib.loads(config_path.read_text('utf-8'))
    output=ROOT/config['output_dir']
    summary=json.loads((output/'summary.json').read_text('utf-8'))
    if summary['status']!='complete':
        raise ValueError(f'{args.arm} has not completed training')
    bundle=ROOT/'outputs/v45-upload'/args.arm
    if bundle.exists():
        raise ValueError(f'{bundle} exists; remove it to re-export')
    saved=torch.load(output/f'epoch-{MIXTURE_EPOCHS}.pt',map_location='cpu',weights_only=True)
    if saved['epoch']!=MIXTURE_EPOCHS-1 or saved['step']!=summary['max_steps'] or saved['config']!=config:
        raise ValueError('Epoch-20 checkpoint does not match the completed run')
    history=saved['training_state']['epoch_results']
    if [r['epoch'] for r in history]!=list(range(1,MIXTURE_EPOCHS+1)):
        raise ValueError('Incomplete epoch validation history')
    del saved['optimizer']
    (bundle/'checkpoints').mkdir(parents=True)
    torch.save(saved,bundle/f'checkpoints/epoch-{MIXTURE_EPOCHS}.pt')
    records=bundle/'records'
    records.mkdir()
    for name in RECORDS:
        if (output/name).exists():
            shutil.copyfile(output/name,records/name)
    write_json(records/'epoch-validation-history.json',history)
    log=ROOT/f'logs/v45/{args.arm}.log'
    if log.exists():
        shutil.copyfile(log,records/'train.log')
    mixture=ROOT/config['data_mixture']['directory']
    (bundle/'mixture').mkdir()
    for name in ('request.json','manifest.json'):
        shutil.copyfile(mixture/name,bundle/'mixture'/name)
    if file_sha(bundle/'mixture/manifest.json')!=config['dataset_sha256']:
        raise ValueError('Mixture manifest differs from the pinned SHA')
    (bundle/'source').mkdir()
    for relative in (config_path.relative_to(ROOT).as_posix(),*SOURCES):
        shutil.copyfile(ROOT/relative,bundle/'source'/relative.replace('/','__'))
    tracking=ROOT/'artifacts/tracking'/output.name
    if tracking.exists():
        shutil.copytree(tracking,bundle/'tracking',ignore=shutil.ignore_patterns('cache','data','config','netrc'))
    files=sorted(p for p in bundle.rglob('*') if p.is_file())
    (bundle/'SHA256SUMS').write_bytes(''.join(f'{file_sha(p)}  {p.relative_to(bundle).as_posix()}\n'
                                              for p in files).encode())
    last=history[-1]
    print(json.dumps({'arm':args.arm,'bundle':str(bundle),'files':len(files)+1,
        'epoch20_step':summary['max_steps'],'lm_bpc':last['lm']['bpc'],'chat_bpc':last['chat']['bpc'],
        'checkpoint_sha256':file_sha(bundle/f'checkpoints/epoch-{MIXTURE_EPOCHS}.pt')}),flush=True)


if __name__=='__main__':
    main()