# Copyright 2025 Softwell S.r.l.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Registry for active WSX WebSocket connections.

Tracks connected clients, supports lookup and broadcast.
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Any, Iterator

from .protocol import build_wsx_response

if TYPE_CHECKING:
    from ..websocket import WebSocket
    from ..types import Scope

__all__ = ["WsxRegistry", "WsxConnectionInfo"]

logger = logging.getLogger("genro_asgi.wsx")


class WsxConnectionInfo:
    """Metadata for a registered WSX connection."""

    __slots__ = ("connection_id", "websocket", "scope", "auth")

    def __init__(
        self,
        connection_id: str,
        websocket: WebSocket,
        scope: Scope,
        auth: dict[str, Any] | None = None,
    ) -> None:
        """Args:
            connection_id: Unique identifier for this connection.
            websocket: Active WebSocket instance.
            scope: ASGI scope dict for this connection.
            auth: Authentication result dict (identity, tags), or None.
        """
        self.connection_id = connection_id
        self.websocket = websocket
        self.scope = scope
        self.auth = auth

    @property
    def identity(self) -> str | None:
        """Authenticated user identity, if available."""
        if self.auth is None:
            return None
        return self.auth.get("identity")

    @property
    def client(self) -> tuple[str, int] | None:
        """Client address as (host, port)."""
        return self.scope.get("client")


class WsxRegistry:
    """Registry of active WSX WebSocket connections."""

    __slots__ = ("_connections",)

    def __init__(self) -> None:
        self._connections: dict[str, WsxConnectionInfo] = {}

    def register(
        self,
        connection_id: str,
        websocket: WebSocket,
        scope: Scope,
        auth: dict[str, Any] | None = None,
    ) -> WsxConnectionInfo:
        """Register a new WSX connection."""
        info = WsxConnectionInfo(connection_id, websocket, scope, auth)
        self._connections[connection_id] = info
        logger.debug("WSX connection registered: %s", connection_id)
        return info

    def unregister(self, connection_id: str) -> WsxConnectionInfo | None:
        """Unregister a WSX connection. Returns info if found."""
        info = self._connections.pop(connection_id, None)
        if info is not None:
            logger.debug("WSX connection unregistered: %s", connection_id)
        return info

    def get(self, connection_id: str) -> WsxConnectionInfo | None:
        """Lookup connection by ID."""
        return self._connections.get(connection_id)

    def find_by_identity(self, identity: str) -> list[WsxConnectionInfo]:
        """Find all connections for a given authenticated identity."""
        return [c for c in self._connections.values() if c.identity == identity]

    async def broadcast(self, data: Any, *, exclude: str | None = None) -> int:
        """Send a WSX response to all connected clients.

        Args:
            data: Payload to send.
            exclude: Connection ID to skip.

        Returns:
            Number of clients the message was sent to.
        """
        msg = build_wsx_response(id="_broadcast", status=200, data=data)
        sent = 0
        for cid, info in list(self._connections.items()):
            if cid == exclude:
                continue
            try:
                await info.websocket.send_text(msg)
                sent += 1
            except Exception:
                logger.warning("Broadcast failed for connection %s", cid)
        return sent

    async def send_to(self, identity: str, data: Any) -> int:
        """Send a WSX response to all connections of a given identity.

        Returns:
            Number of clients the message was sent to.
        """
        msg = build_wsx_response(id="_directed", status=200, data=data)
        sent = 0
        for info in self.find_by_identity(identity):
            try:
                await info.websocket.send_text(msg)
                sent += 1
            except Exception:
                logger.warning(
                    "Send to %s failed for connection %s",
                    identity,
                    info.connection_id,
                )
        return sent

    def __len__(self) -> int:
        return len(self._connections)

    def __iter__(self) -> Iterator[WsxConnectionInfo]:
        return iter(self._connections.values())

    def __contains__(self, connection_id: str) -> bool:
        return connection_id in self._connections


if __name__ == "__main__":
    registry = WsxRegistry()
    print(f"Empty registry: {len(registry)} connections")
