"""Mesh topology — decentralized peer-averaging / gossip over FABNetv4.

No central aggregator. Each peer trains locally, publishes its weights on a tiny
HTTP server, pulls its neighbors' weights **directly over FABNetv4**, and
peer-averages. This is the decentralized setting RQ1/RQ3 center on; the network
carries the real per-round weight exchange so convergence reflects true testbed
behavior.

Sync (``--sync sync``): per-round barrier — wait for each neighbor to reach the
current round before averaging. Async (``--sync async``): gossip — average with
whatever each neighbor has published most recently, no barrier (models staleness).

Churn (RQ3): a peer given ``--drop-at-round R`` exits at round R; its neighbors
time out on it and proceed with the reachable subset (graceful degradation).

    python3 -m fl.mesh_peer --name client-ercot --peer-index 0 --num-clients 4 \
        --neighbors 10.0.0.2:8090,10.0.0.3:8090 --listen-port 8090 --rounds 20
"""

from __future__ import annotations

import argparse
import io
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import List, Optional
from urllib.request import urlopen

import numpy as np
import torch

from . import data as fldata
from . import metrics as flmetrics
from . import model as flmodel
from . import trainer as fltrainer


# --- weight blob (de)serialization -----------------------------------------

def serialize_weights(weights: List[np.ndarray]) -> bytes:
    buf = io.BytesIO()
    np.savez(buf, **{f"w{i:03d}": np.asarray(a) for i, a in enumerate(weights)})
    return buf.getvalue()


def deserialize_weights(blob: bytes) -> List[np.ndarray]:
    with np.load(io.BytesIO(blob), allow_pickle=False) as npz:
        return [npz[k] for k in sorted(npz.files)]


class _PublishState:
    """Thread-safe holder for the weights this peer currently serves."""

    def __init__(self):
        self._lock = threading.Lock()
        self._round = -1
        self._blob = b""

    def publish(self, rnd: int, blob: bytes) -> None:
        with self._lock:
            self._round = rnd
            self._blob = blob

    def snapshot(self):
        with self._lock:
            return self._round, self._blob


def _make_handler(state: _PublishState):
    class Handler(BaseHTTPRequestHandler):
        def log_message(self, *_):  # silence per-request logging
            pass

        def do_GET(self):
            if self.path.startswith("/health"):
                self.send_response(200)
                self.end_headers()
                self.wfile.write(b"ok")
                return
            if self.path.startswith("/weights"):
                rnd, blob = state.snapshot()
                if rnd < 0:
                    self.send_response(404)
                    self.end_headers()
                    return
                self.send_response(200)
                self.send_header("Content-Type", "application/octet-stream")
                self.send_header("X-Round", str(rnd))
                self.send_header("Content-Length", str(len(blob)))
                self.end_headers()
                self.wfile.write(blob)
                return
            self.send_response(404)
            self.end_headers()

    return Handler


def _fetch(url: str, min_round: int, sync: bool, timeout_s: int):
    """Return ``(weights, blob_len)`` from a neighbor, or ``None``.

    Sync: poll until the neighbor's published round >= ``min_round`` (barrier).
    Async: take whatever the neighbor has published right now (or None).
    """
    deadline = time.time() + timeout_s
    while True:
        try:
            with urlopen(url, timeout=10) as resp:
                rnd = int(resp.headers.get("X-Round", "-1"))
                blob = resp.read()
            if not sync or rnd >= min_round:
                return deserialize_weights(blob), len(blob)
        except Exception:
            pass  # neighbor not up / round not ready yet
        if time.time() >= deadline:
            return None
        time.sleep(1.0)


def main() -> None:
    ap = argparse.ArgumentParser(description="Decentralized mesh FL peer")
    ap.add_argument("--name", required=True)
    ap.add_argument("--peer-index", type=int, required=True)
    ap.add_argument("--num-clients", type=int, default=4)
    ap.add_argument("--neighbors", default="", help="comma list of host:port")
    ap.add_argument("--listen-port", type=int, default=8090)
    ap.add_argument("--rounds", type=int, default=20)
    ap.add_argument("--local-epochs", type=int, default=1)
    ap.add_argument("--batch-size", type=int, default=64)
    ap.add_argument("--lr", type=float, default=0.01)
    ap.add_argument("--dataset", default="pq_synth")
    ap.add_argument("--class-weights", choices=["balanced", "none"], default="balanced",
                    help="'balanced' up-weights rare classes (needed for imbalanced vsb)")
    ap.add_argument("--alpha", type=float, default=0.5)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--sync", choices=["sync", "async"], default="sync")
    ap.add_argument("--aggregation", choices=["peer_avg", "gossip"], default="peer_avg")
    ap.add_argument("--drop-at-round", type=int, default=0, help="0 = never churn")
    ap.add_argument("--neighbor-timeout", type=int, default=120)
    ap.add_argument("--prefer-gpu", type=int, default=1)
    ap.add_argument("--out", default="")
    args = ap.parse_args()

    prefer_gpu = bool(args.prefer_gpu)
    is_sync = args.sync == "sync"
    neighbors = [n.strip() for n in args.neighbors.split(",") if n.strip()]
    out_path = args.out or f"results/metrics-{args.name}.json"
    flmetrics.reset(out_path)  # fresh file per run — never append onto a prior run
    sp = fldata.spec(args.dataset)

    device = flmodel.get_device(prefer_gpu=prefer_gpu)
    Xtr, ytr, Xte, yte = fldata.load_partition(
        client_index=args.peer_index, num_clients=args.num_clients,
        dataset=args.dataset, alpha=args.alpha, seed=args.seed,
    )
    train_dl, test_dl = fltrainer.make_loaders(Xtr, ytr, Xte, yte, args.batch_size)
    class_weights = (fltrainer.balanced_class_weights(ytr, sp.num_classes)
                     if args.class_weights == "balanced" else None)

    # Common initialization across all peers so peer-averaging is meaningful.
    torch.manual_seed(args.seed)
    model = flmodel.build_model(sp.num_classes, sp.n_channels)
    fp = flmodel.footprint(model)
    payload = fp["param_bytes"] + fp["buffer_bytes"]

    # Start the weight server.
    state = _PublishState()
    httpd = ThreadingHTTPServer(("0.0.0.0", args.listen_port), _make_handler(state))
    threading.Thread(target=httpd.serve_forever, daemon=True).start()
    print(f"### PROGRESS: [mesh] {args.name} serving weights on :{args.listen_port} "
          f"(dataset={sp.name}, device={device.type}, neighbors={neighbors or 'none'}, "
          f"mix={fldata.class_distribution(ytr, sp)})")

    cum_bytes = 0
    for r in range(1, args.rounds + 1):
        if args.drop_at_round and r >= args.drop_at_round:
            print(f"### PROGRESS: [mesh] {args.name} churning out at round {r}")
            break

        fltrainer.train_local(model, train_dl, device,
                              epochs=args.local_epochs, lr=args.lr,
                              class_weights=class_weights)
        # Publish this round's weights before waiting on neighbors (avoids deadlock).
        state.publish(r, serialize_weights(flmodel.get_weights(model)))

        collected = [flmodel.get_weights(model)]
        reached, round_bytes = 0, 0
        for nb in neighbors:
            got = _fetch(f"http://{nb}/weights", min_round=r,
                         sync=is_sync, timeout_s=args.neighbor_timeout)
            if got is not None:
                weights, blen = got
                collected.append(weights)
                reached += 1
                round_bytes += blen
        cum_bytes += round_bytes

        averaged = flmodel.average_weights(collected)
        flmodel.set_weights(model, averaged)

        _, m = fltrainer.evaluate(model, test_dl, device, sp)
        record = {
            "topology": "mesh",
            "peer": args.name,
            "peer_index": args.peer_index,
            "round": r,
            "sync": args.sync,
            "neighbors_reached": reached,
            "neighbors_total": len(neighbors),
            "comm_bytes_round": int(round_bytes),
            "comm_bytes_cum": int(cum_bytes),
            "model_params": fp["params"],
            "model_size_kb": fp["size_kb"],
            **m,
        }
        flmetrics.append_record(out_path, record)
        print(f"### PROGRESS: [mesh] {args.name} round {r}/{args.rounds} "
              f"acc={m['accuracy']} fault_f1={m['fault_f1']} "
              f"reached={reached}/{len(neighbors)} cum_MB={cum_bytes / 1e6:.2f}")

    # Keep serving briefly so lagging neighbors can still fetch the final round.
    time.sleep(10)
    httpd.shutdown()
    print(f"### PROGRESS: [mesh] {args.name} finished")


if __name__ == "__main__":
    main()
