"""Bounded CPU orchestration; never mutate the source run or production."""
import concurrent.futures
import csv
import json
from pathlib import Path
import shutil
import sys

LOCAL = Path(__file__).resolve().parent.parent
BASE = LOCAL.parent
SOURCE = BASE / '2026-10-30_contextra-fight-integrated-37clips_r01'
L01 = BASE / '2026-10-29_contextra-self-filming-l01-pose-cache_r01'
sys.path.insert(0, str(BASE / '2026-10-29_contextra-fight-b-refine_r01/helpers'))
import control_refine as c
c.LOCAL = LOCAL
c.REMOTE = c.old.RB + '/' + LOCAL.name
c.PREFIX = 'contextra-space-1030-r01-'
c.REPORT = LOCAL / 'reports'


def sync():
    for name in ('helpers', 'specs.json', 'self_maps.json', 'cases.json', 'candidates.json',
                 'plan.json', 'execution_contract.json', 'checkpoint.ko.md'):
        c.old.scp(str(LOCAL / name), 'si@58.233.207.31:' + c.REMOTE + '/')


def setup():
    c.REPORT.mkdir(exist_ok=True)
    c.checked_resources()
    assert c.old.checked(['docker', 'image', 'inspect', c.old.IMAGE, '--format', '{{.Id}}']).decode().strip() == c.old.IMAGE
    c.old.checked(['test', '!', '-e', c.REMOTE])
    specs = c.read(SOURCE / 'specs.json')
    maps = c.read(SOURCE / 'self_maps.json')
    package = SOURCE / 'Contextra_Fight_Integrated_Review_20261030'
    rendered = sorted(p.parent.name for p in package.glob('*/B_Integrated_Review.mp4'))
    remaining = sorted(set(specs) - set(rendered))
    assert len(specs) == 37 and len(rendered) == 9 and len(remaining) == 28
    catalog = c.read(L01 / 'metadata/catalog.json')['L01']
    parts = []
    for part in catalog['parts']:
        pid = part['part_id']; name = 'SELF_' + pid
        check = c.read(L01 / part['reader_data'] / 'cache_validation.json')
        assert check['status'] == 'PASS' and check['observations'] == part['observations']
        rgb = L01 / part['reader_data'] / 'input_rgb_sha256.json'
        assert rgb.exists()
        specs[name] = dict(cache=c.old.RB+'/'+L01.name+'/'+part['cache'], cache_sha256=check['cache_sha256'],
            selected_samples=check['selected_samples'], rgb_hashes=str(rgb),
            rgb_hashes_remote=c.old.RB+'/'+L01.name+'/'+part['reader_data']+'/input_rgb_sha256.json',
            input_sha256=catalog['input_sha256'], source_file='L01.mp4', domain='self', clip_id='L01',
            part_id=pid, baseline=None)
        with (L01 / part['observation_map']).open(encoding='utf-8-sig', newline='') as f:
            maps[name] = list(csv.DictReader(f))
        parts.append(name)
    assert sum(len(specs[k]['selected_samples']) for k in parts) == 8305
    for name in ('cases.json', 'candidates.json'):
        shutil.copy2(SOURCE / name, LOCAL / name)
    for name, value in (('specs.json', specs), ('self_maps.json', maps), ('plan.json',
        dict(remaining=remaining, preserved=rendered, parts=parts, output_inputs=remaining+['SELF_L01']))):
        (LOCAL / name).write_text(json.dumps(value, ensure_ascii=False, indent=2))
    contract = c.read(SOURCE / 'execution_contract.json')
    contract.update(source_run=SOURCE.name, frozen=True, L01='PART_INDEPENDENT_CPU_ONLY',
        max_workers=2, cpus=2, memory_gib=8, timeout_seconds=1800, rendering_only_inputs=remaining,
        new_diagnostic_inputs=parts, expected_videos=29, video_gap='neutral slate, not an observation')
    (LOCAL / 'execution_contract.json').write_text(json.dumps(contract, indent=2))
    c.snapshot('before')
    c.old.checked(['mkdir', c.REMOTE], sudo=True)
    c.old.checked(['chown', 'si:si', c.REMOTE], sudo=True)
    c.old.checked(['chmod', '700', c.REMOTE])
    c.old.checked(['mkdir', '-p', c.REMOTE+'/work', c.REMOTE+'/reports', c.REMOTE+'/download/validation'])
    sync()
    print('SETUP_PASS 28 reused + L01/8305 new CPU diagnostics', flush=True)


def many(stage, names):
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
        for f in concurrent.futures.as_completed([pool.submit(c.run_job, stage, n) for n in names]):
            f.result()


if __name__ == '__main__':
    action = sys.argv[1]
    if action == 'setup': setup()
    elif action == 'sync': sync()
    elif action == 'one': c.run_job(sys.argv[2], sys.argv[3] if len(sys.argv)>3 else 'all')
    elif action == 'remaining': many('render_space', c.read(LOCAL/'plan.json')['remaining'])
    elif action == 'render-all':
        jobs=[('render_space',n) for n in c.read(LOCAL/'plan.json')['remaining']]
        jobs += [('render_l01_chunk',f'{i:05d}') for i in range(0,8306,1000)]
        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
            for f in concurrent.futures.as_completed([pool.submit(c.run_job,s,n) for s,n in jobs]):f.result()
    elif action == 'parts': many('diagnose_l01', c.read(LOCAL/'plan.json')['parts'])
    elif action == 'l01-corrected': many('render_l01_corrected',[f'{i:05d}' for i in range(0,8306,1000)])
    elif action == 'finish': c.finish()
    elif action == 'pull': c.old.scp('si@58.233.207.31:'+c.REMOTE+'/download/.', str(LOCAL/'Contextra_Fight_Space_Review_20261030'))
    else: raise ValueError(action)