import requests
import base64
import os
import tensorflow as tf
import numpy as np
import pickle
from io import BytesIO
import logging

import time
from cifer.rate_limiter import RateLimiter


logger = logging.getLogger("cifer.server")
logging.basicConfig(level=logging.INFO)

DEFAULT_TIMEOUT = 15

try:
    from phe import paillier
    PHE_AVAILABLE = True
except ImportError:
    PHE_AVAILABLE = False


class CiferServer:
    """Federated aggregation server.

    Additive capabilities (no behavior change by default):
    - structured logging + timing metrics
    - last_aggregation_summary
    - dry_run mode (skip upload) controlled by flag
    - verbose flag (keeps existing prints ON by default)
    """

    def __init__(
        self,
        encoded_project_id,
        encoded_company_id,
        encoded_client_id,
        base_api="https://workspace.cifer.ai/FederatedApi",
        dataset_path=None,
        model_path=None,
        use_encryption=False,
        timeout=DEFAULT_TIMEOUT,
        dry_run=False,
        verbose=True,

        # 🔐 Rate Limiting Config
        connection_limit=5,
        connection_window=10,
        message_limit=20,
        message_window=5,
    ):
        self.project_id = encoded_project_id
        self.company_id = encoded_company_id
        self.client_id = encoded_client_id
        self.base_api = base_api
        self.dataset_path = dataset_path
        self.model_path = model_path
        self.use_encryption = use_encryption
        self.timeout = timeout
        self.dry_run = dry_run
        self.verbose = verbose

        self.session = requests.Session()
        self.session.headers.update({"User-Agent": "Cifer-FL-Server/1.0"})

        # Additive: keep a summary of the last run for dashboards/audit
        self.last_aggregation_summary = None

        # -------------------------------------------------
        # 🔐 Rate Limiting (WebSocket Layer Protection)
        # -------------------------------------------------
        self.connection_limiter = RateLimiter(
            max_events=connection_limit,
            window_seconds=connection_window,
        )

        self.message_limiter = RateLimiter(
            max_events=message_limit,
            window_seconds=message_window,
        )

        logger.info(
            "Rate limiter initialized",
            extra={
                "connection_limit": connection_limit,
                "connection_window": connection_window,
                "message_limit": message_limit,
                "message_window": message_window,
            },
        )

        self._p(
            f"🚀 Server Initialized! Encryption: {'ENABLED' if self.use_encryption else 'DISABLED'}"
        )

        if self.use_encryption and not PHE_AVAILABLE:
            raise ImportError("Please install 'phe' for encrypted aggregation: pip install phe")

    # -----------------------------------------------------
    # Small helpers (additive, keeps default behavior)
    # -----------------------------------------------------
    def _p(self, msg: str) -> None:
        """Print helper. Default verbose=True preserves existing behavior."""
        if self.verbose:
            print(msg)

    def _set_last_summary(self, *, stage: str, clients: int, layers: int, encrypted: bool, time_sec: float):
        self.last_aggregation_summary = {
            "project_id": self.project_id,
            "company_id": self.company_id,
            "client_id": self.client_id,
            "stage": stage,
            "clients": clients,
            "layers": layers,
            "encrypted": encrypted,
            "time_sec": round(float(time_sec), 4),
            "ts": int(time.time()),
        }

    # -----------------------------------------------------
    # Model Loading
    # -----------------------------------------------------
    def load_model(self):
        if self.model_path and os.path.exists(self.model_path):
            self._p(f"✅ Loading Local Model: {self.model_path}")
            return tf.keras.models.load_model(self.model_path)

        self._p("🔄 No Local Model Found. Fetching from Clients...")
        return self.fetch_client_models()

    def fetch_client_models(self):
        url = f"{self.base_api}/get_client_models/{self.project_id}"
        response = self.session.get(url, timeout=self.timeout)
        response.raise_for_status()

        try:
            data = response.json()
            if data.get("status") == "success":
                return self.load_models(data.get("models", []))
            else:
                self._p("❌ ERROR: No models found for aggregation.")
                return None
        except Exception as e:
            self._p(f"❌ ERROR: {e}")
            return None

    def fetch_encrypted_models(self):
        url = f"{self.base_api}/get_client_models/{self.project_id}"
        response = self.session.get(url, timeout=self.timeout)
        response.raise_for_status()

        try:
            data = response.json()
            if data.get("status") == "success":
                models = []
                for m in data["models"]:
                    binary = base64.b64decode(m["model_data"])
                    encrypted_weights = pickle.loads(binary)
                    models.append(encrypted_weights)
                return models
            else:
                self._p("❌ ERROR: No encrypted models found.")
                return []
        except Exception as e:
            self._p(f"❌ ERROR loading encrypted models: {e}")
            return []

    def load_models(self, model_data_list):
        models = []
        for i, model_info in enumerate(model_data_list):
            try:
                model_bytes = base64.b64decode(model_info["model_data"])
                bio = BytesIO(model_bytes)
                model = tf.keras.models.load_model(bio)
                models.append(model)
            except Exception as e:
                self._p(f"❌ ERROR: Failed to load model {i} - {e}")
        return models

    # -----------------------------------------------------
    # Validation
    # -----------------------------------------------------
    def _validate_models(self, models):
        ref_weights = models[0].get_weights()
        for i, m in enumerate(models[1:], start=1):
            w = m.get_weights()
            if len(w) != len(ref_weights):
                raise ValueError(f"Model {i} layer count mismatch")
            for a, b in zip(w, ref_weights):
                if a.shape != b.shape:
                    raise ValueError("Model weight shape mismatch")

    # -----------------------------------------------------
    # Aggregation
    # -----------------------------------------------------
    def fed_avg(self, models):
        self._p("🔄 Performing FedAvg Aggregation...")
        if not models:
            self._p("❌ ERROR: No models to aggregate.")
            return None

        start_time = time.time()
        weights = [model.get_weights() for model in models]
        avg_weights = [np.mean(w, axis=0) for w in zip(*weights)]
        models[0].set_weights(avg_weights)
        elapsed = time.time() - start_time

        layers = len(models[0].get_weights())
        logger.info(
            "FedAvg completed",
            extra={
                "project_id": self.project_id,
                "clients": len(models),
                "layers": layers,
                "time_sec": round(elapsed, 4),
                "encrypted": False,
            },
        )
        self._set_last_summary(stage="fed_avg", clients=len(models), layers=layers, encrypted=False, time_sec=elapsed)
        return models[0]

    def encrypted_fed_avg(self, encrypted_models):
        self._p("🔐 Performing Encrypted FedAvg Aggregation...")

        if not encrypted_models:
            self._p("❌ ERROR: No encrypted weights received.")
            return None

        start_time = time.time()
        averaged = []
        for layers in zip(*encrypted_models):
            sum_layer = []
            for weights in zip(*layers):
                summed = weights[0]
                for w in weights[1:]:
                    summed += w
                sum_layer.append(summed / len(encrypted_models))
            averaged.append(sum_layer)

        elapsed = time.time() - start_time
        logger.info(
            "Encrypted FedAvg completed",
            extra={
                "project_id": self.project_id,
                "clients": len(encrypted_models),
                "layers": len(averaged),
                "time_sec": round(elapsed, 4),
                "encrypted": True,
            },
        )
        self._set_last_summary(
            stage="encrypted_fed_avg",
            clients=len(encrypted_models),
            layers=len(averaged),
            encrypted=True,
            time_sec=elapsed,
        )

        self._p("✅ Encrypted aggregation complete (still encrypted)")
        return averaged

    # -----------------------------------------------------
    # Upload
    # -----------------------------------------------------
    def upload_aggregated_model(self, model):
        if self.dry_run:
            # Additive capability: skip upload when explicitly enabled
            self._p("🧪 Dry-run enabled: skipping upload_aggregated_model")
            logger.info(
                "Dry-run: upload skipped",
                extra={"project_id": self.project_id},
            )
            return

        filename = "aggregated_model.h5"
        model.save(filename)

        with open(filename, "rb") as f:
            model_data = f.read()

        files = {"aggregated_model": (filename, model_data)}
        data = {
            "project_id": self.project_id,
            "aggregation_method": "FedAvg",
            # Keep original behavior to avoid downstream changes
            "num_clients": 1,
            "weight_statistics": '{"accuracy": null}',
        }

        api_url = f"{self.base_api}/upload_aggregated_model"
        self._p(f"📡 Uploading aggregated model to {api_url}...")
        start_time = time.time()
        response = self.session.post(api_url, files=files, data=data, timeout=self.timeout)

        if response.status_code == 200:
            self._p("✅ Aggregated model uploaded successfully!")
            elapsed = time.time() - start_time
            logger.info(
                "Aggregated model uploaded",
                extra={
                    "project_id": self.project_id,
                    "clients": 1,
                    "time_sec": round(elapsed, 4),
                },
            )
        else:
            self._p(f"❌ Upload failed: {response.text}")

    # -----------------------------------------------------
    # Run
    # -----------------------------------------------------
    def run(self):
        self._p("✅ Server is running...")

        if self.use_encryption:
            encrypted_models = self.fetch_encrypted_models()
            if not encrypted_models:
                self._p("❌ No encrypted models to aggregate.")
                return

            aggregated_weights = self.encrypted_fed_avg(encrypted_models)

            filename = "aggregated_encrypted_weights.pkl"
            with open(filename, "wb") as f:
                pickle.dump(aggregated_weights, f)

            self._p(f"📦 Encrypted aggregated weights saved to {filename}")
            self._p("🛑 Waiting for client to decrypt and use aggregated model.")
            return

        models = self.load_model()
        if not models:
            self._p("❌ No model available for aggregation.")
            return

        self._validate_models(models)
        aggregated_model = self.fed_avg(models)
        if aggregated_model:
            self.upload_aggregated_model(aggregated_model)
