"""Web chat channel — outbound WebSocket relay (no SSH, no password)."""

import asyncio
import json
import os
import uuid
from typing import Any

import websockets
from loguru import logger

from flowly.bus.events import InboundMessage, OutboundMessage
from flowly.bus.queue import MessageBus
from flowly.channels.base import BaseChannel
from flowly.config.schema import WebChannelConfig


class WebChannel(BaseChannel):
    """
    Web chat channel using an outbound relay WebSocket.

    The VPS connects OUTWARD to the proxy (like Telegram polls Telegram servers).
    The browser connects to the proxy with Firebase JWT — no SSH, no password.

    Protocol:
      VPS  → proxy /relay?token=<agent_jwt>  (persistent outbound connection)
      Browser → proxy /?token=<browser_jwt>  (routed through agent connection)

    Messages forwarded from proxy have a `sessionId` field identifying the browser.
    Responses sent back include the same `sessionId` so the proxy can route them.
    """

    name = "web"

    def __init__(self, config: WebChannelConfig, bus: MessageBus):
        super().__init__(config, bus)
        self.config: WebChannelConfig = config
        self._ws = None
        self._reconnect_delay = 5  # seconds
        self._max_reconnect_delay = 60
        # Track active browser sessions: sessionId → asyncio.Event (response ready)
        self._pending: dict[str, asyncio.Queue] = {}

    async def start(self) -> None:
        """Connect outbound to relay proxy and keep reconnecting (Telegram-like)."""
        if not self.config.enabled:
            return
        if not self.config.auth_token:
            logger.error("[WebChannel] auth_token not set — cannot connect to relay")
            return
        if not self.config.server_id:
            # Try env fallback
            self.config.server_id = os.environ.get("FLOWLY_SERVER_ID", "")
        if not self.config.server_id:
            logger.error("[WebChannel] server_id not set — cannot connect to relay")
            return

        self._running = True
        delay = self._reconnect_delay

        while self._running:
            try:
                await self._connect_and_run()
                delay = self._reconnect_delay  # reset on clean disconnect
            except Exception as e:
                if not self._running:
                    break
                logger.warning(f"[WebChannel] Disconnected ({e}), reconnecting in {delay}s...")
                await asyncio.sleep(delay)
                delay = min(delay * 2, self._max_reconnect_delay)

    async def stop(self) -> None:
        self._running = False
        if self._ws:
            await self._ws.close()
            self._ws = None

    async def send(self, msg: OutboundMessage) -> None:
        """Send agent response back to the browser via the relay proxy."""
        if not self._ws:
            logger.warning("[WebChannel] Cannot send — not connected to relay")
            return

        session_id = msg.chat_id  # chat_id = sessionId for web channel
        run_id = msg.metadata.get("run_id", str(uuid.uuid4()))

        # Send final chat event (browser MoltbotClient listens for this)
        event_msg = {
            "type": "event",
            "sessionId": session_id,
            "event": "chat",
            "data": {
                "state": "final",
                "runId": run_id,
                "message": {
                    "content": [{"type": "text", "text": msg.content}]
                },
            },
        }
        try:
            await self._ws.send(json.dumps(event_msg))
        except Exception as e:
            logger.error(f"[WebChannel] Failed to send response: {e}")

    async def _connect_and_run(self) -> None:
        """Open one WebSocket connection to the relay proxy and process messages."""
        from jose import jwt as jose_jwt
        import time

        jwt_secret = os.environ.get("MOLTBOT_PROXY_JWT_SECRET")
        if not jwt_secret:
            logger.error("[WebChannel] MOLTBOT_PROXY_JWT_SECRET not set — cannot connect to relay")
            return

        # Build agent JWT
        now = int(time.time())
        payload = {
            "type": "agent",
            "serverId": self.config.server_id,
            "gatewayAuthToken": self.config.auth_token,
            "iat": now,
            "exp": now + 3600 * 24,  # 24h — long-lived agent token
            "iss": "flowly",
            "aud": "moltbot-proxy",
        }
        token = jose_jwt.encode(payload, jwt_secret, algorithm="HS256")
        url = f"{self.config.relay_url}?token={token}"

        logger.info(f"[WebChannel] Connecting to relay: {self.config.relay_url}")

        async with websockets.connect(url, ping_interval=30, ping_timeout=10) as ws:
            self._ws = ws
            logger.info("[WebChannel] Connected to relay proxy")

            async for raw in ws:
                if not self._running:
                    break
                try:
                    msg = json.loads(raw)
                    await self._handle_relay_message(ws, msg)
                except json.JSONDecodeError:
                    logger.warning("[WebChannel] Invalid JSON from relay")
                except Exception as e:
                    logger.error(f"[WebChannel] Error handling relay message: {e}")

        self._ws = None

    async def _handle_relay_message(self, ws, msg: dict) -> None:
        """Handle a message forwarded by the relay proxy."""
        msg_type = msg.get("type")

        if msg_type == "ready":
            logger.info("[WebChannel] Relay confirmed agent ready")

        elif msg_type == "browser-connected":
            session_id = msg.get("sessionId", "")
            logger.info(f"[WebChannel] Browser connected: {session_id}")

        elif msg_type == "browser-disconnected":
            session_id = msg.get("sessionId", "")
            logger.info(f"[WebChannel] Browser disconnected: {session_id}")
            self._pending.pop(session_id, None)

        elif msg_type == "rpc":
            await self._handle_rpc(ws, msg)

        elif msg_type == "ping":
            session_id = msg.get("sessionId")
            pong = {"type": "pong", "timestamp": msg.get("timestamp")}
            if session_id:
                pong["sessionId"] = session_id
            await ws.send(json.dumps(pong))

        else:
            logger.debug(f"[WebChannel] Unhandled relay message type: {msg_type}")

    async def _handle_rpc(self, ws, msg: dict) -> None:
        """Handle an RPC call from the browser (forwarded by proxy)."""
        method = msg.get("method", "")
        rpc_id = msg.get("id", "")
        params = msg.get("params", {})
        session_id = msg.get("sessionId", "")

        if method == "chat.send":
            message_text = params.get("message", "")
            session_key = f"web:{session_id}"
            idempotency_key = params.get("idempotencyKey") or str(uuid.uuid4())
            run_id = idempotency_key

            # ACK immediately with runId
            ack = {
                "type": "rpc",
                "id": rpc_id,
                "sessionId": session_id,
                "result": {"runId": run_id},
            }
            await ws.send(json.dumps(ack))

            # Build streaming callback — sends delta events to this browser session.
            # Captures ws and session_id at call time (safe even if self._ws reconnects).
            async def stream_callback(delta: str) -> None:
                await ws.send(json.dumps({
                    "type": "event",
                    "sessionId": session_id,
                    "event": "chat",
                    "data": {"state": "streaming", "runId": run_id, "delta": delta},
                }))

            # Process message asynchronously (don't block the recv loop)
            asyncio.create_task(
                self._process_message(session_id, session_key, message_text, run_id, stream_callback)
            )

        elif method == "chat.abort":
            run_id = params.get("runId", "")
            ack = {"type": "rpc", "id": rpc_id, "sessionId": session_id, "result": {}}
            await ws.send(json.dumps(ack))

        elif method == "chat.history":
            # History is managed by Firestore on the client side — return empty
            ack = {
                "type": "rpc",
                "id": rpc_id,
                "sessionId": session_id,
                "result": {"messages": []},
            }
            await ws.send(json.dumps(ack))

        else:
            logger.warning(f"[WebChannel] Unknown RPC method: {method}")

    async def _process_message(
        self,
        session_id: str,
        session_key: str,
        content: str,
        run_id: str,
        stream_callback=None,
    ) -> None:
        """Push message to bus and wait for agent response."""
        inbound_with_session = _WebInboundMessage(
            channel="web",
            sender_id=session_id,
            chat_id=session_id,
            content=content,
            metadata={"session_key": session_key, "run_id": run_id, "stream_callback": stream_callback},
            _session_key=session_key,
        )

        await self.bus.publish_inbound(inbound_with_session)


# ---------------------------------------------------------------------------
# InboundMessage subclass that allows overriding session_key
# ---------------------------------------------------------------------------

from dataclasses import dataclass, field
from datetime import datetime
from flowly.bus.events import InboundMessage as _Base


@dataclass
class _WebInboundMessage(_Base):
    _session_key: str = ""

    @property
    def session_key(self) -> str:  # type: ignore[override]
        return self._session_key or f"web:{self.chat_id}"
