from __future__ import annotations

import asyncio
import json
import logging
from typing import Awaitable, Callable

import websockets
from websockets.asyncio.client import ClientConnection

from termchat.config import WS_URL

logger = logging.getLogger(__name__)

MAX_RECONNECT_DELAY = 30
PING_INTERVAL = 25  # seconds
PING_TIMEOUT = 10  # seconds


class WsClient:
    def __init__(
        self,
        token: str,
        on_message: Callable[[dict], Awaitable[None]],
        on_reconnect: Callable[[], Awaitable[None]] | None = None,
        on_disconnect: Callable[[], Awaitable[None]] | None = None,
    ) -> None:
        self._token = token
        self._on_message = on_message
        self._on_reconnect = on_reconnect
        self._on_disconnect = on_disconnect
        self._ws: ClientConnection | None = None
        self._running = False
        self._reconnect_delay = 2
        self._connected_before = False
        self._queue: list[dict] = []

    async def connect(self) -> None:
        url = f"{WS_URL}?token={self._token}"
        self._ws = await websockets.connect(
            url,
            ping_interval=PING_INTERVAL,
            ping_timeout=PING_TIMEOUT,
        )
        self._running = True
        self._reconnect_delay = 2

        if self._connected_before and self._on_reconnect:
            await self._on_reconnect()
        self._connected_before = True

        # Flush queued messages
        if self._queue:
            pending = self._queue[:]
            self._queue.clear()
            for msg in pending:
                try:
                    await self._ws.send(json.dumps(msg))
                except Exception:
                    self._queue.append(msg)
                    break

    async def send_message(
        self,
        conversation_id: str,
        content: str,
        message_type: str = "text",
        reply_to_id: str | None = None,
    ) -> None:
        payload = {
            "type": "message",
            "conversation_id": conversation_id,
            "content": content,
            "message_type": message_type,
        }
        if reply_to_id:
            payload["reply_to_id"] = reply_to_id
        if not self._ws:
            self._queue.append(payload)
            logger.debug("WS not connected, queued message (%d in queue)", len(self._queue))
            return
        try:
            await self._ws.send(json.dumps(payload))
        except Exception:
            self._queue.append(payload)
            logger.debug("WS send failed, queued message (%d in queue)", len(self._queue))

    async def mark_read(self, conversation_id: str) -> None:
        if self._ws:
            await self._ws.send(
                json.dumps({
                    "type": "read",
                    "conversation_id": conversation_id,
                })
            )

    async def send_typing(self, conversation_id: str) -> None:
        if self._ws:
            try:
                await self._ws.send(
                    json.dumps({
                        "type": "typing",
                        "conversation_id": conversation_id,
                    })
                )
            except Exception:
                pass

    @property
    def is_connected(self) -> bool:
        return self._ws is not None and self._running

    async def listen(self) -> None:
        while self._running:
            try:
                if self._ws is None:
                    await self.connect()
                async for raw in self._ws:
                    if not self._running:
                        break
                    try:
                        data = json.loads(raw)
                        await self._on_message(data)
                    except json.JSONDecodeError:
                        continue
            except websockets.ConnectionClosed:
                if not self._running:
                    break
                logger.debug("WS closed, reconnecting in %ds", self._reconnect_delay)
                self._ws = None
                if self._on_disconnect:
                    try:
                        await self._on_disconnect()
                    except Exception:
                        pass
                await asyncio.sleep(self._reconnect_delay)
                self._reconnect_delay = min(self._reconnect_delay * 2, MAX_RECONNECT_DELAY)
            except Exception:
                if not self._running:
                    break
                logger.debug("WS error, reconnecting in %ds", self._reconnect_delay)
                self._ws = None
                if self._on_disconnect:
                    try:
                        await self._on_disconnect()
                    except Exception:
                        pass
                await asyncio.sleep(self._reconnect_delay)
                self._reconnect_delay = min(self._reconnect_delay * 2, MAX_RECONNECT_DELAY)

    async def disconnect(self) -> None:
        self._running = False
        if self._ws:
            try:
                await self._ws.close()
            except Exception:
                pass
            self._ws = None
