"""End-to-end test of the TCP protocol with worker subprocesses on the local host."""
import subprocess, sys, threading, time, pathlib, socket, os
import numpy as np
from pamle.protocol import coordinator_run

ROOT = pathlib.Path(__file__).resolve().parents[1]


def _free_port():
    s = socket.socket(); s.bind(("", 0)); port = s.getsockname()[1]; s.close(); return port


def test_protocol_local_workers():
    port = _free_port(); K = 3; res = {}
    def co():
        res.update(coordinator_run(port, K, "logistic", 4, 3 * 2048, (1, 2), 2, "exact", seed=5, log=lambda *a: None))
    t = threading.Thread(target=co, daemon=True); t.start(); time.sleep(0.3)
    env = dict(os.environ, PYTHONPATH=str(ROOT / "src"))
    procs = [subprocess.Popen([sys.executable, "-m", "pamle.worker", "--connect", f"127.0.0.1:{port}", "--id", str(k)], env=env) for k in range(K)]
    t.join(timeout=120); assert not t.is_alive()
    for pr in procs:
        assert pr.wait(timeout=30) == 0
    assert res["stats"]["messages_up"] == K * (1 + 2) and res["stats"]["messages_down"] == K * (1 + 2 + 1)
    assert res["stats"]["bytes_up"] > 0 and res["stats"]["bytes_down"] > 0
    assert res["newton"][1]["grad_norm"] < res["newton"][0]["grad_norm"]
    assert len(res["theta_path"]) == 3 and res["n_hosts"] == 1
    assert res["stats"]["stages"]["A_B_local_fit_and_upload"]["valid_workers"] == K