"""Coordinator-side connection to one expert worker."""
from __future__ import annotations
import os, socket, time
import numpy as np
from transport.protocol import Framer


class WorkerClient:
    def __init__(self, host: str, port: int, token: bytes, model_type: str, timeout: float = 30.0, name: str = None):
        self.host, self.port, self.name = host, port, name or host
        self.sock = socket.create_connection((host, port), timeout=timeout)
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.f = Framer(self.sock, token)
        self.f.send("hello", {"proto": "TX1", "model_type": model_type})
        kind, h, _ = self.f.recv()
        if kind != "hello_ack":
            raise ConnectionError(f"{host}: {h}")
        self.info = h
        self.exec_calls = 0; self.exec_wall = 0.0; self.rows_sent = 0

    def ping(self, n=5):
        rtts = []
        for _ in range(n):
            t = time.perf_counter(); self.f.send("ping", {}); self.f.recv(); rtts.append((time.perf_counter() - t) * 1e3)
        return rtts

    def load(self, assignments):
        self.f.send("load", {"assignments": {str(l): list(map(int, es)) for l, es in assignments.items()}})
        kind, h, _ = self.f.recv()
        if kind != "load_reply":
            raise RuntimeError(h)
        return h

    def exec_async(self, layer: int, rows: np.ndarray, pairs):
        """Send an exec request (non-blocking on the reply)."""
        self.f.send("exec", {"layer": int(layer), "rows": int(rows.shape[0]), "dim": int(rows.shape[1]), "pairs": [[int(r), int(e)] for r, e in pairs]}, rows.tobytes())
        self._t0 = time.perf_counter(); self.rows_sent += rows.shape[0]

    def exec_wait(self):
        kind, h, payload = self.f.recv()
        if kind != "exec_reply":
            raise RuntimeError(h)
        self.exec_calls += 1; self.exec_wall += time.perf_counter() - self._t0
        return np.frombuffer(payload, dtype=np.float16).reshape(h["n"], -1), h

    def stats(self):
        return {"worker": self.name, "exec_calls": self.exec_calls, "exec_wall_s": self.exec_wall, "bytes_sent": self.f.bytes_sent, "bytes_recv": self.f.bytes_recv}

    def close(self):
        self.sock.close()