#!/usr/bin/env python3
"""Persistent resident QNN worker for Fast-WAM (one NPU session per process).

Loads the `action_step` HTP context ONCE, stays alive, serves inference over a TCP
socket (no per-inference reload). Fast-WAM's on-device unit here is the MoT **action
expert** denoise step — it cross-attends to the video KV cache produced by the video
world-model DiT (`video_prefill`), which runs host-side for now (see README).

role:
  action_step : one flow-matching denoise step of the action expert.
                inputs  : latents_action [1,32,7], timestep [1], video_keys [30,1,98,3072],
                          video_values [30,1,98,3072], context [1,129,4096], context_mask [1,129]
                output  : pred_noise [1,32,7]   (host applies the Euler/flow-match update)

Usage: resident_worker.py <role> <port>
"""
import sys, os, socket, struct, time
import numpy as np

WS = os.environ.get("FASTWAM_WS", "/root/fastwam_workspace")


def recvall(s, n):
    b = bytearray()
    while len(b) < n:
        c = s.recv(min(n - len(b), 8 << 20))
        if not c:
            raise ConnectionError()
        b += c
    return bytes(b)


def recv_tensors(s):
    n = struct.unpack("<I", recvall(s, 4))[0]
    out = {}
    for _ in range(n):
        nl = struct.unpack("<I", recvall(s, 4))[0]; name = recvall(s, nl).decode()
        dl = struct.unpack("<I", recvall(s, 4))[0]; dt = recvall(s, dl).decode()
        nd = struct.unpack("<I", recvall(s, 4))[0]; shape = struct.unpack(f"<{nd}I", recvall(s, nd * 4))
        nb = struct.unpack("<Q", recvall(s, 8))[0]; data = recvall(s, nb)
        out[name] = np.frombuffer(data, dtype=dt).reshape(shape).copy()
    return out


def send_tensors(s, d):
    parts = [struct.pack("<I", len(d))]
    for name, a in d.items():
        a = np.ascontiguousarray(a)
        nb = name.encode(); dt = str(a.dtype).encode()
        parts.append(struct.pack("<I", len(nb)) + nb)
        parts.append(struct.pack("<I", len(dt)) + dt)
        parts.append(struct.pack("<I", a.ndim) + struct.pack(f"<{a.ndim}I", *a.shape))
        parts.append(struct.pack("<Q", a.nbytes) + a.tobytes())
    s.sendall(b"".join(parts))


def feed(ctx, m):
    """Order inputs to match the context; QNN consumes float32 raws (bool mask too)."""
    out = []
    for nm, shp in zip(ctx.getInputName(), ctx.getInputShapes()):
        a = np.asarray(m[nm], np.float32)
        if tuple(a.shape) != tuple(shp) and a.size == int(np.prod(shp)):
            a = a.reshape(shp)
        out.append(np.ascontiguousarray(a, np.float32))
    return out


def named(ctx, o):
    return {nm: np.asarray(o[i], np.float32) for i, nm in enumerate(ctx.getOutputName())}


def main():
    role, port = sys.argv[1], int(sys.argv[2])
    from qai_appbuilder import QNNConfig, QNNContext, PerfProfile, Runtime, LogLevel
    QNNConfig.Config(qnn_lib_path=f"{WS}/qnn_libs", runtime=Runtime.HTP, log_level=LogLevel.ERROR)

    suffix = os.environ.get("ACTION_CTX_SUFFIX", "")
    ctx = QNNContext(role, f"{WS}/ctx/{role}{suffix}_socid77_archv73.bin")
    PerfProfile.SetPerfProfileGlobal(PerfProfile.BURST)

    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("127.0.0.1", port)); srv.listen(1)
    open(f"{WS}/tmp/w_{role}.ready", "w").write("1")
    conn, _ = srv.accept()
    conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

    while True:
        try:
            req = recv_tensors(conn)
        except (ConnectionError, struct.error):
            break
        _t = time.time()
        o = ctx.Inference(feed(ctx, req))
        print(f"[infer-ms] {role} {(time.time()-_t)*1000:.1f}", flush=True)
        send_tensors(conn, named(ctx, o))

    os._exit(0)  # dodge an appbuilder<->libs teardown double-free (see README device gotchas)


if __name__ == "__main__":
    main()