from typing import Any, Dict, Optional, Type
import flwr as fl
import optuna


class Strategy:
    def __init__(
        self,
        name: str = "FedAvg",
        fraction_fit: float = 1.0,
        fraction_evaluate: float = 1.0,
        min_fit_clients: int = 2,
        min_evaluate_clients: int = 2,
        min_available_clients: int = 2,
        initial_parameters=None,
        evaluation_methode: str = "centralized",
        strategy_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        self.fraction_fit = fraction_fit
        self.fraction_evaluate = fraction_evaluate
        self.min_fit_clients = min_fit_clients
        self.min_evaluate_clients = min_evaluate_clients
        self.min_available_clients = min_available_clients
        self.initial_parameters = initial_parameters if initial_parameters is not None else []
        self.evaluate_fn = None
        self.name = name
        self.evaluation_methode = evaluation_methode
        self.strategy_kwargs = strategy_kwargs or {}

    def optuna_fed_optimization(self, direction: str, hpo_rate: int, params_config):
        self.study = optuna.create_study(direction=direction)
        self.hpo_rate = hpo_rate
        self.params_config = params_config

    def get_strategy_by_name(self):
        try:
            return getattr(fl.server.strategy, self.name)
        except AttributeError as e:
            available = [n for n in dir(fl.server.strategy) if n and n[0].isupper()]
            raise ValueError(f"Unknown strategy '{self.name}'. Available: {available}") from e

    def create_strategy(self):
        BaseStrategy: Type[fl.server.strategy.Strategy] = self.get_strategy_by_name()

        # ✅ Define a wrapper strategy that stores per-client metrics
        class TrackingStrategy(BaseStrategy):  # type: ignore[misc]
            def __init__(self, *args, **kwargs):
                super().__init__(*args, **kwargs)
                self.client_fit_history: Dict[str, list] = {}   # cid -> list of metrics dict per round
                self.client_eval_history: Dict[str, list] = {}  # cid -> list of metrics dict per round

            def aggregate_fit(self, server_round, results, failures):
                # results: List[Tuple[ClientProxy, FitRes]]
                for client_proxy, fit_res in results:
                    cid = getattr(client_proxy, "cid", "unknown")
                    self.client_fit_history.setdefault(cid, []).append(dict(fit_res.metrics or {}))
                return super().aggregate_fit(server_round, results, failures)

            def aggregate_evaluate(self, server_round, results, failures):
                # results: List[Tuple[ClientProxy, EvaluateRes]]
                for client_proxy, eval_res in results:
                    cid = getattr(client_proxy, "cid", "unknown")
                    self.client_eval_history.setdefault(cid, []).append(dict(eval_res.metrics or {}))
                return super().aggregate_evaluate(server_round, results, failures)

        kwargs = dict(
            fraction_fit=self.fraction_fit,
            fraction_evaluate=self.fraction_evaluate,
            min_fit_clients=self.min_fit_clients,
            min_evaluate_clients=self.min_evaluate_clients,
            min_available_clients=self.min_available_clients,
        )

        if self.initial_parameters is not None and len(self.initial_parameters) != 0:
            kwargs["initial_parameters"] = fl.common.ndarrays_to_parameters(self.initial_parameters)

        if self.evaluate_fn is not None:
            kwargs["evaluate_fn"] = self.evaluate_fn

        kwargs.update(self.strategy_kwargs)

        self.strategy_object = TrackingStrategy(**kwargs)
        return self.strategy_object
