# 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.

"""WSX handler — bridges WebSocket connections to the request/route/response cycle.

Accepts a WebSocket connection, authenticates via handshake headers,
then enters a receive loop where each WSX message is dispatched as
an independent request through the server's router.

Messages with path starting with ``_wsx/`` are control messages
(ping/pong, auth) handled inline and never routed.
"""

from __future__ import annotations

import asyncio
import logging
import uuid
from typing import TYPE_CHECKING, Any

from genro_routes import is_result_wrapper
from genro_toolbox.smartasync import smartasync

from .protocol import build_wsx_response, is_wsx_message, parse_wsx_message
from .registry import WsxRegistry
from ..exceptions import HTTPException
from ..request import set_current_request
from ..websocket import WebSocket

if TYPE_CHECKING:
    from ..server.server import AsgiServer
    from ..types import Receive, Scope, Send

__all__ = ["WsxHandler"]

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

# Same error mapping used by Dispatcher
ROUTER_ERRORS: dict[str, type[Exception]] = {
    "not_found": HTTPException,
    "not_authorized": HTTPException,
    "not_authenticated": HTTPException,
    "not_available": HTTPException,
    "validation_error": HTTPException,
}


class WsxHandler:
    """Handles WebSocket connections speaking the WSX protocol.

    Each connection spawns a receive loop. Each WSX message becomes
    an independent request dispatched through the server router.
    Messages are processed concurrently via asyncio tasks.
    """

    __slots__ = ("server", "registry", "max_concurrent")

    def __init__(self, server: AsgiServer, max_concurrent: int = 10) -> None:
        """Args:
            server: Parent AsgiServer instance.
            max_concurrent: Max concurrent WSX messages per connection.
        """
        self.server = server
        self.registry = WsxRegistry()
        self.max_concurrent = max_concurrent

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        """ASGI entry point — accept WebSocket, authenticate, run message loop.

        Args:
            scope: ASGI WebSocket scope dict.
            receive: ASGI receive callable.
            send: ASGI send callable.
        """
        ws = WebSocket(scope, receive, send)
        await ws.accept()
        connection_id = str(uuid.uuid4())

        # Parse headers from scope (same logic as middleware.headers_dict)
        scope["_headers"] = {
            name.decode("latin-1").lower(): value.decode("latin-1")
            for name, value in scope.get("headers", [])
        }

        # Authenticate using the auth middleware if available
        auth = self._authenticate(scope)
        scope["auth"] = auth
        scope["_filters"] = {
            "auth_tags": auth["tags"] if auth else [],
        }

        self.registry.register(connection_id, ws, scope, auth)
        logger.info(
            "WSX connection %s from %s (identity=%s)",
            connection_id,
            scope.get("client"),
            auth.get("identity") if auth else None,
        )

        try:
            await self._message_loop(ws, scope, connection_id)
        finally:
            self.registry.unregister(connection_id)
            logger.info("WSX connection %s closed", connection_id)

    def _authenticate(self, scope: Scope) -> dict[str, Any] | None:
        """Authenticate WebSocket using handshake headers.

        Delegates to AuthMiddleware._authenticate() if the middleware
        is present in the server's middleware chain.
        """
        # Walk the middleware chain to find AuthMiddleware
        app: Any = self.server.dispatcher
        while app is not None:
            if hasattr(app, "_authenticate") and hasattr(app, "_auth_config"):
                result: dict[str, Any] | None = app._authenticate(scope)
                return result
            app = getattr(app, "app", None)
        return None

    async def _message_loop(
        self, ws: WebSocket, scope: Scope, connection_id: str
    ) -> None:
        """Receive loop — dispatches each WSX message as a concurrent task."""
        pending: set[asyncio.Task[None]] = set()
        semaphore = asyncio.Semaphore(self.max_concurrent)

        try:
            async for raw in ws.iter_text():
                if not is_wsx_message(raw):
                    logger.debug("Non-WSX message ignored on %s", connection_id)
                    continue

                parsed = parse_wsx_message(raw)
                path = parsed.get("path", "/")

                if path.startswith("_wsx/"):
                    await self._handle_control(ws, parsed)
                else:
                    await semaphore.acquire()
                    task = asyncio.create_task(
                        self._dispatch_message(ws, scope, raw, parsed)
                    )
                    pending.add(task)
                    def _on_task_done(t: asyncio.Task[None]) -> None:
                        pending.discard(t)
                        semaphore.release()

                    task.add_done_callback(_on_task_done)
        finally:
            # Wait for in-flight tasks to complete before cleanup
            if pending:
                _, still_pending = await asyncio.wait(
                    pending, timeout=5.0
                )
                for task in still_pending:
                    task.cancel()
                if still_pending:
                    await asyncio.gather(*still_pending, return_exceptions=True)

    async def _handle_control(self, ws: WebSocket, parsed: dict[str, Any]) -> None:
        """Handle WSX control messages (_wsx/* paths)."""
        path = parsed.get("path", "")
        msg_id = parsed.get("id", "")

        if path == "_wsx/ping":
            response = build_wsx_response(id=msg_id, status=200, data="pong")
            await ws.send_text(response)
        else:
            response = build_wsx_response(
                id=msg_id,
                status=404,
                data={"error": f"Unknown control path: {path}"},
            )
            await ws.send_text(response)

    async def _dispatch_message(
        self,
        ws: WebSocket,
        scope: Scope,
        raw: str,
        parsed: dict[str, Any],
    ) -> None:
        """Dispatch a single WSX message through the router."""
        msg_id = parsed.get("id", str(uuid.uuid4()))

        try:
            request = await self.server.request_registry.create(
                scope, ws._receive, ws._send,
                message=raw,
                websocket=ws,
                server=self.server,
            )
            set_current_request(request)

            try:
                node = self.server.router.node(
                    request.path,
                    auth_tags=request.auth_tags,
                    env_capabilities=request.env_capabilities,
                    errors=ROUTER_ERRORS,
                )

                result = await smartasync(node)(**dict(request.query))

                if is_result_wrapper(result):
                    metadata = {**node.metadata, **result.metadata}
                    data = result.value
                else:
                    metadata = node.metadata
                    data = result

                # Build response headers from metadata if present
                response_headers: dict[str, str] | None = None
                content_type = metadata.get("content_type")
                if content_type:
                    response_headers = {"content-type": content_type}

                response = build_wsx_response(
                    id=msg_id,
                    status=200,
                    headers=response_headers,
                    data=data,
                )
                await ws.send_text(response)

            except HTTPException as exc:
                response = build_wsx_response(
                    id=msg_id,
                    status=exc.status_code,
                    data={"error": exc.detail},
                )
                await ws.send_text(response)
            except Exception as exc:
                logger.exception("Handler error for message %s", msg_id)
                response = build_wsx_response(
                    id=msg_id,
                    status=500,
                    data={"error": str(exc)},
                )
                await ws.send_text(response)
            finally:
                set_current_request(None)
                self.server.request_registry.unregister()

        except Exception as exc:
            logger.exception("Request creation error for message %s", msg_id)
            response = build_wsx_response(
                id=msg_id,
                status=500,
                data={"error": f"Request creation failed: {exc}"},
            )
            await ws.send_text(response)


if __name__ == "__main__":
    pass
