import copy
import importlib.util
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
import time
import traceback

from prepare_model import R,REPO,SOURCE,MODEL,TARGET,ACTIVE,ENV,REG,API,CONTAINER,save,sha,cmd,inspect,gpu
from tools.stream_benchmark_runtime import Runtime
import yaml

CONFIG=Path('/workspace/contextra-runtime-config/test-run/medium6-batch4.yaml')
plan=json.loads((R/'plan.json').read_text())
COMPOSE=['docker','compose','--env-file',str(ENV),'-f',str(REPO/'docker-compose.yml'),'-f',str(ACTIVE)]


def atomic(path,data,expected):
    assert sha(path)==expected
    st=path.stat();fd,temp=tempfile.mkstemp(prefix='.medium-b4-',dir=path.parent)
    try:
        os.fchmod(fd,st.st_mode&0o777);os.fchown(fd,st.st_uid,st.st_gid)
        with os.fdopen(fd,'wb') as handle:handle.write(data);handle.flush();os.fsync(handle.fileno())
        assert sha(path)==expected;os.replace(temp,path)
    finally:
        if os.path.exists(temp):os.unlink(temp)


def configuration(version=3):
    example=yaml.safe_load((REPO/'tools/examples/stream_benchmark.example.yaml').read_text())
    example.update(expected_image=plan['analytics']['Image'],identity_manifest='/workspace/contextra-runtime-config/test-run/medium16-registry.json',
        registration_file=str(REG),lock_file='/workspace/contextra-runtime-config/test-run/benchmark.lock',
        redis=dict(access_file='/workspace/contextra-runtime-config/metadata-active/access.json',secrets_file='/workspace/contextra-runtime-config/metadata-active/secrets.json'))
    example['model'].update(version=version,instances=6 if version==3 else 4,max_batch=4 if version==3 else 1,queue_delay_us=5000 if version==3 else 0)
    example['artifacts']=[dict(name='engine',path=str((TARGET/'3/model.plan') if version==3 else MODEL/'2/model.plan'),
        sha256=sha((TARGET/'3/model.plan') if version==3 else MODEL/'2/model.plan')),
        dict(name='model-config',path=str((TARGET if version==3 else MODEL)/'config.pbtxt'),sha256=sha((TARGET if version==3 else MODEL)/'config.pbtxt')),
        dict(name='registry',path=str(REG),sha256=sha(REG))]
    return example


def reflect(label,target):
    save(label+'-reflection.json',dict(at=time.time(),target=target))
    cmd(COMPOSE+['config','--quiet'])
    cmd(COMPOSE+['up','--no-start','--no-deps','--force-recreate','--pull','never','analytics'],120)
    cmd(['docker','start',CONTAINER])
    end=time.monotonic()+120
    while time.monotonic()<end:
        current=inspect(CONTAINER)
        if current['State'].get('Health',{}).get('Status')=='healthy':break
        time.sleep(1)
    else:raise RuntimeError('Analytics readiness deadline')
    assert current['Image']==plan['analytics']['Image']
    assert dict(v.split('=',1) for v in current['Config']['Env'])==dict(v.split('=',1) for v in plan['analytics']['Config']['Env'])
    for key in ('NanoCpus','CpuQuota','CpusetCpus','Memory','MemorySwap','RestartPolicy'):
        assert current['HostConfig'][key]==plan['analytics']['HostConfig'][key]
    old={x['Destination']:x for x in plan['analytics']['Mounts']};new={x['Destination']:x for x in current['Mounts']}
    model_dest='/opt/contextra/models/triton/RFDETRMedium'
    assert set(old)==set(new)
    for dest in old:
        expected=dict(old[dest])
        if dest==model_dest and target==3:expected['Source']=str(TARGET)
        assert new[dest]==expected, dest
    runtime=Runtime(configuration(target),R)
    for name,cfg in plan['configs'].items():
        actual=runtime.triton('v2/models/'+name+'/config')
        if name!='RFDETRMedium':assert actual==cfg
        else:
            assert actual['max_batch_size']==(4 if target==3 else 1)
            assert actual['instance_group'][0]['count']==(6 if target==3 else 4)
            assert [int(v) for v in actual['version_policy']['specific']['versions']]==[target]
    runtime.resources()
    save(label+'-ready.json',dict(at=time.time(),image=current['Image'],container=current['Id'],model=runtime.triton('v2/models/RFDETRMedium/config')))
    return runtime


def direct_check():
    code=(R/'direct_check.py').read_text().replace("R=Path('/evidence')",'R=Path('+repr(str(R))+')')
    # Only bounded fixture tensors cross docker exec stdin; no service mount/SDK sharing.
    destination='/tmp/codex/ltb1021/2026-09-16_contextra-medium6-batch4_r01/direct'
    cmd(['docker','exec',CONTAINER,'mkdir','-p',destination])
    for name in ('direct_check.py','input-20.bin','input-80.bin','input-160.bin','new-batch-4.json'):
        cmd(['docker','cp',str(R/name),CONTAINER+':'+destination+'/'+name])
    code=(R/'direct_check.py').read_text().replace("R=Path('/evidence')",'R=Path('+repr(destination)+')')
    result=subprocess.run(['docker','exec','-i',CONTAINER,'python3','-'],input=code,capture_output=True,text=True,timeout=35)
    save('direct-check-command.json',dict(exit=result.returncode,stdout=result.stdout,stderr=result.stderr[-2000:]))
    assert result.returncode==0,'TRITON_REQUEST_CONTRACT_FAILED'
    cmd(['docker','cp',CONTAINER+':'+destination+'/direct-contract.json',str(R/'direct-contract.json')])
    for i in range(4):cmd(['docker','cp',CONTAINER+':'+destination+f'/triton-request-{i}.npz',str(R/f'triton-request-{i}.npz')])
    # Exact task-owned files only, not a shared temporary-directory cleanup.
    names=['direct_check.py','input-20.bin','input-80.bin','input-160.bin','new-batch-4.json','direct-contract.json']+[f'triton-request-{i}.npz' for i in range(4)]
    cmd(['docker','exec',CONTAINER,'rm',*[(destination+'/'+n) for n in names]])
    cmd(['docker','exec',CONTAINER,'rmdir',destination])


def run(resume=False):
    if resume:
        assert (R/'reflection-start.json').exists() and not (R/'reflection-resume.json').exists()
        assert not (R/'ascent-start.json').exists()
        shutil.copy2(R/'failure.json',R/'reflection-preparation-failure.json')
        save('reflection-resume.json',dict(at=time.time(),reason='Docker Env order differed; dictionary content and image were identical. No channel Start or ascent occurred.'))
    else:
        assert not (R/'reflection-start.json').exists()
    assert json.loads((R/'engine-check.json').read_text())['contract_pass']
    assert sha(ACTIVE)==plan['fingerprints'][str(ACTIVE)] and sha(REG)==plan['fingerprints'][str(REG)]
    if not resume:
        assert not TARGET.exists() and not CONFIG.exists()
        (TARGET/'3').mkdir(parents=True,mode=0o700)
        shutil.copy2(R/'medium-b4.plan',TARGET/'3/model.plan');shutil.copy2(R/'medium-config.candidate',TARGET/'config.pbtxt')
        save('artifact.json',dict(engine_sha256=sha(TARGET/'3/model.plan'),config_sha256=sha(TARGET/'config.pbtxt'),
            onnx_sha256=sha(SOURCE/'model.onnx'),target=str(TARGET),visual_quality='pending',profile=[1,4,4],precision='FP16'))
        cfg=configuration();CONFIG.write_text(yaml.safe_dump(cfg,sort_keys=False));CONFIG.chmod(0o600)
    else:
        assert sha(TARGET/'3/model.plan')==sha(R/'medium-b4.plan')
        assert sha(TARGET/'config.pbtxt')==sha(R/'medium-config.candidate')
        cfg=configuration();assert yaml.safe_load(CONFIG.read_text())==cfg
    original=ACTIVE.read_bytes();compose=yaml.safe_load(original)
    volumes=compose['services']['analytics'].setdefault('volumes',[])
    dest='/opt/contextra/models/triton/RFDETRMedium'
    assert not any(v.get('target')==dest for v in volumes)
    volumes.append(dict(type='bind',source=str(TARGET),target=dest,read_only=True,bind={'create_host_path':False}))
    candidate=yaml.safe_dump(compose).encode();candidate_hash=__import__('hashlib').sha256(candidate).hexdigest()
    if not resume:
        save('reflection-start.json',dict(at=time.time(),config_after=candidate_hash,tool_hashes={p.name:sha(p) for p in (REPO/'tools').glob('stream_benchmark*.py')}))
    runtime=Runtime(cfg,R);changed=False;retain=False;failed=None
    originals=[r['stream_id'] for r in plan['rows'] if r['original']]
    try:
        for sid in originals:
            q=runtime.stop(sid);save('original-stop-'+sid+'.json',q)
        atomic(ACTIVE,candidate,plan['fingerprints'][str(ACTIVE)]);changed=True
        runtime=reflect('candidate-corrected' if resume else 'candidate',3)
        direct_check()
        validation=cmd(['python3',str(REPO/'tools/stream_benchmark.py'),'validate',str(CONFIG)],60)
        save('tool-validation.json',json.loads(validation))
        args=['python3','-B',str(REPO/'tools/stream_benchmark.py'),'run',str(CONFIG),'--output',str(R/'ascent')]
        save('ascent-start.json',dict(at=time.time(),command=args))
        with (R/'ascent.log').open('x') as log:
            completed=subprocess.run(args,stdout=log,stderr=log,timeout=1500)
        assert completed.returncode==0,'ASCENT_RUNTIME_OR_CLEANUP_FAILED'
        result=json.loads((R/'ascent/result.json').read_text())
        last=result['stages'][-1];baseline=json.loads(Path('/tmp/codex/ltb1021/2026-09-16_contextra-medium4-sixteen_r01/stage16-summary.json').read_text())
        mean=sum(r['result_fps'] for r in last['rows'])/len(last['rows']);prior=sum(r['result_fps'] for r in baseline['rows'])/len(baseline['rows'])
        retain=last['complete'] and len(last['rows'])==16 and (last['accepted'] or mean>prior*1.05)
        save('retention-decision.json',dict(retain=retain,mean_result_fps=mean,prior_mean_result_fps=prior,quality='pending',
            reason='Stable complete ascent and accepted/improved throughput' if retain else 'No meaningful improvement or incomplete ascent'))
    except Exception as error:
        failed=dict(type=type(error).__name__,message=str(error)[:300],trace=traceback.format_exc())
        save('failure.json',failed)
        print('Runtime stopped:',type(error).__name__,str(error)[:160],flush=True)
    finally:
        if not retain:
            for row in plan['rows']:
                q=runtime.states()[row['stream_id']]
                cleanup=q.get('last_cleanup') or {}
                if q.get('child_pid') is None and cleanup.get('child_reaped') and cleanup.get('writer_released'):
                    continue
                if q['state']!='stopped' or q.get('child_pid'):
                    # Do not repeat Stop already owned by the benchmark cleanup.
                    prefix='stop-'+__import__('hashlib').sha256(row['stream_id'].encode()).hexdigest()[:16]+'-*.json'
                    assert not list((R/'ascent').glob(prefix)),'UNCLEAN_BENCHMARK_CHILD_NO_REPEAT'
                    runtime.stop(row['stream_id'])
            if changed:
                atomic(ACTIVE,original,candidate_hash);runtime=reflect('restored',2)
            for sid in originals:runtime.start(sid)
        final=runtime.states()
        assert all(final[s]['state']=='running' for s in originals)
        assert all(final[r['stream_id']]['state']=='stopped' and final[r['stream_id']].get('child_pid') is None for r in plan['rows'] if not r['original'])
        assert sha(REG)==plan['fingerprints'][str(REG)] and sha(ENV)==plan['fingerprints'][str(ENV)]
        assert sha(MODEL/'2/model.plan')==plan['fingerprints'][str(MODEL/'2/model.plan')]
        assert inspect('contextra-redis-1')['Id']==plan['redis']
        sender=json.loads(cmd(['python3',str(REPO/'tools/rtsp_test.py'),'list']))
        assert all(sender[0].get(k)==plan['sender'][0].get(k) for k in ('executionId','uri','video','mode','state'))
        save('final.json',dict(at=time.time(),candidate_retained=retain,error=failed,states=final,
            model=runtime.triton('v2/models/RFDETRMedium/config'),resources=runtime.resources(),
            registry_preserved=True,original_engine_preserved=True,sender_redis_preserved=True,image_unchanged=True,
            actual_persistent_consistent=sha(ACTIVE)==(candidate_hash if retain else plan['fingerprints'][str(ACTIVE)])))
        print('Final:', 'candidate retained' if retain else 'instance4/batch1 restored',flush=True)


if __name__=='__main__':run(resume=len(sys.argv)>1 and sys.argv[1]=='resume-reflection')