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

"""WebSocket/WSX connection client.

Connects to a genro-asgi server via WebSocket, sends WSX messages,
and receives responses matched by correlation ID.
"""

from __future__ import annotations

import asyncio
import base64
import logging
import uuid
from typing import Any

import websockets
from websockets.asyncio.client import ClientConnection

from genro_asgi.wsx.protocol import (
    build_wsx_message,
    is_wsx_message,
    parse_wsx_message,
)

__all__ = ["WsxConnection"]

logger = logging.getLogger("genro_cli")


class WsxConnection:
    """WebSocket client that speaks the WSX protocol.

    Manages connection lifecycle, authentication, and request/response
    correlation via message IDs.
    """

    __slots__ = (
        "url",
        "_ws",
        "_pending",
        "_receive_task",
        "_connected",
        "_auth_header",
    )

    def __init__(
        self,
        url: str,
        username: str | None = None,
        password: str | None = None,
        token: str | None = None,
    ) -> None:
        self.url = url
        self._ws: ClientConnection | None = None
        self._pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
        self._receive_task: asyncio.Task[None] | None = None
        self._connected = False

        # Build auth header for handshake
        if token:
            self._auth_header = f"Bearer {token}"
        elif username and password:
            b64 = base64.b64encode(f"{username}:{password}".encode()).decode()
            self._auth_header = f"Basic {b64}"
        else:
            self._auth_header = ""

    @property
    def connected(self) -> bool:
        return self._connected

    async def connect(self) -> None:
        """Open WebSocket connection and start receive loop."""
        headers: dict[str, str] = {}
        if self._auth_header:
            headers["Authorization"] = self._auth_header

        self._ws = await websockets.connect(
            self.url,
            additional_headers=headers,
        )
        self._connected = True
        self._receive_task = asyncio.create_task(self._receive_loop())
        logger.info("Connected to %s", self.url)

    async def close(self) -> None:
        """Close the connection and cancel receive loop."""
        self._connected = False
        if self._receive_task and not self._receive_task.done():
            self._receive_task.cancel()
            try:
                await self._receive_task
            except asyncio.CancelledError:
                pass
        if self._ws:
            await self._ws.close()
            self._ws = None

        # Resolve any pending futures with error
        for msg_id, future in self._pending.items():
            if not future.done():
                future.set_exception(ConnectionError("Connection closed"))
        self._pending.clear()
        logger.info("Disconnected from %s", self.url)

    async def request(
        self,
        method: str,
        path: str,
        query: dict[str, Any] | None = None,
        data: Any = None,
        headers: dict[str, str] | None = None,
        timeout: float = 30.0,
    ) -> dict[str, Any]:
        """Send a WSX request and wait for the response.

        Args:
            method: HTTP method (GET, POST, etc.)
            path: Request path.
            query: Query parameters.
            data: Request payload.
            headers: Additional headers.
            timeout: Response timeout in seconds.

        Returns:
            Parsed WSX response dict with id, status, data, headers.

        Raises:
            ConnectionError: If not connected.
            asyncio.TimeoutError: If response not received within timeout.
        """
        if not self._ws or not self._connected:
            raise ConnectionError("Not connected")

        msg_id = str(uuid.uuid4())
        msg = build_wsx_message(
            id=msg_id,
            method=method,
            path=path,
            query=query,
            data=data,
            headers=headers,
        )

        future: asyncio.Future[dict[str, Any]] = asyncio.get_event_loop().create_future()
        self._pending[msg_id] = future

        try:
            await self._ws.send(msg)
            return await asyncio.wait_for(future, timeout=timeout)
        except asyncio.TimeoutError:
            self._pending.pop(msg_id, None)
            raise
        except Exception:
            self._pending.pop(msg_id, None)
            raise

    async def ping(self) -> dict[str, Any]:
        """Send a WSX ping and wait for pong."""
        return await self.request("GET", "_wsx/ping")

    async def _receive_loop(self) -> None:
        """Background task that reads responses and resolves pending futures."""
        try:
            assert self._ws is not None
            async for raw in self._ws:
                if not isinstance(raw, str):
                    continue
                if not is_wsx_message(raw):
                    continue

                parsed = parse_wsx_message(raw)
                msg_id = parsed.get("id", "")

                future = self._pending.pop(msg_id, None)
                if future and not future.done():
                    future.set_result(parsed)
                else:
                    # Unsolicited message (broadcast, push)
                    logger.debug("Unsolicited message: %s", msg_id)
        except websockets.ConnectionClosed:
            logger.info("Server closed the connection")
        except asyncio.CancelledError:
            return
        finally:
            self._connected = False


if __name__ == "__main__":
    pass
