"""Sequential preparation, scratch smoke, fifty training segments and fifty probes."""
from pathlib import Path
import os,sys,json,time,subprocess,fcntl,traceback
BASE=Path(__file__).resolve().parent.parent;OUT=BASE/'scratch_pcqm_cosmo'
os.environ.update(CUDA_VISIBLE_DEVICES='0',TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD='1',
    CUBLAS_WORKSPACE_CONFIG=':4096:8',OMP_NUM_THREADS='8',MKL_NUM_THREADS='8',PYTHONUNBUFFERED='1')


def write(path,value):
    temp=path.with_suffix('.partial');temp.write_text(json.dumps(value,indent=2));temp.replace(path)


def run(args):
    for attempt in range(3):
        phase='smoke_validation' if '--smoke' in args else ('preparing_probe' if args[0]=='scratch_probe.py' else 'loading_training_segment')
        write(OUT/'status.json',dict(phase=phase,pid=os.getpid(),args=args,updated_unix=time.time()))
        proc=subprocess.Popen([sys.executable,'-u',str(BASE/'source'/args[0]),*args[1:]],env=os.environ.copy())
        write(OUT/'active_process.json',dict(pid=proc.pid,args=args,attempt=attempt+1))
        code=proc.wait()
        if code==0:return
        if attempt==2:raise RuntimeError(f'{args} failed after three attempts; inspect driver.log')
        time.sleep(10)


def main():
    OUT.mkdir(parents=True,exist_ok=True)
    lock=(OUT/'driver.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    if (OUT/'complete.json').exists():return
    while True:
        p=OUT/'preparation_status.json'
        state={}
        if p.exists():
            state=json.loads(p.read_text())
            if state.get('phase')=='failed':raise RuntimeError(state.get('error'))
        if (OUT/'preparation_complete.json').exists() and state.get('phase')=='ready_for_training':break
        time.sleep(10)
    if not (OUT/'smoke_validation.json').exists():run(['scratch_train.py','--smoke'])
    for point in range(1,51):
        marker=OUT/'evaluation'/f'probe_{point:03d}_complete.json'
        if marker.exists():continue
        _saved_checkpoint = OUT/'checkpoints'/f'probe_{point:03d}'
        if not (_saved_checkpoint/'hub_verified.json').exists():
            run(['scratch_train.py','--point',str(point)])
        run(['scratch_probe.py',str(point)])
    import pandas as pd
    results=pd.read_csv(OUT/'evaluation/results.csv')
    assert sorted(results.point.tolist())==list(range(1,51))
    assert all(sum((results.point-1)//5==epoch)==5 for epoch in range(10))
    proof=json.loads((OUT/'evaluation/validation.json').read_text());assert proof['passed'] and proof['probes']==50
    hub=json.loads((OUT/'report_hub.json').read_text())
    write(OUT/'training_complete.json',dict(epochs=10,probes=50))
    write(OUT/'complete.json',dict(completed=True,epochs=10,probes=50,hub=hub,completed_unix=time.time()))
    write(OUT/'status.json',dict(phase='complete',epochs=10,probes=50,report=hub['report']))
    from dotenv import load_dotenv
    load_dotenv('/marimo/.env',override=False)
    from huggingface_hub import HfApi,CommitOperationAdd
    api=HfApi();repo='Queimo/rem3di-scratch-pcqm-cosmo'
    api.create_commit(repo_id=repo,operations=[CommitOperationAdd(path_in_repo=n,path_or_fileobj=str(OUT/n)) for n in ['complete.json','training_complete.json','status.json','latest_point.json','latest_probe.json']],commit_message='Mark ten scratch epochs and all fifty verified BigSolDB Kumo probes complete')


if __name__=='__main__':
    try:main()
    except Exception:
        write(OUT/'driver_failure.json',dict(error=traceback.format_exc(),time=time.time()));raise