#!/usr/bin/env python3
"""Expert worker: holds a set of (layer, expert) pairs resident (loaded from its local copy of the checkpoint) and executes them
on activations sent by the coordinator. Exact: the same gather_qmm + swiglu sequence as the local block, per (row, k) slot.

  python -m transport.worker --model <dir> --listen 0.0.0.0:9411 --token-file ~/.titan_token [--allow 192.168.2.80]
"""
from __future__ import annotations
import argparse, json, os, socket, sys, threading, time
import numpy as np
import mlx.core as mx
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from mlx_lm.models.activations import swiglu
from titan.registry import load_spec
from titan.store import ExpertStore
from titan.slots import SlotTable
from transport.protocol import Framer


class Worker:
    def __init__(self, model_dir, capacity_per_layer=None):
        self.spec = load_spec(model_dir)
        self.store = ExpertStore(self.spec)
        q = self.spec.quantization
        self.gs, self.bits = q["group_size"], q["bits"]
        self.tables = {}
        self.cap = capacity_per_layer or self.spec.num_experts
        self.lock = threading.Lock()
        self.exec_count = 0; self.exec_time = 0.0

    def table(self, l):
        if l not in self.tables:
            self.tables[l] = SlotTable(l, self.cap, self.store)
        return self.tables[l]

    def load(self, assignments):
        """assignments: {layer: [experts]} -> make them resident."""
        n = 0
        for l, es in assignments.items():
            t = self.table(int(l)); t.ensure([int(e) for e in es]); n += len(es)
        mx.eval(*[b for t in self.tables.values() for b in t.buf.values()])
        return n

    def execute(self, layer, rows_f16: np.ndarray, pairs):
        """rows_f16: [R, D] activations; pairs: list of (row, expert). Returns [len(pairs), D] float16 outputs."""
        t0 = time.perf_counter()
        tab = self.table(layer)
        experts = [e for _, e in pairs]
        slots = tab.ensure(list(dict.fromkeys(experts)))
        slot_of = dict(zip(dict.fromkeys(experts), slots))
        x = mx.array(rows_f16)  # [R, D]
        ridx = mx.array([r for r, _ in pairs], dtype=mx.uint32)
        sidx = mx.array([slot_of[e] for e in experts], dtype=mx.uint32)
        xs = x[ridx]                                   # [P, D]
        xe = mx.expand_dims(xs, (-2, -3))              # [P, 1, 1, D]
        b = tab.buf
        def qmm(h, m):
            return mx.gather_qmm(h, b[f"{m}.weight"], b[f"{m}.scales"], b.get(f"{m}.biases"), rhs_indices=sidx.reshape(-1, 1),
                                 transpose=True, group_size=self.gs, bits=self.bits)
        up = qmm(xe, "up"); gate = qmm(xe, "gate")
        y = qmm(swiglu(gate, up), "down").squeeze(-2).squeeze(-2)  # [P, D]
        out = np.array(y.astype(mx.float16))
        self.exec_count += 1; self.exec_time += time.perf_counter() - t0
        return out


def serve(args):
    token = open(os.path.expanduser(args.token_file), "rb").read().strip()
    allow = set(args.allow.split(",")) if args.allow else None
    w = Worker(args.model, args.capacity)
    host, port = args.listen.rsplit(":", 1)
    srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM); srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind((host, int(port))); srv.listen(4)
    print(f"[worker] {socket.gethostname()} listening on {args.listen}, model {w.spec.model_type}, cap {w.cap}", flush=True)
    while True:
        conn, addr = srv.accept()
        if allow and addr[0] not in allow:
            print(f"[worker] refused {addr}", flush=True); conn.close(); continue
        conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        conn.settimeout(args.timeout)
        threading.Thread(target=handle, args=(conn, addr, token, w), daemon=True).start()


def handle(conn, addr, token, w: Worker):
    f = Framer(conn, token)
    try:
        kind, h, _ = f.recv()
        if kind != "hello" or h.get("proto") != "TX1" or h.get("model_type") != w.spec.model_type:
            f.send("error", {"msg": "bad hello"}); conn.close(); return
        f.send("hello_ack", {"host": socket.gethostname(), "num_layers": w.spec.num_layers, "num_experts": w.spec.num_experts,
                              "cap": w.cap, "resident": {str(l): t.resident() for l, t in w.tables.items()}})
        print(f"[worker] coordinator {addr} connected", flush=True)
        while True:
            kind, h, payload = f.recv()
            if kind == "ping":
                f.send("pong", {"t": time.time()})
            elif kind == "load":
                n = w.load(h["assignments"])
                f.send("load_reply", {"loaded": n, "resident": {str(l): t.resident() for l, t in w.tables.items()}})
            elif kind == "exec":
                R, D = h["rows"], h["dim"]
                x = np.frombuffer(payload, dtype=np.float16).reshape(R, D)
                out = w.execute(h["layer"], x, [tuple(p) for p in h["pairs"]])
                f.send("exec_reply", {"layer": h["layer"], "n": out.shape[0], "exec_ms": w.exec_time * 1e3}, out.tobytes())
            else:
                f.send("error", {"msg": f"unexpected {kind}"})
    except (ConnectionError, socket.timeout, OSError) as e:
        print(f"[worker] {addr} closed: {e}", flush=True)
    finally:
        conn.close()


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--listen", default="0.0.0.0:9411")
    ap.add_argument("--token-file", default="~/.titan_token")
    ap.add_argument("--allow", default=None, help="comma-separated coordinator IPs allowed to connect")
    ap.add_argument("--capacity", type=int, default=None)
    ap.add_argument("--timeout", type=float, default=600.0)
    serve(ap.parse_args())