import asyncio
import inspect
import logging
import threading
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, Callable
from fastapi import WebSocket, WebSocketDisconnect
from fastapi.websockets import WebSocketState
from .config_manager import ConfigManager
import importlib
import json
from collections import defaultdict
from uuid import uuid4

logger = logging.getLogger(__name__)

class ConnectionManager:
    _instance = None
    _lock = threading.Lock()

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._init()
        return cls._instance

    def _init(self):
        self.active_connections: Dict[str, WebSocket] = {}
        self.rooms: Dict[str, Dict[str, WebSocket]] = defaultdict(dict)
        self.namespaces: Dict[str, Dict[str, WebSocket]] = defaultdict(dict)
        self.event_handlers: Dict[str, Dict[str, Callable]] = defaultdict(dict)
        self._config = ConfigManager()
        self._prefixes = {
            'event': {},
            'room': {},
            'namespace': {}
        }

    async def connect(self, websocket: WebSocket, client_id: str, namespace: str = "/"):
        await websocket.accept()
        self.active_connections[client_id] = websocket
        self.namespaces[namespace][client_id] = websocket
        return client_id

    def disconnect(self, client_id: str, namespace: str = "/"):
        if client_id in self.active_connections:
            del self.active_connections[client_id]
        if namespace in self.namespaces and client_id in self.namespaces[namespace]:
            del self.namespaces[namespace][client_id]
        
        # Remove from all rooms
        for room in list(self.rooms.keys()):
            if client_id in self.rooms[room]:
                del self.rooms[room][client_id]

    async def send_personal_message(self, message: str, client_id: str):
        if client_id in self.active_connections:
            websocket = self.active_connections[client_id]
            if websocket.client_state == WebSocketState.CONNECTED:
                await websocket.send_text(message)

    async def broadcast(self, message: str, namespace: str = "/"):
        for client_id, websocket in self.namespaces[namespace].items():
            if websocket.client_state == WebSocketState.CONNECTED:
                await websocket.send_text(message)

    async def send_to_room(self, room: str, message: str, namespace: str = "/"):
        for client_id, websocket in self.rooms[room].items():
            if (websocket.client_state == WebSocketState.CONNECTED and 
                client_id in self.namespaces[namespace]):
                await websocket.send_text(message)

    async def add_to_room(self, room: str, client_id: str):
        if client_id in self.active_connections:
            self.rooms[room][client_id] = self.active_connections[client_id]

    def remove_from_room(self, room: str, client_id: str):
        if room in self.rooms and client_id in self.rooms[room]:
            del self.rooms[room][client_id]

    def register_event(self, event: str, handler: Callable, namespace: str = "/", module: str = None):
        prefixed_event = self._apply_prefix(event, 'event', module)
        self.event_handlers[namespace][prefixed_event] = handler

    async def handle_event(self, event: str, data: Any, client_id: str, namespace: str = "/", module: str = None):
        prefixed_event = self._apply_prefix(event, 'event', module)
        if namespace in self.event_handlers and prefixed_event in self.event_handlers[namespace]:
            handler = self.event_handlers[namespace][prefixed_event]
            if inspect.iscoroutinefunction(handler):
                return await handler(data, client_id)
            return handler(data, client_id)
        return None

    def _apply_prefix(self, name: str, prefix_type: str, module: str = None):
        if module and module in self._prefixes[prefix_type]:
            return f"{self._prefixes[prefix_type][module]}-{name}"
        return name

    def load_prefixes(self):
        modules_metadata = self._config.get_modules_metadata()
        for module_name, config in self._config._configs.items():
            if module_name == 'global' or module_name not in modules_metadata:
                continue
                
            websocket_config = config.get('websocket', {})
            if 'event-prefix' in websocket_config:
                self._prefixes['event'][module_name] = websocket_config['event-prefix']
            if 'room-prefix' in websocket_config:
                self._prefixes['room'][module_name] = websocket_config['room-prefix']
            if 'namespace-prefix' in websocket_config:
                self._prefixes['namespace'][module_name] = websocket_config['namespace-prefix']

class WebSocketManager:
    _instance = None
    _lock = threading.Lock()

    def __new__(cls):
        if cls._instance is None:
            with cls._lock:
                if cls._instance is None:
                    cls._instance = super().__new__(cls)
                    cls._instance._init()
        return cls._instance

    def _init(self):
        self._project_root = None
        self._config = ConfigManager()
        self._connection_manager = ConnectionManager()
        self._initialized = False
        self._handlers = {}
        self._lock = threading.RLock()

    def set_project_root(self, path: Union[str, Path]):
        """Définit la racine du projet"""
        self._project_root = Path(path).resolve()

    def initialize(self):
        """Initialise le WebSocketManager"""
        with self._lock:
            if self._initialized:
                return
                
            self._connection_manager.load_prefixes()
            self._load_handlers()
            self._initialized = True

    def _load_handlers(self):
        """Charge les handlers WebSocket depuis les fichiers ws/"""
        # Global handlers
        global_ws = self._project_root / 'ws'
        if global_ws.exists():
            self._load_ws_handlers_from_dir(global_ws, 'global')

        # Module handlers
        modules_metadata = self._config.get_modules_metadata()
        modules_dir = self._project_root / 'modules'
        if modules_dir.exists():
            for module_name in modules_metadata.keys():
                module_ws = modules_dir / module_name / 'ws'
                if module_ws.exists():
                    self._load_ws_handlers_from_dir(module_ws, module_name)

    def _load_ws_handlers_from_dir(self, directory: Path, context: str):
        """Charge les handlers depuis un répertoire ws/"""
        for handler_file in directory.glob('**/*.py'):
            if handler_file.name.startswith('_'):
                continue

            try:
                module_path = '.'.join([
                    'modules', context, 'ws', 
                    *handler_file.relative_to(directory.parent).with_suffix('').parts
                ]) if context != 'global' else '.'.join([
                    'ws', 
                    *handler_file.relative_to(self._project_root).with_suffix('').parts
                ])

                module = importlib.import_module(module_path)
                
                # Enregistre tous les handlers du module
                for name, obj in inspect.getmembers(module):
                    if name.startswith('handle_') and callable(obj):
                        event_name = name[7:]  # Enlève 'handle_'
                        namespace = getattr(obj, '__namespace__', '/')
                        self._connection_manager.register_event(
                            event_name, obj, namespace, context
                        )
                        logger.info(f"Registered WS handler: {namespace}/{event_name} from {context}")

            except Exception as e:
                logger.error(f"Failed to load WS handler from {handler_file}: {e}")

    async def handle_connection(self, websocket: WebSocket, namespace: str = "/"):
        """Gère une connexion WebSocket"""
        client_id = str(uuid4())
        await self._connection_manager.connect(websocket, client_id, namespace)
        
        try:
            while True:
                data = await websocket.receive_text()
                try:
                    message = json.loads(data)
                    if isinstance(message, dict) and 'event' in message:
                        await self._handle_ws_message(message, client_id, namespace)
                except json.JSONDecodeError:
                    logger.warning(f"Invalid JSON received from {client_id}")
        except WebSocketDisconnect:
            self._connection_manager.disconnect(client_id, namespace)
            logger.info(f"Client disconnected: {client_id}")

    async def _handle_ws_message(self, message: dict, client_id: str, namespace: str = "/"):
        """Traite un message WebSocket"""
        event = message['event']
        data = message.get('data')
        module = message.get('module')
        
        result = await self._connection_manager.handle_event(
            event, data, client_id, namespace, module
        )
        
        if result is not None:
            await self.send_to_client(client_id, {
                'event': f"{event}_response",
                'data': result
            }, namespace)

    async def send_to_client(self, client_id: str, message: dict, namespace: str = "/"):
        """Envoie un message à un client spécifique"""
        await self._connection_manager.send_personal_message(
            json.dumps(message), client_id
        )

    async def broadcast(self, message: dict, namespace: str = "/"):
        """Diffuse un message à tous les clients d'un namespace"""
        await self._connection_manager.broadcast(
            json.dumps(message), namespace
        )

    async def send_to_room(self, room: str, message: dict, namespace: str = "/"):
        """Envoie un message à tous les clients d'une room"""
        await self._connection_manager.send_to_room(
            room, json.dumps(message), namespace
        )

    async def add_to_room(self, room: str, client_id: str):
        """Ajoute un client à une room"""
        self._connection_manager.add_to_room(room, client_id)

    def remove_from_room(self, room: str, client_id: str):
        """Retire un client d'une room"""
        self._connection_manager.remove_from_room(room, client_id)

    def get_client_ids(self, namespace: str = "/") -> List[str]:
        """Retourne les IDs des clients connectés à un namespace"""
        return list(self._connection_manager.namespaces[namespace].keys())

    def get_room_clients(self, room: str) -> List[str]:
        """Retourne les IDs des clients dans une room"""
        return list(self._connection_manager.rooms[room].keys())

def websocket_handler(namespace: str = "/"):
    """Décorateur pour enregistrer un handler WebSocket"""
    def decorator(func):
        func.__namespace__ = namespace
        return func
    return decorator