from __future__ import annotations

import asyncio
import json
import logging
import uuid
from datetime import datetime, timezone
from typing import Any, Awaitable, Callable

import websockets
from websockets.asyncio.client import ClientConnection

from mrp.auth import Keypair, _b64url_encode
from mrp.models import Message, SendResult

logger = logging.getLogger("mrp.websocket")


class WSClient:
    """WebSocket connection to MRP relay."""

    def __init__(self, relay: str, keypair: Keypair):
        self.relay = relay.rstrip("/")
        self.keypair = keypair
        self._ws: ClientConnection | None = None
        self._on_message: Callable[[Message], Awaitable[dict | None]] | None = None
        self._pending_sends: dict[str, asyncio.Future[dict]] = {}
        self._closed = False

    def _ws_url(self) -> str:
        url = self.relay
        if url.startswith("https://"):
            url = "wss://" + url[8:]
        elif url.startswith("http://"):
            url = "ws://" + url[7:]
        return url + "/v1/ws"

    async def connect(self) -> None:
        """Connect and authenticate with the relay."""
        self._ws = await websockets.connect(self._ws_url())
        self._closed = False

        # Send auth frame
        timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
        canonical = f"WS\n{timestamp}"
        signature = _b64url_encode(self.keypair.sign(canonical.encode()))

        auth_frame = {
            "type": "auth",
            "public_key": self.keypair.public_key_b64,
            "timestamp": timestamp,
            "signature": signature,
        }
        await self._ws.send(json.dumps(auth_frame))

        # Wait for auth result
        try:
            raw = await asyncio.wait_for(self._ws.recv(), timeout=10.0)
        except asyncio.TimeoutError:
            await self._ws.close()
            raise ConnectionError("Auth response timeout")

        result = json.loads(raw)
        if result.get("type") != "auth_result" or result.get("status") != "ok":
            error = result.get("error", "unknown error")
            await self._ws.close()
            raise ConnectionError(f"Auth failed: {error}")

    async def send_message(
        self,
        recipient_key: str,
        body: dict,
        content_type: str = "application/json",
        thread_id: str | None = None,
        in_reply_to: str | None = None,
        expires_in: int | None = None,
        metadata: dict | None = None,
    ) -> SendResult:
        """Send a message via WebSocket and wait for send_result."""
        if not self._ws:
            raise ConnectionError("Not connected")

        request_id = str(uuid.uuid4())
        frame: dict[str, Any] = {
            "type": "send",
            "request_id": request_id,
            "recipient_key": recipient_key,
            "body": body,
        }
        if content_type != "application/json":
            frame["content_type"] = content_type
        if thread_id is not None:
            frame["thread_id"] = thread_id
        if in_reply_to is not None:
            frame["in_reply_to"] = in_reply_to
        if expires_in is not None:
            frame["expires_in"] = expires_in
        if metadata is not None:
            frame["metadata"] = metadata

        future: asyncio.Future[dict] = asyncio.get_event_loop().create_future()
        self._pending_sends[request_id] = future

        await self._ws.send(json.dumps(frame))

        try:
            result = await asyncio.wait_for(future, timeout=30.0)
        except asyncio.TimeoutError:
            self._pending_sends.pop(request_id, None)
            raise TimeoutError(f"No response for send request {request_id}")

        if result.get("status") != "ok":
            raise ConnectionError(f"Send failed: {result.get('error', 'unknown')}")

        return SendResult(
            message_id=result["message_id"],
            status="sent",
            created_at=datetime.now(timezone.utc),
            expires_at=datetime.now(timezone.utc),
        )

    async def _read_loop(self) -> None:
        """Read frames and dispatch."""
        if not self._ws:
            return
        try:
            async for raw in self._ws:
                if self._closed:
                    break
                frame = json.loads(raw)
                frame_type = frame.get("type")

                if frame_type == "ping":
                    await self._ws.send(json.dumps({"type": "pong"}))

                elif frame_type == "message":
                    msg = Message.from_dict(frame)
                    # Auto-ack
                    await self._send_ack(msg.message_id)
                    # Dispatch to handler
                    if self._on_message:
                        try:
                            reply = await self._on_message(msg)
                            if reply is not None and isinstance(reply, dict):
                                await self.send_message(
                                    recipient_key=msg.sender_key,
                                    body=reply,
                                    in_reply_to=msg.message_id,
                                    thread_id=msg.thread_id,
                                )
                        except Exception:
                            logger.exception("Error in message handler")

                elif frame_type == "send_result":
                    request_id = frame.get("request_id")
                    future = self._pending_sends.pop(request_id, None)
                    if future and not future.done():
                        future.set_result(frame)

        except websockets.ConnectionClosed:
            logger.debug("WebSocket connection closed")

    async def _send_ack(self, message_id: str) -> None:
        """Send message acknowledgement."""
        if self._ws:
            await self._ws.send(json.dumps({"type": "ack", "message_id": message_id}))

    def on_message(self, handler: Callable[[Message], Awaitable[dict | None]]) -> None:
        """Register message handler."""
        self._on_message = handler

    async def run(self) -> None:
        """Connect and run the read loop."""
        await self.connect()
        await self._read_loop()

    async def close(self) -> None:
        """Graceful close."""
        self._closed = True
        if self._ws:
            await self._ws.close()
            self._ws = None


class ReconnectingWSClient(WSClient):
    """WSClient with automatic reconnection and exponential backoff."""

    def __init__(
        self,
        relay: str,
        keypair: Keypair,
        max_retries: int = 10,
        base_delay: float = 1.0,
        max_delay: float = 60.0,
    ):
        super().__init__(relay, keypair)
        self.max_retries = max_retries
        self.base_delay = base_delay
        self.max_delay = max_delay

    async def run(self) -> None:
        """Connect, read loop, reconnect on disconnect with backoff."""
        retries = 0
        delay = self.base_delay

        while not self._closed and retries < self.max_retries:
            try:
                await self.connect()
                retries = 0
                delay = self.base_delay
                await self._read_loop()
            except (ConnectionError, OSError, websockets.ConnectionClosed) as e:
                if self._closed:
                    break
                retries += 1
                logger.warning(
                    "WebSocket disconnected (attempt %d/%d): %s",
                    retries, self.max_retries, e,
                )
                if retries < self.max_retries:
                    await asyncio.sleep(delay)
                    delay = min(delay * 2, self.max_delay)

        if retries >= self.max_retries:
            logger.error("Max reconnection attempts reached")
