"""Star topology — Flower client (local training + report).

Runs on each ``client-*`` node. Trains the shared 1D-CNN on this site's non-IID
partition and reports weights to the star server over FABNetv4. Uses a background
retry loop so a client can start before the server is reachable.

    python3 -m fl.star_client --server 10.x.x.x:8080 --client-index 0 \
        --num-clients 4 --local-epochs 1 --batch-size 64 --lr 0.01
"""

from __future__ import annotations

import argparse
import time

import flwr as fl

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


class GridClient(fl.client.NumPyClient):
    def __init__(self, args):
        self.args = args
        self.device = flmodel.get_device(prefer_gpu=args.prefer_gpu)
        self.spec = fldata.spec(args.dataset)
        Xtr, ytr, Xte, yte = fldata.load_partition(
            client_index=args.client_index,
            num_clients=args.num_clients,
            dataset=args.dataset,
            alpha=args.alpha,
            seed=args.seed,
        )
        self.train_dl, self.test_dl = fltrainer.make_loaders(
            Xtr, ytr, Xte, yte, args.batch_size
        )
        self.model = flmodel.build_model(self.spec.num_classes, self.spec.n_channels)
        self.class_weights = (
            fltrainer.balanced_class_weights(ytr, self.spec.num_classes)
            if args.class_weights == "balanced" else None)
        print(f"### PROGRESS: [star] client {args.client_index} ready "
              f"(dataset={self.spec.name}, device={self.device.type}, train={len(ytr)}, "
              f"class_weights={args.class_weights}, "
              f"mix={fldata.class_distribution(ytr, self.spec)})")

    def get_parameters(self, config):
        return flmodel.get_weights(self.model)

    def fit(self, parameters, config):
        flmodel.set_weights(self.model, parameters)
        proximal_mu = float(config.get("proximal_mu", 0.0))
        n_seen = fltrainer.train_local(
            self.model, self.train_dl, self.device,
            epochs=self.args.local_epochs, lr=self.args.lr,
            proximal_mu=proximal_mu, class_weights=self.class_weights,
        )
        return flmodel.get_weights(self.model), n_seen, {}

    def evaluate(self, parameters, config):
        flmodel.set_weights(self.model, parameters)
        loss, m = fltrainer.evaluate(self.model, self.test_dl, self.device, self.spec)
        return float(loss), m["n_eval"], m


def _start(server: str, client: GridClient) -> None:
    """Start the client, tolerating both new and old Flower entrypoints."""
    try:
        fl.client.start_client(server_address=server, client=client.to_client())
    except AttributeError:
        # Older Flower: no to_client()/start_client
        fl.client.start_numpy_client(server_address=server, client=client)


def main() -> None:
    ap = argparse.ArgumentParser(description="Flower star client")
    ap.add_argument("--server", required=True, help="server_address host:port")
    ap.add_argument("--client-index", type=int, required=True)
    ap.add_argument("--num-clients", type=int, default=4)
    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("--prefer-gpu", type=int, default=1)
    ap.add_argument("--connect-retries", type=int, default=30)
    args = ap.parse_args()
    args.prefer_gpu = bool(args.prefer_gpu)

    client = GridClient(args)
    for attempt in range(1, args.connect_retries + 1):
        try:
            _start(args.server, client)
            break
        except Exception as exc:  # server not up yet / transient grpc error
            if attempt == args.connect_retries:
                raise
            print(f"### PROGRESS: [star] client {args.client_index} waiting for "
                  f"server ({attempt}/{args.connect_retries}): {exc}")
            time.sleep(5)
    print(f"### PROGRESS: [star] client {args.client_index} finished")


if __name__ == "__main__":
    main()
