"""
Federated Learning adapter layer for model_lib.

Design goals
------------
* Decoupled from any specific FL framework — the FL framework imports *this*,
  not the other way around.
* Works over SQS (messages carry serialised model parameter diffs).
* Supports heterogeneous model registries (not just WastePredictor).
* Neural-network parameters use FedAvg; tree-based models use quality-weighted
  model selection (trees are not additively averageable).

Architecture
------------
                        ┌──────────────────────┐
                        │  FederatedModelAdapter│  ← abstract contract
                        │  (ABC)               │
                        └──────────┬───────────┘
                                   │
                 ┌─────────────────┼─────────────────┐
                 │                 │                 │
     WastePredictorFLAdapter   (your next model)  …
                 │
         ┌───────┴────────┐
         │ get_params()   │  → dict[str, np.ndarray]  (for SQS payload)
         │ set_params()   │  ← dict[str, np.ndarray]
         │ local_train()  │  trains on local DataFrame, returns updated params
         │ serialize()    │  → bytes  (full model, for server aggregation)
         │ deserialize()  │  ← bytes
         └────────────────┘

SQS message contract
--------------------
Each SQS message body is JSON with fields:
    {
        "event":        "FL_PARAMS_UPDATE" | "FL_MODEL_BROADCAST",
        "model_id":     "waste_predictor_v4",
        "client_id":    "<uuid>",
        "round":        <int>,
        "num_samples":  <int>,          # for weighted FedAvg
        "params":       "<base64>",     # base64(pickle(dict[str, np.ndarray]))
        "metrics":      { "r2": … }
    }
"""
from __future__ import annotations

import abc
import base64
import io
import pickle
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Dict, List, Optional, Type

import numpy as np
import pandas as pd


# ── SQS message schema ────────────────────────────────────────────────────────

class FLEventType(str, Enum):
    PARAMS_UPDATE = "FL_PARAMS_UPDATE"        # client → server
    MODEL_BROADCAST = "FL_MODEL_BROADCAST"    # server → clients
    ROUND_START = "FL_ROUND_START"            # server → clients (trigger training)
    ROUND_COMPLETE = "FL_ROUND_COMPLETE"      # server → orchestrator


@dataclass
class FLMessage:
    """Typed wrapper for every FL SQS message."""

    event: FLEventType
    model_id: str
    client_id: str
    round: int
    num_samples: int
    params: Optional[Dict[str, np.ndarray]] = None
    metrics: Dict[str, float] = field(default_factory=dict)
    extra: Dict[str, Any] = field(default_factory=dict)

    # ── Serialisation to / from SQS JSON body ────────────────────────

    def to_sqs_body(self) -> str:
        import json

        payload = {
            "event": self.event.value,
            "model_id": self.model_id,
            "client_id": self.client_id,
            "round": self.round,
            "num_samples": self.num_samples,
            "metrics": self.metrics,
            "extra": self.extra,
        }
        if self.params is not None:
            payload["params"] = base64.b64encode(
                pickle.dumps(self.params, protocol=4)
            ).decode("ascii")
        return json.dumps(payload)

    @classmethod
    def from_sqs_body(cls, body: str) -> "FLMessage":
        import json

        d = json.loads(body)
        params: Optional[Dict[str, np.ndarray]] = None
        if "params" in d:
            params = pickle.loads(base64.b64decode(d["params"]))

        return cls(
            event=FLEventType(d["event"]),
            model_id=d["model_id"],
            client_id=d["client_id"],
            round=int(d["round"]),
            num_samples=int(d["num_samples"]),
            params=params,
            metrics=d.get("metrics", {}),
            extra=d.get("extra", {}),
        )


# ── Abstract adapter ──────────────────────────────────────────────────────────

class FederatedModelAdapter(abc.ABC):
    """
    Abstract contract every model must implement to participate in FL via SQS.

    Concrete subclasses are registered in :data:`MODEL_REGISTRY` so the
    FL framework can instantiate the right adapter from a model_id string.
    """

    #: Unique stable identifier for this model type (used in SQS messages).
    MODEL_ID: str = ""

    # ── Parameter exchange ────────────────────────────────────────────

    @abc.abstractmethod
    def get_params(self) -> Dict[str, np.ndarray]:
        """
        Return the trainable parameters as a flat dict of numpy arrays.

        For neural networks: PyTorch state-dict tensors → numpy.
        For tree models: serialised bytes wrapped in a 1-D uint8 array.
        """

    @abc.abstractmethod
    def set_params(self, params: Dict[str, np.ndarray]) -> None:
        """Apply parameter dict received from the aggregation server."""

    # ── Local training ────────────────────────────────────────────────

    @abc.abstractmethod
    def local_train(
        self,
        train_df: pd.DataFrame,
        *,
        epochs: int = 50,
        verbose: bool = False,
    ) -> Dict[str, np.ndarray]:
        """
        Train on local data and return updated parameters.

        Returns the *same* structure as ``get_params()``.
        """

    # ── Full-model serialisation (for server-side aggregation) ────────

    @abc.abstractmethod
    def serialize(self) -> bytes:
        """Serialise the full model to bytes."""

    @classmethod
    @abc.abstractmethod
    def deserialize(cls, data: bytes) -> "FederatedModelAdapter":
        """Reconstruct from bytes produced by :meth:`serialize`."""

    # ── FedAvg aggregation (server-side, static) ──────────────────────

    @staticmethod
    @abc.abstractmethod
    def aggregate(
        params_list: List[Dict[str, np.ndarray]],
        weights: Optional[List[float]] = None,
    ) -> Dict[str, np.ndarray]:
        """
        Aggregate a list of parameter dicts into one.

        Parameters
        ----------
        params_list:
            One dict per client.
        weights:
            Per-client weights (e.g. num_samples).  Uniform if None.
        """

    # ── Metrics ───────────────────────────────────────────────────────

    @abc.abstractmethod
    def evaluate(self, test_df: pd.DataFrame) -> Dict[str, float]:
        """Return evaluation metrics on an held-out DataFrame."""

    # ── SQS helper ───────────────────────────────────────────────────

    def build_update_message(
        self,
        params: Dict[str, np.ndarray],
        round_: int,
        num_samples: int,
        client_id: Optional[str] = None,
        metrics: Optional[Dict[str, float]] = None,
    ) -> FLMessage:
        return FLMessage(
            event=FLEventType.PARAMS_UPDATE,
            model_id=self.MODEL_ID,
            client_id=client_id or str(uuid.uuid4()),
            round=round_,
            num_samples=num_samples,
            params=params,
            metrics=metrics or {},
        )


# ── Model registry ────────────────────────────────────────────────────────────

MODEL_REGISTRY: Dict[str, Type[FederatedModelAdapter]] = {}


def register_model(cls: Type[FederatedModelAdapter]) -> Type[FederatedModelAdapter]:
    """Class decorator that registers a FederatedModelAdapter subclass."""
    if not cls.MODEL_ID:
        raise ValueError(f"{cls.__name__} must define MODEL_ID")
    MODEL_REGISTRY[cls.MODEL_ID] = cls
    return cls


def get_adapter(model_id: str) -> Type[FederatedModelAdapter]:
    try:
        return MODEL_REGISTRY[model_id]
    except KeyError:
        raise KeyError(
            f"No FL adapter registered for model_id={model_id!r}. "
            f"Available: {list(MODEL_REGISTRY)}"
        )


# ── WastePredictor FL adapter ─────────────────────────────────────────────────

@register_model
class WastePredictorFLAdapter(FederatedModelAdapter):
    """
    Federated Learning adapter for the WastePredictor V4 ensemble.

    Parameter strategy
    ------------------
    * **Neural network** (``NeuralNetworkTrainer``):
      Full PyTorch ``state_dict`` — tensors → ``float32`` numpy arrays.
      FedAvg applies weighted averaging.

    * **Tree models** (``GradientBoostingWasteModel``, ``StackedEnsembleModel``):
      Serialised as ``uint8`` byte arrays.  FedAvg uses *quality-weighted model
      selection* (best R² client wins — trees cannot be additively averaged).

    Aggregation  (server calls ``WastePredictorFLAdapter.aggregate()``)
    --------------------------------------------------------------------------
    1. Average NeuralNetwork weights across clients (standard FedAvg).
    2. For each tree component, select the client with the highest reported R².
    """

    MODEL_ID = "waste_predictor_v4"

    def __init__(self, model_path: Optional[str] = None) -> None:
        from model_lib.predictor import WastePredictor

        self._predictor = WastePredictor(model_path)
        self._predictor._ensure_loaded()

    # ── Helpers ───────────────────────────────────────────────────────

    def _nn_trainer(self):
        return self._predictor.model.models.get("neural_network")

    def _gb_model(self):
        return self._predictor.model.models.get("gradient_boosting")

    def _stacked_model(self):
        return self._predictor.model.models.get("stacked")

    # ── Parameter exchange ────────────────────────────────────────────

    def get_params(self) -> Dict[str, np.ndarray]:
        params: Dict[str, np.ndarray] = {}

        # Neural network weights
        nn = self._nn_trainer()
        if nn is not None and hasattr(nn, "model"):
            for k, tensor in nn.model.state_dict().items():
                params[f"nn::{k}"] = tensor.cpu().float().numpy()

        # Tree models as pickled bytes
        gb = self._gb_model()
        if gb is not None:
            params["tree::gradient_boosting"] = np.frombuffer(
                pickle.dumps(gb, protocol=4), dtype=np.uint8
            )

        stacked = self._stacked_model()
        if stacked is not None:
            params["tree::stacked"] = np.frombuffer(
                pickle.dumps(stacked, protocol=4), dtype=np.uint8
            )

        return params

    def set_params(self, params: Dict[str, np.ndarray]) -> None:
        import torch

        nn = self._nn_trainer()
        if nn is not None and hasattr(nn, "model"):
            nn_keys = {k[len("nn::"):]: v for k, v in params.items() if k.startswith("nn::")}
            if nn_keys:
                state_dict = {
                    k: torch.tensor(v).to(nn.device) for k, v in nn_keys.items()
                }
                nn.model.load_state_dict(state_dict, strict=True)

        if "tree::gradient_boosting" in params:
            from model_lib._core.loader import load_raw_bytes
            gb = load_raw_bytes(params["tree::gradient_boosting"].tobytes())
            self._predictor.model.models["gradient_boosting"] = gb

        if "tree::stacked" in params:
            from model_lib._core.loader import load_raw_bytes
            stacked = load_raw_bytes(params["tree::stacked"].tobytes())
            self._predictor.model.models["stacked"] = stacked

    # ── Local training ────────────────────────────────────────────────

    def local_train(
        self,
        train_df: pd.DataFrame,
        *,
        epochs: int = 50,
        verbose: bool = False,
    ) -> Dict[str, np.ndarray]:
        """
        Fine-tune the neural network component on local data.
        Tree models are *re-trained* from scratch on local data.
        """
        from model_lib._core.architecture import (
            GradientBoostingWasteModel,
            NeuralNetworkTrainer,
        )

        # Retrain NN with fewer epochs (local adaptation)
        nn = NeuralNetworkTrainer(epochs=epochs, lr=5e-4, patience=max(10, epochs // 5))
        nn.fit(train_df, verbose=verbose)
        self._predictor.model.models["neural_network"] = nn

        # Retrain fast GB
        gb = GradientBoostingWasteModel(n_estimators=200, learning_rate=0.05, max_depth=5)
        gb.fit(train_df, verbose=False)
        self._predictor.model.models["gradient_boosting"] = gb

        return self.get_params()

    # ── Full-model serialisation ──────────────────────────────────────

    def serialize(self) -> bytes:
        buf = io.BytesIO()
        pickle.dump(
            {
                "models": self._predictor.model.models,
                "weights": self._predictor.model.weights,
                "feature_names": self._predictor.model.feature_engineer.feature_names,
                "output_cols": self._predictor.model.output_cols,
                "metrics": self._predictor.model.test_metrics,
            },
            buf,
            protocol=4,
        )
        return buf.getvalue()

    @classmethod
    def deserialize(cls, data: bytes) -> "WastePredictorFLAdapter":
        from model_lib._core.loader import load_raw_bytes

        raw = load_raw_bytes(data)
        adapter = cls.__new__(cls)
        from model_lib.predictor import WastePredictor
        from model_lib._core.architecture import AdvancedFeatureEngineer, ProductionWastePredictor

        predictor = ProductionWastePredictor.__new__(ProductionWastePredictor)
        predictor.models = raw["models"]
        predictor.weights = raw["weights"]
        predictor.output_cols = raw.get("output_cols", AdvancedFeatureEngineer.OUTPUT_COLS)
        predictor.feature_engineer = AdvancedFeatureEngineer()
        predictor.feature_engineer.feature_names = raw.get("feature_names", [])
        predictor.test_metrics = raw.get("metrics", {})

        wp = WastePredictor.__new__(WastePredictor)
        import threading
        wp._model = predictor
        wp._model_path = None
        wp._lock = threading.Lock()
        adapter._predictor = wp
        return adapter

    # ── FedAvg aggregation ────────────────────────────────────────────

    @staticmethod
    def aggregate(
        params_list: List[Dict[str, np.ndarray]],
        weights: Optional[List[float]] = None,
    ) -> Dict[str, np.ndarray]:
        if not params_list:
            raise ValueError("params_list is empty")

        n = len(params_list)
        w = np.array(weights, dtype=float) if weights else np.ones(n)
        w = w / w.sum()

        aggregated: Dict[str, np.ndarray] = {}
        all_keys = set(params_list[0].keys())

        for key in all_keys:
            if key.startswith("nn::"):
                # Weighted FedAvg on NN parameters
                aggregated[key] = sum(
                    w[i] * params_list[i][key].astype(float)
                    for i in range(n)
                ).astype(np.float32)

            elif key.startswith("tree::"):
                # Quality-weighted model selection: best client wins
                best_idx = int(np.argmax(w))
                aggregated[key] = params_list[best_idx][key]

        return aggregated

    # ── Metrics ───────────────────────────────────────────────────────

    def evaluate(self, test_df: pd.DataFrame) -> Dict[str, float]:
        from sklearn.metrics import mean_absolute_error, r2_score, mean_squared_error

        output_cols = self._predictor.OUTPUT_COLS
        y_true = test_df[output_cols].values
        y_pred = self._predictor.predict_batch(test_df).values
        r2_per = r2_score(y_true, y_pred, multioutput="raw_values")
        return {
            "r2": float(r2_score(y_true, y_pred)),
            "mae": float(mean_absolute_error(y_true, y_pred)),
            "rmse": float(np.sqrt(mean_squared_error(y_true, y_pred))),
            "r2_per_target": dict(zip(output_cols, r2_per.tolist())),
        }
