"""Record policy inputs before wrapper stage updates; serve ordinary reference control."""
import argparse
import asyncio
import http
import logging
import time
import jax
import numpy as np
from websockets.asyncio.server import serve
from openpi_client import msgpack_numpy
from serve_checkpoint import load_policy
from common import ROOT, EVAL, PORT, local_only, progress, save_npz, write_json, sha

async def main(variant):
    local_only()
    policy = load_policy(EVAL.parent / 'b1k_sft' / variant)
    dummy = {'robot_r1::proprio': np.zeros(61, np.float32), 'task_id': np.array([0])}
    for name in ['zed_link', 'left_realsense_link', 'right_realsense_link']:
        dummy[f'robot_r1::robot_r1:{name}:Camera:0::rgb'] = np.zeros((224, 224, 3), np.uint8)
    for _ in range(27):
        assert np.isfinite(policy.act(dummy).numpy()).all()
    policy.reset()
    policy.base_policy._rng = jax.random.key(0)
    last = None

    def snapshot(obs, meta, stage):
        ident = meta['condition']
        path = ROOT / 'observations' / f'{ident}.npz'
        if path.exists():
            return
        arrays = {k: np.asarray(v) for k, v in obs.items() if k == 'robot_r1::proprio' or k.endswith('::rgb') or k == 'task_id'}
        assert arrays['robot_r1::proprio'].shape == (61,)
        save_npz(path, **arrays)
        write_json(path.with_suffix('.json'), dict(meta, stage=int(stage), source_model=variant,
                   seed=0, task=0, sha256=sha(path), created=time.time(),
                   step_convention='number of completed environment steps; input before next control action',
                   shapes={k:list(v.shape) for k,v in arrays.items()}))
        print('SAVED', ident, 'stage', stage, flush=True)

    async def handler(ws):
        nonlocal last
        await ws.send(msgpack_numpy.packb({'task_id': 0, 'variant': variant}))
        async for message in ws:
            obs = msgpack_numpy.unpackb(message)
            if 'reset' in obs:
                policy.reset()
                policy.base_policy._rng = jax.random.key(0)
                last = None
                continue
            meta = obs.pop('_capture', {})
            if meta.get('fallback'):
                assert last is not None
                old_obs, old_stage, old_meta = last
                snapshot(old_obs, dict(meta, actual_step=old_meta['step'], reason='reference episode terminated before requested step'), old_stage)
                await ws.send(msgpack_numpy.packb({'action': np.zeros(23, np.float32)}))
                continue
            if meta.get('condition'):
                snapshot(obs, dict(meta, actual_step=meta['step']), policy.current_stage)
            last = (obs, policy.current_stage, meta)
            action = policy.act(obs).numpy()
            assert action.shape == (23,) and np.isfinite(action).all()
            if meta.get('step', 0) % 100 == 0:
                progress('capture_progress.json', variant=variant, instance=meta.get('instance'), step=meta.get('step'), stage=policy.current_stage)
            await ws.send(msgpack_numpy.packb({'action': action}))

    def health(connection, request):
        if request.path == '/healthz':
            return connection.respond(http.HTTPStatus.OK, 'OK\n')
    async with serve(handler, '127.0.0.1', PORT, compression=None, max_size=None,
                     ping_interval=None, process_request=health) as server:
        progress(f'logs/{variant}_ready.json', state='ready')
        await server.serve_forever()

if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    p = argparse.ArgumentParser(); p.add_argument('variant')
    asyncio.run(main(p.parse_args().variant))