"""Publish executed trainable variants beside immutable inference families."""
import json,sys,hashlib,shutil,fcntl
from pathlib import Path
from huggingface_hub import HfApi
root=Path('/workspace/vds-litert-batch2')
api=HfApi(token=(Path('/workspace/.secrets/hf_token').read_text().strip()))
manifest=json.loads((root/'package/model_manifest.json').read_text());specs={s['key']:s for g in ('general_models','fireviewer_models') for s in manifest[g]}
upstreams=json.loads((root/'prefetch_manifest.json').read_text())['models']
upstreams.update(json.loads((root/'previous-prefetch.json').read_text()))
for s in json.loads((root/'previous-models.json').read_text()):specs[s['key']]={**s,'title':s['key']}
def sha(p):return hashlib.sha256(p.read_bytes()).hexdigest()
for key in sys.argv[1:]:
    directory=root/'learning'/key;report=json.loads((directory/'training_validation.json').read_text());source=json.loads((directory/'source.json').read_text())
    assert report['validation_passed'] and report['optimizer_resume_exact'] and report['restart_logits_exact']
    assert sha(directory/'model.tflite')==report['sha256']
    spec=specs[key];upstream=upstreams[key]
    destination='fireviewer/litert-models' if key.startswith('fireviewer_') else 'Charlbi/Lite_rt_prepared_for_android_dataset_builder'
    if key.startswith('fireviewer_') and not spec.get('publish_default',False):
        # User explicitly requested FireViewer delivery; the pilot permits a
        # private research copy, never a public release under this workflow.
        assert source.get('private_pilot_only') and api.repo_info(destination).private
    stage=directory/'publish';stage.mkdir(exist_ok=True)
    shutil.copy2(directory/'model.tflite',stage/'model.tflite')
    labels=source.get('labels') or json.loads((root/'learning/imagenet-labels.json').read_text())
    labels=[label if labels.count(label)==1 else f'{label} [class {i}]' for i,label in enumerate(labels)]
    prep=source.get('pretrained_cfg',{})
    config={'task':'classification' if prep else 'object_detection','adapter':'classification' if prep else 'yolo','inputWidth':224,'inputHeight':224,'inputLayout':'NCHW','inputType':'FLOAT32','mean':0.,'std':1.,'resizeMode':'stretch','labels':labels,'threads':2,'outputIndex':0,'threshold':.1,'scoreActivation':'softmax' if prep else 'none','dynamicMinSize':128,'dynamicMaxSize':1280,'dynamicStride':32 if key=='convformer_s18' or not prep else 1}
    if prep:
        config.update(channelMean=[v*255 for v in prep['mean']],channelStd=[v*255 for v in prep['std']])
        training={'targetEncoding':'one_hot','targetShape':[1,len(labels)],'inferOutputs':['scores']}
    else:
        config.update(inputWidth=source.get('app_size',320),inputHeight=source.get('app_size',320),mean=0.,std=255.,resizeMode='letterbox' if source['kind']=='yolo' else 'stretch',outputLayout='BCN',coordinates='pixels',outputIndexBoxes=0,outputIndexScores=1)
        if source.get('channelMean'):config.update(channelMean=source['channelMean'],channelStd=source['channelStd'])
        training={'targetEncoding':'boxes_xyxy_class_mask','targetShape':[1,100,6],'inferOutputs':['detections']}
    training.update(trainSignature='train',inferSignature='infer',saveSignature='save',restoreSignature='restore',imageInput='x',targetInput='y',lossOutput='loss',checkpointInput='checkpoint_path',learningRateInput='learning_rate',scope=report['training_scope'])
    config['training']=training
    config['isRgb']=source.get('isRgb',True)
    config['dynamicStride']=report.get('variable_image_preprocessing',{}).get('stride',32 if key=='convformer_s18' else 1)
    config['dynamicMinSize']=source.get('dynamic_min_size',128)
    if (directory/'android_model_config.json').exists():config=json.loads((directory/'android_model_config.json').read_text())
    # The default graph's output order is not necessarily the named detection
    # signature order. Resolve actual tensor indices for the Android adapter.
    import tensorflow as tf
    engine=tf.lite.Interpreter(model_path=str(directory/'model.tflite'),experimental_op_resolver_type=tf.lite.experimental.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES)
    named=engine.get_signature_runner('infer').get_output_details()
    main=engine.get_output_details()
    output_indices={name:next(j for j,d in enumerate(main) if d['index']==details['index']) for name,details in named.items()}
    config['namedOutputIndices']=output_indices
    if config['adapter']=='yolo':
        config['outputIndex']=output_indices['detections']
        config.pop('outputIndexBoxes',None);config.pop('outputIndexScores',None)
    del engine
    (stage/'android_model_config.json').write_text(json.dumps(config,indent=2)+'\n')
    report['runtime_requirement']='LiteRT builtin CPU; disable default XNNPACK delegates; Select TF Ops required for checkpoint Save/Restore'
    report['android_runtime_dependencies']=['org.tensorflow:tensorflow-lite:2.16.1','org.tensorflow:tensorflow-lite-select-tf-ops:2.16.1']
    report['frozen_backbone_sha256']=sha(Path(source['backbone']))
    report['head_initialization']=source.get('head_initialization','source_model_or_identity_adapter')
    builder='build_trainable_fire_multitask.py' if source.get('kind')=='multitask' else 'build_trainable_classifier.py' if prep else 'build_trainable_heatmap.py' if source.get('kind')=='heatmap' else 'build_trainable_detector.py'
    report['conversion_source_sha256']={n:sha(root/n) for n in ['merge_training_graph.py',builder]}
    if source.get('parity_report'):
        parity=Path(source['parity_report']);shutil.copy2(parity,stage/'source_parity.json')
        report['source_parity_report']='source_parity.json'
    if (directory/'runtime216_validation.json').exists():
        compatible=json.loads((directory/'runtime216_validation.json').read_text());assert compatible['sha256']==report['sha256']
        shutil.copy2(directory/'runtime216_validation.json',stage/'runtime216_validation.json')
    (stage/'training_validation.json').write_text(json.dumps(report,indent=2)+'\n')
    contract={'schema':2,'key':key+'_learning','source_model':key,'task':config['task'],'upstream':spec.get('repo') or spec.get('checkpoint_url'),'upstream_revision':upstream.get('sha') or upstream.get('sha256'),'backbone_frozen':True,'training_scope':report['training_scope'],'report':report,'optimizer':'momentum SGD, momentum=0.9','checkpoint_contains':['trainable parameters','momentum buffers','global step'],'requires_select_tf_ops':True,'android_configuration':'android_model_config.json'}
    (stage/'runtime_contract.json').write_text(json.dumps(contract,indent=2)+'\n')
    local=Path(upstream['local_dir'])
    for p in local.iterdir():
        if p.is_file() and (p.name.lower().startswith(('license','notice','copying')) or p.name in ('CHECKPOINT_LICENSE.md','CHECKPOINT_RIGHTS.md','CITATION.md','COMMERCIAL_USE.md','MODEL_RIGHTS_POLICY.md','RIGHTS_AND_ATTRIBUTION.md','THIRD_PARTY_NOTICES.md')):shutil.copy2(p,stage/p.name)
    if (local/'README.md').is_file():shutil.copy2(local/'README.md',stage/'UPSTREAM_README.md')
    if key=='rtmdet_tiny':shutil.copy2(root/'vendor/mmyolo/LICENSE',stage/'LICENSE')
    (stage/'README.md').write_text(f"# {spec['title']} — Android learning variant\n\nSource: `{spec.get('repo')}`, revision `{upstream.get('sha')}`.\n\nThis is a real mutable LiteRT model with `infer`, `train`, `save`, and `restore` signatures. Training executes on Android using the image and reviewed targets. Trainable scope: **{report['training_scope']}**. The visual backbone remains frozen. {report.get('limitation','')}\n\nThis learning variant is stored separately under `models/{key}_learning` and has its own explicit Android configuration.\n\nVariable image dimensions: `{report.get('variable_image_preprocessing', {'mode':'native_dynamic'})}`. The Android configuration chooses a bounded input canvas; the image may have any original size. Images are processed individually, so a dataset need not have homogeneous dimensions.\n\nCheckpoints include parameters, momentum and step. Validation executed actual updates, changed parameters, closed/reopened the interpreter, restored identical predictions and reproduced the next optimizer update. See `training_validation.json`. Tests use synthetic supervision to verify execution and persistence; no accuracy or forgetting benchmark is claimed. Android device execution: `{report.get('android_device_tested',False)}`.\n\nUse LiteRT CPU with XNNPACK disabled and Select TF Ops for checkpoint operations. The default label space is retained; changing class count requires rebuilding the head with the supplied conversion sources. Use only human-reviewed labels for training. Retain replay examples and validate a candidate before activating it.\n\nUpstream licenses and attributions continue to apply.\n")
    (stage/'artifact_manifest.json').write_text(json.dumps({p.name:{'sha256':sha(p),'bytes':p.stat().st_size} for p in stage.iterdir() if p.is_file() and p.name!='artifact_manifest.json'},indent=2)+'\n')
    with (root/'.publish.lock').open('a') as lock:
        fcntl.flock(lock,fcntl.LOCK_EX);parent=api.repo_info(destination,repo_type='model').sha
        commit=api.upload_folder(repo_id=destination,repo_type='model',folder_path=str(stage),path_in_repo=f'models/{key}_learning',parent_commit=parent,commit_message=f'Add on-device learning and optimizer resume: {key}')
    path=f'models/{key}_learning/model.tflite';remote=api.get_paths_info(destination,[path],revision=commit.oid)[0]
    assert remote.size==report['bytes'] and remote.lfs.sha256==report['sha256']
    receipt={'repo':destination,'commit':commit.oid,'path':path,'sha256':report['sha256'],'bytes':report['bytes'],'remote_verified':True}
    (directory/'publication_receipt.json').write_text(json.dumps(receipt,indent=2)+'\n');print(json.dumps(receipt),flush=True)