import copy
import csv
import hashlib
import importlib.util
import json
from pathlib import Path
import shutil
import subprocess
import time
import traceback
import urllib.request
import yaml

from ops import R, REPO, FILES, p

ORIGINAL=list(p.CHANNELS.values())
MODEL=Path('/data/contextra/model-artifacts/deployments/pose-batch-c869cc4-87a0ca96/RFDETRKeypointPreview')
DEST='/opt/contextra/models/triton/RFDETRKeypointPreview'
PHASE='preparation'
STOP_SENT=set()

def api(path,method='GET',body=None):
    if method!='GET':p.ledger('request',phase=PHASE,method=method,path=path)
    return p.api(path,method,body)

def states():
    return {x['stream_id']:x for x in api('/api/v1/streams')['streams']}

def stop(sid):
    before=states()[sid]
    if p.quiescent(before):return before
    identity=(sid,before.get('run_id'),before.get('run_generation'))
    assert identity not in STOP_SENT,'STOP_ALREADY_SENT'
    STOP_SENT.add(identity)
    try:api('/api/v1/streams/'+sid+'/stop','POST')
    except Exception as error:p.save(sid+'-stop-response-error.json',{'type':type(error).__name__})
    end=time.monotonic()+90
    while time.monotonic()<end:
        q=states()[sid]
        assert (q.get('run_id'),q.get('run_generation'))==identity[1:],'RUN_CHANGED'
        if p.quiescent(q):return q
        if q.get('cleanup_error'):raise RuntimeError('CLEANUP_ERROR')
        time.sleep(.5)
    raise RuntimeError('STOP_TIMEOUT_NO_RETRY')

def start(sid):
    assert states()[sid]['state']=='stopped'
    api('/api/v1/streams/'+sid+'/start','POST')
    end=time.monotonic()+90
    while time.monotonic()<end:
        q=states()[sid]
        if q['state']=='running' and (q.get('metadata_path') or {}).get('processed',0)>0:return q
        if q['state']=='error':raise RuntimeError('START_FAILED')
        time.sleep(.5)
    raise RuntimeError('FIRST_FRAME_TIMEOUT')

def patch(sid,changes):
    current=api('/api/v1/streams/'+sid+'/settings')['configured']
    changes={k:v for k,v in changes.items() if current.get(k)!=v}
    if changes:api('/api/v1/streams/'+sid+'/settings','PATCH',dict(expected_revision=current['revision'],**changes))
    result=api('/api/v1/streams/'+sid+'/settings')['configured']
    assert all(result[k]==v for k,v in changes.items()),'SETTINGS_READBACK'
    return result

def triton(suffix):
    code="import json,urllib.request; print(urllib.request.urlopen('http://127.0.0.1:8000/v2/models/RFDETRKeypointPreview/"+suffix+"',timeout=5).read().decode())"
    result=subprocess.run(['docker','exec','contextra-analytics-1','python3','-c',code],capture_output=True,text=True,timeout=10)
    if result.returncode:raise RuntimeError('TRITON_READ_FAILED')
    return json.loads(result.stdout)

def resource():
    container=p.base.inspect('contextra-analytics-1');pid=container['State']['Pid']
    cg=Path('/sys/fs/cgroup'+Path(f'/proc/{pid}/cgroup').read_text().strip().split('::')[1])
    mem={x.split(':')[0]:int(x.split()[1])*1024 for x in Path('/proc/meminfo').read_text().splitlines()}
    cpu={k:int(v) for k,v in (x.split() for x in (cg/'cpu.stat').read_text().splitlines())}
    gpu=json.loads(json.dumps(p.prior.ops.prior.prev.gpu()))
    value=dict(at=time.monotonic(),memory=int((cg/'memory.current').read_text()),available=mem['MemAvailable'],
               cpu=cpu,gpu=gpu,disk=shutil.disk_usage('/var/lib/docker').free)
    assert value['memory']<16*1024**3 and value['available']>8*1024**3,'RAM_BOUND'
    assert value['disk']>15*1024**3,'DISK_BOUND'
    assert all(v[0]<min(40000,v[1]*.9) for v in gpu),'GPU_BOUND'
    return value

def redis_sample(row,frame=False):
    key='contextra:metadata:v1:camera:'+hashlib.sha256(row['source_id'].encode()).hexdigest()+':frames'
    q=p.base.m.query;at=time.monotonic()
    try:info=q('reader','XINFO','STREAM',key)
    except RuntimeError as error:
        if 'no such key' not in str(error):raise
        return dict(at=at,missing=True)
    if not isinstance(info,dict):return dict(at=at,missing=True)
    value=dict(at=at,entries=info['entries-added'],length=info['length'],pttl=q('reader','PTTL',key))
    if frame:
        data=q('reader','XREVRANGE',key,'+','-','COUNT','1')
        if data:
            fields=data[0][1]
            if not isinstance(fields,dict):fields=dict(zip(fields[::2],fields[1::2]))
            value['frame']=json.loads(fields['json'])
    return value

def snapshot(rows,frames=False):
    values=states();at=time.monotonic();channels={}
    for row in rows:
        sid=row['stream_id'];q=values[sid];path=q.get('metadata_path') or {};perf=path.get('performance') or {}
        channels[sid]=dict(state=q['state'],run_id=q.get('run_id'),generation=q.get('run_generation'),
            error=q.get('last_error'),processed=path.get('processed',0),queued=path.get('queued',0),
            stale=path.get('stale',0),drop=path.get('queue_drop',0),rejected=path.get('publisher_rejected',0),
            counts=perf.get('counts',{}),observation_missed=perf.get('observation_missed',0),
            model=q.get('inference_profile'),redis=redis_sample(row,frames))
    return dict(at=at,channels=channels,publisher=api('/api/v1/metadata/status')['metrics'],resources=resource())

def prepare_model():
    assert p.read('engine.json')['exit']==0 and p.read('engine-semantic-check.json')['success']
    target=MODEL/'3/model.plan';target.parent.mkdir(exist_ok=True)
    assert not target.exists()
    shutil.copyfile(R/'pose-b4.plan',target)
    assert p.sha(target)==p.read('engine.json')['sha']
    p.save('model-prepared.json',dict(root=str(MODEL),config_sha=p.sha(MODEL/'config.pbtxt'),engine_sha=p.sha(target)))

def apply():
    global PHASE
    PHASE='apply';plan=p.read('plan.json');image=p.read('image.json')
    assert p.read('resume-approval.json')['status']=='USER_APPROVED_PERFORMANCE_TRIAL'
    assert p.read('batch4-contract.json')['success']
    assert p.sha(MODEL/'3/model.plan')==p.read('engine.json')['sha']
    assert p.read('pose-image.json')['success'] and not (R/'apply-start.json').exists()
    assert p.sha(p.ACTIVE)==plan['active_hash'] and p.sha(p.ENV)==plan['env_hash']
    assert p.sha(p.prior.REG)==plan['registry_hash']
    assert all(p.sha(REPO/name)==digest for name,digest in image['hashes'].items())
    assert all(sid in ORIGINAL or p.quiescent(q) for sid,q in states().items())
    for sid in ORIGINAL:p.save(sid+'-baseline-cleanup.json',stop(sid))
    p.save('apply-start.json',dict(at=time.time(),image=image,model=str(MODEL),active_sha=p.sha(p.ACTIVE),env_sha=p.sha(p.ENV)))
    cfg=yaml.safe_load(p.ACTIVE.read_text());service=cfg['services']['analytics'];service['image']=image['image']
    volumes=[v for v in service['volumes'] if v['target']==DEST];assert len(volumes)==1
    volumes[0]['source']=str(MODEL)
    p.atomic(p.ACTIVE,yaml.safe_dump(cfg).encode(),plan['active_hash'])
    lines=p.ENV.read_text().splitlines(keepends=True)
    found=[i for i,v in enumerate(lines) if v.startswith('CONTEXTRA_IMAGE=')];assert len(found)==1
    lines[found[0]]='CONTEXTRA_IMAGE='+image['tag']+'\n';p.atomic(p.ENV,''.join(lines).encode(),plan['env_hash'])
    command=['docker','compose','--env-file',str(p.ENV),'-f',str(REPO/'docker-compose.yml'),'-f',str(p.ACTIVE)]
    p.cmd(command+['config','--quiet']);p.cmd(command+['up','--no-start','--no-deps','--force-recreate','--pull','never','analytics'])
    p.cmd(['docker','start','contextra-analytics-1'])
    end=time.monotonic()+120
    while time.monotonic()<end:
        current=p.base.inspect('contextra-analytics-1')
        if current['State'].get('Health',{}).get('Status')=='healthy':break
        time.sleep(1)
    else:raise RuntimeError('READINESS_FAILED')
    assert current['Image']==image['image']
    assert dict(x.split('=',1) for x in current['Config']['Env'])==dict(x.split('=',1) for x in plan['analytics']['Config']['Env'])
    assert all(p.quiescent(q) for q in states().values())
    cfg=triton('config');assert cfg['max_batch_size']==4 and cfg['instance_group'][0]['count']==4
    assert int(cfg['dynamic_batching']['max_queue_delay_microseconds'])==5000
    p.save('applied.json',dict(at=time.time(),config=cfg,image=current['Image'],states=states(),resources=resource(),
        active_sha=p.sha(p.ACTIVE),env_sha=p.sha(p.ENV)))
    print('Applied image + Pose4x4; healthy, all channels stopped',flush=True)

def prepare():
    global PHASE
    PHASE='registration';plan=p.read('plan.json');reg=yaml.safe_load(p.prior.REG.read_text())
    regs={v['stream_id']:v for v in reg['added']};rows=[]
    for i in range(6):
        for group,origin in enumerate(ORIGINAL):
            if group and i==5:continue
            original=i==0;sid=origin if original else f'pose16-0916-g{group+1}-c{i+1}'
            source=regs[origin]['source_id'] if original else sid+'-source'
            rows.append(dict(stream_id=sid,source_id=source,origin=origin,original=original))
    assert len(rows)==16
    p.save('trial-plan.json',dict(channels=rows,temporary=[v['stream_id'] for v in rows if not v['original']],
        registration_before=reg,settings_before=plan['settings'],bounds=plan['caps'],at=time.time()))
    assert api('/api/v1/metadata/status')['metrics']['channels']+13<=64
    for row in rows:
        sid=row['stream_id'];base=regs[row['origin']]
        cfg=plan['settings'][base['source_id']]['configured']
        if not row['original']:
            assert sid not in states()
            body={k:base[k] for k in ('server_id','uri','source_type')}
            body.update(stream_id=sid,source_id=row['source_id'],store=False,active_learning=False,autostart=False,
                scenarios='off',analysisFpsTarget=10,inference_profile={'roles':['pose']},frame_time_anchor=base.get('frame_time_anchor'))
            if base.get('source_options'):body['source_options']=base['source_options']
            api('/api/v1/streams?start=false','POST',body)
        changed=patch(sid,dict(pose=dict(cfg['pose'],enabled=True),fire={'enabled':False},
            tracking=cfg['tracking'],intrusion=dict(cfg['intrusion'],enabled=False),analysisFpsTarget=10,
            display_score=cfg['display_score'],display_categories=cfg['display_categories'],
            debug_all_candidates=cfg['debug_all_candidates'],frame_time_anchor=cfg['frame_time_anchor']))
        assert changed['execution_plan']['role']=='pose' and not changed['tracking']['enabled']
    p.save('all-prepared.json',dict(states=states(),at=time.time()))
    print('All16 registrations/settings complete; Start0',flush=True)
    return rows

def cleanup(rows):
    global PHASE
    PHASE='cleanup';report={};plan=p.read('plan.json')
    for row in reversed(rows):
        sid=row['stream_id']
        if sid not in states():continue
        try:
            q=stop(sid);report[sid]={'cleanup':q.get('last_cleanup'),'state':q['state']}
            if not row['original'] and q['state']=='stopped':report[sid]['deleted']=api('/api/v1/streams/'+sid,'DELETE')
        except Exception as error:report[sid]={'error':type(error).__name__,'reason':str(error)[:100]}
        p.save('cleanup.json',report)
    for source,sid in p.CHANNELS.items():
        if not p.quiescent(states()[sid]):continue
        cfg=plan['settings'][source]['configured']
        patch(sid,{k:cfg[k] for k in ('pose','fire','tracking','intrusion','display_score','display_categories','debug_all_candidates','frame_time_anchor','analysisFpsTarget')})
        start(sid)
    p.save('final.json',dict(states=states(),publisher=api('/api/v1/metadata/status'),resources=resource(),
        inventory=p.base.inventory(),settings={s:api('/api/v1/streams/'+s+'/settings') for s in ORIGINAL}))

def trial():
    global PHASE
    assert not (R/'trial-start.json').exists();p.save('trial-start.json',dict(at=time.time()))
    rows=[];summaries=[]
    try:
        rows=prepare()
        for size,seconds in ((4,20),(8,20),(12,20),(16,60)):
            PHASE='start-'+str(size)
            for row in rows[size-4:size]:
                resource();q=start(row['stream_id'])
                assert q['inference_profile']['roles']==['pose'] and q['inference_profile']['active_models']=={'pose':'RFDETRKeypointPreview'}
                print('Started',row['stream_id'],flush=True)
            time.sleep(5);PHASE='observe-'+str(size)
            t0=triton('stats');samples=[snapshot(rows[:size],True)];until=samples[0]['at']+seconds
            p.save(f'window-{size}.json',dict(size=size,samples=samples,triton_before=t0,complete=False))
            while time.monotonic()<until:
                time.sleep(min(5,max(0,until-time.monotonic())))
                sample=snapshot(rows[:size]);samples.append(sample)
                p.save(f'window-{size}.json',dict(size=size,samples=samples,triton_before=t0,complete=False))
                assert all(q['state']=='running' and not q['error'] for q in sample['channels'].values()),'CHANNEL_FAILED'
                if len(samples)>=4 and any(all(s['channels'][sid]['queued']>=16 for s in samples[-3:]) for sid in sample['channels']):
                    raise RuntimeError('SUSTAINED_QUEUE')
            samples[-1]=snapshot(rows[:size],True);t1=triton('stats')
            window=dict(size=size,samples=samples,triton_before=t0,triton_after=t1,complete=True)
            p.save(f'window-{size}.json',window)
            dt=samples[-1]['at']-samples[0]['at'];summary=[]
            for row in rows[:size]:
                sid=row['stream_id'];a=samples[0]['channels'][sid];b=samples[-1]['channels'][sid]
                assert a['run_id']==b['run_id']
                rr=None
                if not a['redis'].get('missing') and not b['redis'].get('missing'):
                    rr=(b['redis']['entries']-a['redis']['entries'])/(b['redis']['at']-a['redis']['at'])
                summary.append(dict(stage=size,stream_id=sid,seconds=dt,result_fps=(b['processed']-a['processed'])/dt,
                    redis_fps=rr,queue_start=a['queued'],queue_end=b['queued'],stale=b['stale']-a['stale'],drop=b['drop']-a['drop'],
                    missed=b['observation_missed']-a['observation_missed'],length=b['redis'].get('length')))
            summaries.extend(summary);p.save('summary.json',summaries)
            print('Stage',size,[(q['stream_id'],round(q['result_fps'],2),None if q['redis_fps'] is None else round(q['redis_fps'],2),q['stale'],q['drop']) for q in summary],flush=True)
            if any((q['stale']>=10 or q['drop']>=10) and q['queue_end']>0 for q in summary):
                raise RuntimeError('STAGE_DROP_NO_FURTHER_START')
    except Exception as error:
        p.save('trial-failure.json',dict(phase=PHASE,type=type(error).__name__,reason=str(error)[:300],trace=traceback.format_exc()))
        print('Trial stopped',PHASE,type(error).__name__,str(error)[:100],flush=True)
    finally:
        if not rows and (R/'trial-plan.json').exists():rows=p.read('trial-plan.json')['channels']
        cleanup(rows or [dict(stream_id=s,original=True) for s in ORIGINAL])
        if summaries:
            with (R/'summary.csv').open('w') as f:
                writer=csv.DictWriter(f,fieldnames=list(summaries[0]));writer.writeheader();writer.writerows(summaries)
        print('Trial complete; restoration recorded',flush=True)

if __name__=='__main__':
    import sys
    {'model':prepare_model,'apply':apply,'trial':trial}[sys.argv[1]]()