"""Transport layer for E2/E6: TCP framing with explicit byte accounting, a remote worker that receives caches/texts and
runs the receiver model, and a client. Loopback and LAN are both supported (labelled by the caller)."""
from __future__ import annotations
import socket, struct, json, time, sys, traceback

MAGIC = b"WP30"

def send_frame(sock: socket.socket, header: dict, payload: bytes) -> tuple[int, float]:
    h = json.dumps(header).encode("utf-8")
    t0 = time.perf_counter()
    sock.sendall(MAGIC + struct.pack("!II", len(h), len(payload)) + h + payload)
    return len(h) + len(payload) + 12, (time.perf_counter() - t0) * 1000

def _recvall(sock, n):
    buf = bytearray()
    while len(buf) < n:
        chunk = sock.recv(min(1 << 20, n - len(buf)))
        if not chunk:
            raise ConnectionError("socket closed")
        buf += chunk
    return bytes(buf)

def recv_frame(sock: socket.socket) -> tuple[dict, bytes, float]:
    t0 = time.perf_counter()
    hdr = _recvall(sock, 12)
    assert hdr[:4] == MAGIC, "bad magic"
    hl, pl = struct.unpack("!II", hdr[4:])
    h = json.loads(_recvall(sock, hl).decode("utf-8"))
    p = _recvall(sock, pl)
    return h, p, (time.perf_counter() - t0) * 1000

class Client:
    def __init__(self, host: str, port: int, timeout: float = 600.0):
        self.sock = socket.create_connection((host, port), timeout=timeout)
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
    def call(self, header: dict, payload: bytes = b"") -> tuple[dict, bytes, dict]:
        t0 = time.perf_counter()
        nbytes, send_ms = send_frame(self.sock, header, payload)
        h, p, recv_ms = recv_frame(self.sock)
        return h, p, {"transport_bytes": nbytes, "send_ms": send_ms, "roundtrip_ms": (time.perf_counter() - t0) * 1000, "reply_bytes": len(p) + 12}
    def close(self):
        self.sock.close()

def serve(port: int, model_repo: str) -> int:
    """Remote receiver worker. Commands: ping | run_cache (payload = serialized cache) | run_text | prefill_time."""
    from .backends.mlx_backend import MLXModel, memory_snapshot, reset_peak
    from .cache.serialize import deserialize_cache, codec_decompress, IncompatibleCache
    from .cache.sparse import cache_tensor_bytes
    import mlx.core as mx
    m = MLXModel(model_repo)
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(("0.0.0.0", port)); srv.listen(4)
    print(json.dumps({"worker": "ready", "model": model_repo, "port": port, "revision": m.info.revision}), flush=True)
    while True:
        conn, addr = srv.accept()
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        try:
            while True:
                h, p, recv_ms = recv_frame(conn)
                t_recv_done = time.perf_counter()
                cmd = h.get("cmd")
                reply = {"cmd": cmd, "recv_ms": recv_ms, "payload_bytes": len(p)}
                try:
                    if cmd == "ping":
                        reply["model_revision"] = m.info.revision; reply["config_hash"] = m.info.config_hash
                    elif cmd == "shutdown":
                        send_frame(conn, {"ok": True}, b""); conn.close(); srv.close(); return 0
                    elif cmd == "run_text":
                        ids = h["token_ids"]
                        reset_peak()
                        g = m.generate_text(ids, max_tokens=h.get("max_tokens", 48))
                        reply.update({"text": g["text"], "prefill_ms": g["prefill_ms"], "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"], "output_tokens": g["output_tokens"], "peak_bytes": int(mx.get_peak_memory())})
                    elif cmd == "run_cache":
                        codec = h.get("codec", "none")
                        t0 = time.perf_counter()
                        raw, dec_ms = codec_decompress(p, codec)
                        try:
                            cache, man = deserialize_cache(raw, m.info, h.get("prefix_hash"))
                        except IncompatibleCache as e:
                            reply.update({"status": "incompatible", "error": str(e)}); send_frame(conn, reply, b""); continue
                        load_ms = (time.perf_counter() - t0) * 1000
                        reset_peak()
                        g = m.generate_from_cache(cache, h["query_ids"], max_tokens=h.get("max_tokens", 48))
                        reply.update({"status": "success", "text": g["text"], "decompress_ms": dec_ms, "load_ms": load_ms, "ttft_ms": g["ttft_ms"], "decode_ms": g["decode_ms"],
                                      "output_tokens": g["output_tokens"], "cache_tensor_bytes": cache_tensor_bytes(cache) , "peak_bytes": int(mx.get_peak_memory())})
                    elif cmd == "prefill_cache":
                        # worker builds a cache from token ids and returns it serialized (for pull-style transfer)
                        from .cache.serialize import build_manifest, serialize_cache, codec_compress
                        ids = h["token_ids"]
                        c, t = m.prefill(ids)
                        man = build_manifest(m.info, ids, c, h.get("document_id", "?"), "B1", {})
                        payload = serialize_cache(c, man)
                        payload, cms = codec_compress(payload, h.get("codec", "none"))
                        reply.update({"prefill_ms": t["prefill_ms"], "serialize_ms": cms, "prefix_hash": man["prefix_token_hash"]})
                        send_frame(conn, reply, payload); continue
                    else:
                        reply["error"] = "unknown cmd"
                except Exception as e:
                    reply.update({"status": "implementation_error", "error": f"{type(e).__name__}: {e}", "trace": traceback.format_exc()[-800:]})
                reply["server_ms"] = (time.perf_counter() - t_recv_done) * 1000
                send_frame(conn, reply, b"")
        except (ConnectionError, AssertionError, OSError):
            pass
        finally:
            try: conn.close()
            except Exception: pass