"""Distributed expert execution: a placement maps (layer, expert) -> worker name (or None = local). For each layer the
execution block splits the (row, k) pairs between local slots and remote workers, sends every worker its batch BEFORE waiting
for any reply (so a layer waits for its slowest worker, not for each in turn), then assembles y[row, k, :] in the original
order so the weighted sum is bit-identical to the single-node path. Workers that time out or fail raise; the caller may
retry locally (fallback) — the expert exists on local NVMe too.
"""
from __future__ import annotations
import time
from typing import Dict, List, Optional, Tuple
import numpy as np
import mlx.core as mx
from mlx_lm.models.activations import swiglu
from transport.client import WorkerClient


class RemoteExecutor:
    def __init__(self, workers: Dict[str, WorkerClient], placement: Dict[Tuple[int, int], Optional[str]], fallback_local: bool = True):
        self.workers = workers
        self.placement = placement
        self.fallback_local = fallback_local
        self.remote_pairs = 0; self.local_pairs = 0; self.layers_remote = 0; self.remote_wait = 0.0; self.failures = 0

    def where(self, layer: int, expert: int) -> Optional[str]:
        return self.placement.get((layer, expert))

    def split(self, layer: int, flat: List[int], k: int):
        """flat: expert ids in (row, k) order. Returns local pairs and per-worker pairs as (position, row, expert)."""
        local, remote = [], {}
        for pos, e in enumerate(flat):
            w = self.where(layer, e)
            if w is None or w not in self.workers:
                local.append((pos, pos // k, e))
            else:
                remote.setdefault(w, []).append((pos, pos // k, e))
        return local, remote

    def run_layer(self, block, x_flat: mx.array, flat: List[int], shape, probs_np=None):
        """Execute one layer's (row, k) pairs locally and remotely; returns y of shape (*shape, D)."""
        k = shape[-1]
        local, remote = self.split(block.table.layer, flat, k)
        R, D = x_flat.shape
        outs = [None] * len(flat)
        # 1. dispatch every remote batch first (activations as float16 rows)
        if remote:
            dt = x_flat.dtype
            x_np = np.array(x_flat.view(mx.uint16)) if dt == mx.bfloat16 else np.array(x_flat)
            for w, pairs in remote.items():
                self.workers[w].exec_async(block.table.layer, x_np, [(r, e) for _, r, e in pairs], dtype=str(dt).split(".")[-1])
            self.layers_remote += 1
        # 2. local experts through the slot table (same code path as single node)
        if local:
            lflat = [e for _, _, e in local]
            rows = [r for _, r, _ in local]
            uniq = list(dict.fromkeys(lflat))
            if block.prefetcher is not None:
                block.prefetcher.before_layer(block.table.layer, uniq); block.prefetcher.after_layer(block.table.layer, uniq, probs_np)
            slot_of = dict(zip(uniq, block.table.ensure(uniq)))
            sidx = mx.array([slot_of[e] for e in lflat], dtype=mx.uint32).reshape(-1, 1)
            xs = mx.expand_dims(x_flat[mx.array(rows, dtype=mx.uint32)], (-2, -3))
            up, gate = block._upgate(xs, sidx) if block.table.fused else (block._qmm(xs, "up", sidx, False), block._qmm(xs, "gate", sidx, False))
            yl = block._qmm(swiglu(gate, up), "down", sidx, False).squeeze(-2).squeeze(-2)  # [P, D]
            for i, (pos, _, _) in enumerate(local):
                outs[pos] = yl[i]
            self.local_pairs += len(local)
        # 3. collect remote replies
        if remote:
            t = time.perf_counter()
            for w, pairs in remote.items():
                try:
                    y, _ = self.workers[w].exec_wait()
                except Exception as exc:
                    self.failures += 1
                    if not self.fallback_local:
                        raise
                    # local fallback: execute the failed pairs locally
                    lflat = [e for _, _, e in pairs]; rows = [r for _, r, _ in pairs]
                    uniq = list(dict.fromkeys(lflat)); slot_of = dict(zip(uniq, block.table.ensure(uniq)))
                    sidx = mx.array([slot_of[e] for e in lflat], dtype=mx.uint32).reshape(-1, 1)
                    xs = mx.expand_dims(x_flat[mx.array(rows, dtype=mx.uint32)], (-2, -3))
                    up, gate = block._upgate(xs, sidx) if block.table.fused else (block._qmm(xs, "up", sidx, False), block._qmm(xs, "gate", sidx, False))
                    yf = block._qmm(swiglu(gate, up), "down", sidx, False).squeeze(-2).squeeze(-2)
                    for i, (pos, _, _) in enumerate(pairs):
                        outs[pos] = yf[i]
                    continue
                ym = mx.array(y).view(mx.bfloat16) if x_flat.dtype == mx.bfloat16 else mx.array(y)
                for i, (pos, _, _) in enumerate(pairs):
                    outs[pos] = ym[i]
                self.remote_pairs += len(pairs)
            self.remote_wait += time.perf_counter() - t
        y = mx.stack(outs, 0).reshape(tuple(shape) + (D,))
        return y

    def stats(self):
        return {"remote_pairs": self.remote_pairs, "local_pairs": self.local_pairs, "layers_remote": self.layers_remote,
                "remote_wait_s": self.remote_wait, "failures": self.failures, "workers": {n: w.stats() for n, w in self.workers.items()}}