"""
WebSocket transport layer for communicating with SimpleX Chat CLI.

The SimpleX Chat CLI runs as a WebSocket server (default port 5225).
This module handles the low-level JSON-based request/response protocol.
"""

from __future__ import annotations

import asyncio
import json
import logging
from dataclasses import dataclass, field
from typing import Any, AsyncIterator, Callable, Awaitable

import websockets
from websockets.asyncio.client import ClientConnection

logger = logging.getLogger("dexgram.transport")


@dataclass
class ChatServerRequest:
    corr_id: str
    cmd: str

    def to_json(self) -> str:
        return json.dumps({"corrId": self.corr_id, "cmd": self.cmd})


@dataclass
class ChatServerResponse:
    corr_id: str | None
    resp: dict[str, Any]


class TransportError(Exception):
    pass


class ChatTransport:
    """
    Manages the WebSocket connection to SimpleX Chat CLI.

    Sends commands as JSON with correlation IDs, routes correlated
    responses back to the caller, and pushes uncorrelated events
    into an async queue for the consumer.
    """

    def __init__(
        self,
        ws: ClientConnection,
        event_queue: asyncio.Queue[dict[str, Any]],
    ):
        self._ws = ws
        self._event_queue = event_queue
        self._pending: dict[str, asyncio.Future[dict[str, Any]]] = {}
        self._reader_task: asyncio.Task | None = None
        self._closed = False

    @classmethod
    async def connect(
        cls,
        url: str = "ws://localhost:5225",
        queue_size: int = 256,
        timeout: float = 10.0,
    ) -> ChatTransport:
        ws = await asyncio.wait_for(
            websockets.connect(url, max_size=10 * 1024 * 1024),
            timeout=timeout,
        )
        event_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=queue_size)
        transport = cls(ws, event_queue)
        transport._reader_task = asyncio.create_task(transport._read_loop())
        return transport

    async def send(self, corr_id: str, cmd: str) -> dict[str, Any]:
        if self._closed:
            raise TransportError("Transport is closed")
        future: asyncio.Future[dict[str, Any]] = asyncio.get_event_loop().create_future()
        self._pending[corr_id] = future
        req = ChatServerRequest(corr_id=corr_id, cmd=cmd)
        await self._ws.send(req.to_json())
        return await future

    async def _read_loop(self) -> None:
        try:
            async for raw in self._ws:
                if isinstance(raw, bytes):
                    raw = raw.decode("utf-8")
                try:
                    data = json.loads(raw)
                except json.JSONDecodeError:
                    logger.warning("Received non-JSON message: %s", raw[:200])
                    continue

                resp = data.get("resp")
                if resp is None or not isinstance(resp, dict):
                    logger.warning("Message missing 'resp' field: %s", raw[:200])
                    continue

                corr_id = data.get("corrId")
                if corr_id and corr_id in self._pending:
                    future = self._pending.pop(corr_id)
                    if not future.done():
                        future.set_result(resp)
                else:
                    try:
                        self._event_queue.put_nowait(resp)
                    except asyncio.QueueFull:
                        logger.warning("Event queue full, dropping event: %s", resp.get("type"))
        except websockets.ConnectionClosed:
            logger.info("WebSocket connection closed")
        except Exception as e:
            logger.error("Reader loop error: %s", e)
        finally:
            self._closed = True
            for future in self._pending.values():
                if not future.done():
                    future.set_exception(TransportError("Connection closed"))
            self._pending.clear()

    async def events(self) -> AsyncIterator[dict[str, Any]]:
        while not self._closed or not self._event_queue.empty():
            try:
                event = await asyncio.wait_for(self._event_queue.get(), timeout=1.0)
                yield event
            except asyncio.TimeoutError:
                if self._closed and self._event_queue.empty():
                    break

    async def close(self) -> None:
        self._closed = True
        if self._reader_task and not self._reader_task.done():
            self._reader_task.cancel()
            try:
                await self._reader_task
            except asyncio.CancelledError:
                pass
        await self._ws.close()
