"""Wire protocol TITAN-X1: length-prefixed frames, HMAC-SHA256 authentication of every frame with a pre-shared token.

Frame = magic(4) | kind(1) | seq(4) | header_len(4) | payload_len(8) | header(JSON) | payload(bytes) | mac(32)
mac = HMAC(token, magic..payload). A worker only answers frames whose MAC verifies; a bad MAC closes the connection.
Kinds: 0 hello, 1 hello-ack, 2 exec request, 3 exec reply, 4 load request (install experts), 5 load reply, 6 ping, 7 pong, 9 error.
Activations are sent as raw float16 rows; results come back as raw float16 rows (one per (row, k) slot), so the coordinator
performs the weighted sum in exactly the local order (bit-for-bit identical aggregation).
"""
from __future__ import annotations
import hashlib, hmac, json, socket, struct

MAGIC = b"TX1\0"
HDR = struct.Struct("<4sBIIQ")
KIND = {"hello": 0, "hello_ack": 1, "exec": 2, "exec_reply": 3, "load": 4, "load_reply": 5, "ping": 6, "pong": 7, "error": 9}
KIND_NAME = {v: k for k, v in KIND.items()}


def _recv_exact(sock: socket.socket, n: int) -> bytes:
    buf = bytearray(n); view = memoryview(buf); got = 0
    while got < n:
        r = sock.recv_into(view[got:], n - got)
        if r == 0:
            raise ConnectionError("peer closed")
        got += r
    return bytes(buf)


class Framer:
    def __init__(self, sock: socket.socket, token: bytes):
        self.sock = sock
        self.token = token
        self.seq = 0
        self.bytes_sent = 0
        self.bytes_recv = 0

    def send(self, kind: str, header: dict, payload: bytes = b""):
        h = json.dumps(header, separators=(",", ":")).encode()
        head = HDR.pack(MAGIC, KIND[kind], self.seq, len(h), len(payload))
        mac = hmac.new(self.token, head + h + payload, hashlib.sha256).digest()
        self.sock.sendall(b"".join([head, h, payload, mac]))
        self.seq += 1
        self.bytes_sent += HDR.size + len(h) + len(payload) + 32

    def recv(self):
        head = _recv_exact(self.sock, HDR.size)
        magic, kind, seq, hl, pl = HDR.unpack(head)
        if magic != MAGIC:
            raise ConnectionError("bad magic")
        h = _recv_exact(self.sock, hl)
        payload = _recv_exact(self.sock, pl)
        mac = _recv_exact(self.sock, 32)
        if not hmac.compare_digest(mac, hmac.new(self.token, head + h + payload, hashlib.sha256).digest()):
            raise ConnectionError("bad MAC")
        self.bytes_recv += HDR.size + hl + pl + 32
        return KIND_NAME[kind], json.loads(h), payload