"""Star topology — Flower aggregation server (FedAvg / FedProx).

Runs on the ``agg`` node. Clients train locally and report; this server averages
and writes per-round aggregated metrics + a communication-bytes estimate to
``metrics.json`` (the RQ1/RQ2 evidence). This is the classic centralized FL
baseline the proposal's RQ2 argues against — one aggregator bottleneck + trust
dependency.

    python3 -m fl.star_server --rounds 20 --num-clients 4 \
        --address 0.0.0.0:8080 --aggregation fedavg --out results/metrics.json
"""

from __future__ import annotations

import argparse
import time
from typing import List, Tuple

import flwr as fl
from flwr.common import Metrics

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


def weighted_average(client_metrics: List[Tuple[int, Metrics]]) -> Metrics:
    """Aggregate client evaluate() metrics weighted by #eval examples."""
    if not client_metrics:
        return {}
    total = sum(n for n, _ in client_metrics)
    keys = [k for k, v in client_metrics[0][1].items() if isinstance(v, (int, float))]
    return {
        k: round(sum(n * float(m[k]) for n, m in client_metrics) / total, 5)
        for k in keys
    }


class _MetricsMixin:
    """Mixin that logs each round's aggregated eval metrics + comm bytes."""

    def setup_logging(self, out_path: str, payload_bytes: int, footprint: dict):
        self.out_path = out_path
        self.payload_bytes = payload_bytes
        self.footprint = footprint
        self.cum_bytes = 0
        self._round_start = time.time()

    def configure_fit(self, server_round, parameters, client_manager):
        self._round_start = time.time()
        return super().configure_fit(server_round, parameters, client_manager)

    def aggregate_evaluate(self, server_round, results, failures):
        loss, agg = super().aggregate_evaluate(server_round, results, failures)
        agg = agg or {}
        n_participating = len(results)
        # Star round: each participating client downloads then uploads a full model.
        round_bytes = n_participating * self.payload_bytes * 2
        self.cum_bytes += round_bytes
        record = {
            "topology": "star",
            "round": server_round,
            "loss": round(float(loss), 5) if loss is not None else None,
            "clients_participating": n_participating,
            "clients_failed": len(failures),
            "comm_bytes_round": int(round_bytes),
            "comm_bytes_cum": int(self.cum_bytes),
            "wall_s": round(time.time() - self._round_start, 3),
            "model_params": self.footprint["params"],
            "model_size_kb": self.footprint["size_kb"],
        }
        record.update(agg)  # accuracy, macro_f1, fault_f1, disturbance_f1, ...
        flmetrics.append_record(self.out_path, record)
        acc = agg.get("accuracy")
        print(f"### PROGRESS: [star] round {server_round}/{self.num_rounds} "
              f"acc={acc} fault_f1={agg.get('fault_f1')} "
              f"cum_MB={self.cum_bytes / 1e6:.2f}")
        return loss, agg


class MetricsFedAvg(_MetricsMixin, fl.server.strategy.FedAvg):
    pass


class MetricsFedProx(_MetricsMixin, fl.server.strategy.FedProx):
    pass


def main() -> None:
    ap = argparse.ArgumentParser(description="Flower star aggregation server")
    ap.add_argument("--rounds", type=int, default=20)
    ap.add_argument("--num-clients", type=int, default=4)
    ap.add_argument("--address", default="0.0.0.0:8080")
    ap.add_argument("--aggregation", choices=["fedavg", "fedprox"], default="fedavg")
    ap.add_argument("--proximal-mu", type=float, default=0.1)
    ap.add_argument("--dataset", default="pq_synth",
                    help="must match the clients' --dataset (sets model shape)")
    ap.add_argument("--out", default="results/metrics.json")
    args = ap.parse_args()

    sp = fldata.spec(args.dataset)
    fp = flmodel.footprint(flmodel.build_model(sp.num_classes, sp.n_channels))
    payload = fp["param_bytes"] + fp["buffer_bytes"]

    common = dict(
        fraction_fit=1.0,
        fraction_evaluate=1.0,
        min_fit_clients=args.num_clients,
        min_evaluate_clients=args.num_clients,
        min_available_clients=args.num_clients,
        evaluate_metrics_aggregation_fn=weighted_average,
    )
    if args.aggregation == "fedprox":
        strategy = MetricsFedProx(proximal_mu=args.proximal_mu, **common)
    else:
        strategy = MetricsFedAvg(**common)
    strategy.num_rounds = args.rounds
    flmetrics.reset(args.out)  # fresh file per run — never append onto a prior run
    strategy.setup_logging(args.out, payload, fp)

    print(f"### PROGRESS: [star] server up on {args.address} "
          f"({args.aggregation}, {args.num_clients} clients, {args.rounds} rounds, "
          f"model {fp['params']} params / {fp['size_kb']} KB)")
    fl.server.start_server(
        server_address=args.address,
        config=fl.server.ServerConfig(num_rounds=args.rounds),
        strategy=strategy,
    )
    print("### PROGRESS: [star] server finished")


if __name__ == "__main__":
    main()
