"""Two-process exactness test: a worker on localhost holds half of every layer's experts; the coordinator executes those
remotely and the rest locally. Logits must equal mlx-lm's bit-for-bit (same kernels, same bytes, same aggregation order)."""
import os, secrets, socket, subprocess, sys, time, json
import mlx.core as mx
import pytest
from pathlib import Path
from mlx_lm.utils import load_model, load_tokenizer
from titan.runtime import TitanModel
from transport.client import WorkerClient

PROMPTS = ["The capital of Canada is", "def fibonacci(n):\n    \"\"\"Return the n-th Fibonacci number.\"\"\"\n"]

@pytest.fixture(scope="module")
def worker(model_dir, tmp_path_factory):
    tok = secrets.token_hex(32).encode()
    tf = tmp_path_factory.mktemp("tok") / "token"; tf.write_bytes(tok)
    port = 9400 + os.getpid() % 100
    p = subprocess.Popen([sys.executable, "-m", "transport.worker", "--model", model_dir, "--listen", f"127.0.0.1:{port}", "--token-file", str(tf), "--allow", "127.0.0.1"],
                         cwd=str(Path(__file__).resolve().parents[1]), stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
    for _ in range(100):
        try:
            s = socket.create_connection(("127.0.0.1", port), timeout=0.5); s.close(); break
        except OSError:
            time.sleep(0.2)
    yield {"port": port, "token": tok, "proc": p}
    p.terminate()

def test_remote_half_experts_bit_identical(model_dir, worker):
    refm, _ = load_model(Path(model_dir)); tok = load_tokenizer(Path(model_dir))
    tm = TitanModel(model_dir, capacity=max(8, tm_cap(model_dir)))
    E = tm.spec.num_experts; L = tm.spec.num_layers
    placement = {(l, e): ("w0" if e % 2 == 1 else None) for l in range(L) for e in range(E)}  # odd experts remote
    w = WorkerClient("127.0.0.1", worker["port"], worker["token"], tm.spec.model_type, name="w0")
    rtt = w.ping(5)
    w.load({l: [e for e in range(E) if e % 2 == 1] for l in range(L)})
    rx = tm.attach_workers({"w0": w}, placement)
    rec = {"rtt_ms_localhost": rtt, "prompts": []}
    for p in PROMPTS:
        ids = mx.array(tok.encode(p))[None]
        ref = refm(ids); mx.eval(ref)
        tm.reset_stats()
        out = tm(ids); mx.eval(out)
        d = float(mx.max(mx.abs(out.astype(mx.float32) - ref.astype(mx.float32))))
        rec["prompts"].append({"prompt": p[:30], "max_abs_dlogit": d, "argmax_equal": bool(mx.array_equal(mx.argmax(out, -1), mx.argmax(ref, -1))), "remote": rx.stats()})
        assert d <= 1e-3, (p, d)
    # greedy decode token-identical
    from mlx_lm import generate
    a = generate(refm, tok, prompt=PROMPTS[0], max_tokens=24, verbose=False)
    b = generate(tm.model, tok, prompt=PROMPTS[0], max_tokens=24, verbose=False)
    rec["greedy_equal"] = a == b; rec["greedy_ref"] = a; rec["greedy_titan"] = b
    assert a == b
    os.makedirs("benchmarks/raw/tests", exist_ok=True)
    json.dump(rec, open(f"benchmarks/raw/tests/distributed_localhost_{tm.spec.model_type}_{time.strftime('%Y%m%dT%H%M%S')}.json", "w"), indent=2)
    w.close()

def tm_cap(model_dir):
    from titan.registry import load_spec
    return load_spec(model_dir).num_experts // 2