"""Resume the existing 1333-frame reconstruction with project-local model caches."""
import json
import os
from pathlib import Path
import subprocess
import time
ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / 'output'
STATUS = OUT / 'dynhamr_auto_status.json'
def status(stage, **extra):
    payload = dict(stage=stage, time=time.strftime('%Y-%m-%d %H:%M:%S'), **extra)
    STATUS.write_text(json.dumps(payload, indent=2) + '\n')
    print(json.dumps(payload), flush=True)

def required_files():
    hub = ROOT / '.torch_cache/hub'
    files = [hub/'sam/sam_vit_b_01ec64.pth', hub/'aot/R50_DeAOTL_PRE_YTB_DAV.pth',
             hub/'checkpoints/groundingdino_swint_ogc.pth', hub/'checkpoints/video_depth_anything_vits.pth',
             ROOT/'weights/unidepth-v2-vitl14/model.safetensors']
    manifest = (OUT/'vipe_weights_downloads.txt').read_text().splitlines()
    for i in range(0, len(manifest), 3):
        files.append(Path(manifest[i+1].strip().removeprefix('dir='))/manifest[i+2].strip().removeprefix('out='))
    return files

files = required_files()
deadline = time.monotonic() + 5400
while True:
    missing = [str(p) for p in files if not p.is_file() or p.stat().st_size < 1000000 or Path(str(p)+'.aria2').exists()]
    if not missing:
        break
    status('WAITING_WEIGHTS', pending=missing)
    if time.monotonic() > deadline:
        status('FAILED', reason='Weight preparation timeout')
        raise SystemExit(1)
    time.sleep(30)

env = os.environ.copy()
env.update(TORCH_HOME=str(ROOT/'.torch_cache'), HF_HOME=str(ROOT/'.hf_cache'),
           HF_HUB_OFFLINE='1', TRANSFORMERS_OFFLINE='1', PYTHONUNBUFFERED='1',
           VIPE_CACHE_DIR=str(OUT/'vipe_frame_cache'),
           MPLCONFIGDIR='/tmp/matplotlib-dynhamr', HYDRA_FULL_ERROR='1', PYOPENGL_PLATFORM='egl')
dyn = ROOT/'third_party/Dyn-HaMR'
env['PATH'] = str(dyn/'.dynhamr/bin') + ':' + env['PATH']
env['PYTHONPATH'] = ':'.join(str(p) for p in [dyn/'dyn-hamr', dyn/'third-party/hamer', dyn/'third-party/hamer/third-party/ViTPose'])
cmd = [str(dyn/'.dynhamr/bin/python'), '-u', 'run_opt.py', 'data=video_vipe',
       'data.root='+str(dyn/'test'), 'data.seq=2047635068',
       'data.vipe_dir='+str(OUT/'vipe_2047635068'), 'is_static=False',
       'run_opt=True', 'run_vis=True', 'run_prior=True', '+data.vipe_pipeline=dynhamr_cameras',
       'HMP.vid_path='+str(dyn/'test/videos/2047635068.mp4')]
status('RUNNING', command=cmd)
with (OUT/'dynhamr_auto_run.log').open('a') as log:
    result = subprocess.run(cmd, cwd=dyn/'dyn-hamr', env=env, stdout=log, stderr=subprocess.STDOUT)
status('PROCESS_EXITED' if result.returncode == 0 else 'FAILED', returncode=result.returncode,
       log=str(OUT/'dynhamr_auto_run.log'))
raise SystemExit(result.returncode)