import asyncio
import websockets
import json
import ssl
import numpy as np
import jwt
import os
import logging
import secrets
from pathlib import Path
from dotenv import load_dotenv
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
import base64

# =====================================================
# Logging
# =====================================================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s"
)
logger = logging.getLogger("cifer_federated_server")

# =====================================================
# ENV SETUP (Auto-create .env if missing)
# =====================================================
ENV_PATH = Path(__file__).resolve().parent.parent / ".env"

if not ENV_PATH.exists():
    strong_secret = secrets.token_hex(32)  # 64 hex chars (32 bytes)
    ENV_PATH.write_text(f"CIFER_JWT_SECRET={strong_secret}\n")
    logger.warning(".env auto-created with secure CIFER_JWT_SECRET")

load_dotenv(dotenv_path=ENV_PATH)

SECRET_KEY = os.getenv("CIFER_JWT_SECRET")
if not SECRET_KEY or len(SECRET_KEY.encode()) < 32:
    raise RuntimeError("CIFER_JWT_SECRET must be at least 32 bytes")

ROUND_TIMEOUT = int(os.getenv("CIFER_ROUND_TIMEOUT", "60"))
TRAIN_ON_SERVER = os.getenv("CIFER_SERVER_TRAIN", "0") in ("1", "true", "True")

# =====================================================
# Federated Server
# =====================================================
class FederatedServer:
    def __init__(
        self,
        ip_address="127.0.0.1",
        port=8765,
        n_round_clients=3,
        use_homomorphic=False,
        use_secure=False,
        certfile="certificate.pem",
        keyfile="private_key.pem"
    ):
        self.ip_address = ip_address
        self.port = port
        self.n_round_clients = n_round_clients
        self.use_homomorphic = use_homomorphic
        self.use_secure = use_secure
        self.connected_clients = set()
        self.round_weights = {}
        self.round_start_time = {}
        self.custom_model = None

        if self.use_secure:
            self.ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
            self.ssl_context.load_cert_chain(certfile=certfile, keyfile=keyfile)
        else:
            self.ssl_context = None

        if self.use_homomorphic:
            self.private_key = RSA.generate(3072)
            self.public_key = self.private_key.publickey()

    # =================================================
    # Authentication
    # =================================================
    async def authenticate_client(self, websocket):
        try:
            token = await websocket.recv()
            payload = jwt.decode(
                token,
                SECRET_KEY,
                algorithms=["HS256"],
                options={"require": ["exp"]}
            )
            logger.info(f"Authenticated user: {payload.get('user_id')}")
            return True
        except Exception as e:
            logger.warning(f"Authentication failed: {e}")
            await websocket.close()
            return False

    # =================================================
    # Message Handler
    # =================================================
    async def receive_message(self, websocket):
        if not await self.authenticate_client(websocket):
            return

        self.connected_clients.add(websocket)

        try:
            async for message in websocket:
                data = json.loads(message)

                if data.get("type") == "request_public_key" and self.use_homomorphic:
                    await websocket.send(self.public_key.export_key().decode())

                elif data.get("type") == "send_model":
                    await self.handle_model(data, websocket)

        except websockets.exceptions.ConnectionClosed:
            logger.warning("Client disconnected")
        finally:
            self.connected_clients.discard(websocket)

    # =================================================
    # Handle Model Aggregation
    # =================================================
    async def handle_model(self, data, websocket):
        try:
            round_id = data.get("round_id")
            if not round_id:
                await websocket.send(json.dumps({"status": "error", "message": "round_id missing"}))
                return

            current_time = asyncio.get_event_loop().time()

            if round_id not in self.round_start_time:
                self.round_start_time[round_id] = current_time

            if current_time - self.round_start_time[round_id] > ROUND_TIMEOUT:
                logger.warning(f"Round {round_id} timeout")
                self.round_weights.pop(round_id, None)
                self.round_start_time.pop(round_id, None)
                return

            if self.use_homomorphic:
                cipher = PKCS1_OAEP.new(self.private_key)
                weights = [
                    np.frombuffer(
                        cipher.decrypt(base64.b64decode(w)),
                        dtype=np.float32
                    )
                    for w in data["weights"]
                ]
            else:
                weights = [np.array(w) for w in data["weights"]]

            self.round_weights.setdefault(round_id, []).append(weights)

            if len(self.round_weights[round_id]) == self.n_round_clients:
                logger.info(f"Aggregating round {round_id}")

                group = self.round_weights[round_id]
                layers = len(group[0])

                aggregated = [
                    np.mean([client[i] for client in group], axis=0)
                    for i in range(layers)
                ]

                self.round_weights.pop(round_id, None)
                self.round_start_time.pop(round_id, None)

                for client in list(self.connected_clients):
                    if client.open:
                        await client.send(json.dumps({
                            "status": "success",
                            "round_id": round_id,
                            "message": "Aggregation complete"
                        }))

        except Exception as e:
            logger.exception(f"Aggregation failed: {e}")
            try:
                await websocket.send(json.dumps({
                    "status": "error",
                    "message": str(e)
                }))
            except:
                pass

    # =================================================
    # Start Server
    # =================================================
    async def start_server(self):
        protocol = "wss" if self.use_secure else "ws"
        async with websockets.serve(
            self.receive_message,
            self.ip_address,
            self.port,
            ssl=self.ssl_context,
            max_size=10 * 1024 * 1024
        ):
            logger.info("======================================")
            logger.info(f"🚀 Cifer Federated Server STARTED")
            logger.info(f"Protocol : {protocol}")
            logger.info(f"Address  : {protocol}://{self.ip_address}:{self.port}")
            logger.info(f"Clients  : waiting for {self.n_round_clients}")
            logger.info("======================================")
            print(f"\n🚀 Cifer Federated Server running at {protocol}://{self.ip_address}:{self.port}\n")
            await asyncio.Future()

# =====================================================
# Runner
# =====================================================
def run_federated_server(
    ip_address="127.0.0.1",
    port=8765,
    n_round_clients=3,
    custom_model=None,   # backward compatibility
    use_homomorphic=False,
    use_secure=False
):
    server = FederatedServer(
        ip_address=ip_address,
        port=port,
        n_round_clients=n_round_clients,
        use_homomorphic=use_homomorphic,
        use_secure=use_secure
    )

    # keep compatibility with older runner scripts
    if custom_model is not None:
        server.custom_model = custom_model

    logger.info("Launching Federated Server...")
    asyncio.run(server.start_server())
