"""
WebSocket Handler Module - Handles WebSocket functionality for the application
"""

import json
import logging
import time
from typing import Optional

from flask import Flask, request
from flask_sock import Sock
from flask_cors import cross_origin

from ressources.blueprint.riedl_blueprint import get_riedls, get_diff

logger = logging.getLogger("median.websocket_handler")


class WebSocketHandler:
    """
    WebSocket Handler class to centralize WebSocket functionality
    """

    def __init__(self):
        self.sock = None
        self.aideplus_clients = []  # Store active AidePlus WebSocket connections

    def init_app(self, app: Flask) -> Sock:
        logger.info("Initializing WebSocket handler")

        app.config["SOCK_SERVER_OPTIONS"] = {
            "ping_interval": 25,
            "ping_timeout": 60,
            "max_message_size": None,  # No limit on message size
        }

        self.sock = Sock(app)

        # Health check endpoint to test connectivity
        @app.route("/health-check", methods=["GET", "OPTIONS"])
        @cross_origin(origins="*", allow_headers=["Content-Type", "Authorization"])
        def health_check():
            """Health check endpoint to verify server is accessible"""
            logger.info(f"Health check request received from: {request.remote_addr}")
            return {"status": "ok", "service": "aideplus-websocket", "timestamp": time.time()}

        # Register WebSocket routes
        self._register_routes()

        logger.info("WebSocket server initialized")

        return self.sock

    def get_sock(self) -> Optional[Sock]:
        return self.sock

    def _register_routes(self):
        """Register all WebSocket routes"""
        if self.sock is None:
            logger.error("Cannot register WebSocket routes: Sock not initialized")
            return

        # Register the Riedl WebSocket route
        @self.sock.route("/ws/riedl")
        def echo(ws):
            try:
                old_rields = get_riedls()
                while 1:
                    new_rields = get_riedls()
                    changes = []

                    for new in new_rields:
                        old = list(filter(lambda x: (x["pk"] == new["pk"]), old_rields))
                        obj = {}
                        if len(old) == 1:
                            obj = get_diff(new, old[0])

                        if len(obj.keys()) > 0:
                            changes.append(obj)

                        if len(changes) > 0:
                            ws.send(data=json.dumps({"data": changes}))
                            old_rields = new_rields

                    time.sleep(30)
            except Exception as e:
                logger.error(f"Error in WebSocket connection: {str(e)}")
                pass
            return ""

        # Register the AidePlus WebSocket route
        @self.sock.route("/ws/aideplus")
        def aideplus_connection(ws):
            try:
                client_info = {"address": request.remote_addr, "time": time.time()}
                self.aideplus_clients.append(ws)
                logger.info(f"New AidePlus WebSocket client connected from: {client_info['address']}")

                # Send initial connection confirmation
                try:
                    ws.send(
                        json.dumps(
                            {
                                "type": "connection",
                                "data": {
                                    "status": "connected",
                                    "server_time": time.time(),
                                    "message": "WebSocket connection established",
                                },
                            }
                        )
                    )
                except Exception as e:
                    logger.error(f"Error sending initial confirmation: {str(e)}")

                # Keep connection alive
                while True:
                    try:
                        message = ws.receive(timeout=30)
                        if message:
                            # Process the client message
                            try:
                                msg_data = json.loads(message)

                                if msg_data.get("action") == "publish":
                                    topic = msg_data.get("topic")
                                    payload = msg_data.get("message")

                                    if topic and payload:

                                        # Create a simpler message structure
                                        mqtt_message = {
                                            "topic": topic,
                                            "device": "server",
                                            "payload": payload,
                                            "timestamp": time.time(),
                                        }

                                        # Broadcast as MQTT message
                                        self.broadcast_to_aideplus("mqtt", mqtt_message)

                                        # Confirm receipt
                                        ws.send(
                                            json.dumps(
                                                {"type": "publish_response", "status": "success", "topic": topic}
                                            )
                                        )

                                    else:
                                        logger.warning("Invalid publish message: missing topic or payload")
                                        ws.send(
                                            json.dumps(
                                                {
                                                    "type": "error",
                                                    "message": "Invalid publish message: missing topic or payload",
                                                }
                                            )
                                        )
                                else:
                                    logger.debug(f"Received non-publish message: {msg_data.get('action', 'unknown')}")
                            except json.JSONDecodeError:
                                logger.warning("Received invalid JSON from client")
                            except Exception as e:
                                logger.error(f"Error processing client message: {str(e)}", exc_info=True)
                    except Exception:
                        # Just a timeout, keep connection alive
                        pass

                    # Send heartbeat every 30 seconds
                    try:
                        ws.send(
                            json.dumps(
                                {
                                    "type": "heartbeat",
                                    "data": {"timestamp": time.time()},
                                }
                            )
                        )
                    except Exception as e:
                        logger.warning(f"Failed to send heartbeat: {str(e)}")
                        break  # Client likely disconnected

            except Exception as e:
                logger.error(f"Error in WebSocket connection: {str(e)}")
            finally:
                # Unregister the client connection on disconnect
                if ws in self.aideplus_clients:
                    self.aideplus_clients.remove(ws)
                    logger.info(
                        f"AidePlus WebSocket client disconnected:"
                        f"{client_info['address'] if 'client_info' in locals() else 'unknown'}"
                    )

            return ""

    def broadcast_to_aideplus(self, message_type, data):
        """
        Broadcast a message to all connected AidePlus WebSocket clients
        """
        if not self.aideplus_clients:
            # No active AidePlus WebSocket clients to broadcast to
            return

        message = json.dumps({"type": message_type, "data": data, "timestamp": time.time()})

        # Send to all connected clients
        disconnected_clients = []
        for client in self.aideplus_clients:
            try:
                client.send(data=message)
            except Exception as e:
                logger.error(f"Error sending to WebSocket client: {str(e)}")
                disconnected_clients.append(client)

        # Remove disconnected clients
        for client in disconnected_clients:
            if client in self.aideplus_clients:
                self.aideplus_clients.remove(client)


# Create a singleton instance
websocket_handler = WebSocketHandler()


def init_websocket(app: Flask) -> Sock:
    """
    Initialize WebSocket with the Flask app

    Args:
        app: Flask application instance

    Returns:
        Sock instance
    """
    return websocket_handler.init_app(app)
