"""
WebSocket Connection Manager

Handles:
- WSS connection to the AWS backend
- Authentication handshake
- Heartbeat keep-alive
- Automatic reconnection with exponential backoff
- Message routing (tool requests, control messages)
"""

import asyncio
import json
import logging
import time
from typing import Callable, Dict, Optional

import websockets
from websockets.exceptions import ConnectionClosed

from .protocol import (
    ToolRequest, ToolResponse, ControlMessage, MessageType,
    make_auth_message, make_heartbeat_ack,
    parse_message,
)
from .credential_vault import CredentialVault, StoredCredential

logger = logging.getLogger(__name__)


class WorkspaceConnection:
    """
    WSS client connecting the workspace agent to the AWS backend.
    
    Lifecycle:
    1. Connect to WSS endpoint
    2. Send auth message with token + version
    3. Receive AUTH_OK → ready
    4. Loop: receive messages, dispatch to handler
    5. Heartbeat keep-alive
    6. Auto-reconnect on disconnect
    """

    def __init__(
        self,
        server_url: str,
        token: str,
        version: str,
        on_tool_request: Callable,  # (ToolRequest) → ToolResponse
        vault: CredentialVault,
        heartbeat_interval: int = 30,
        max_reconnect_delay: int = 300,
    ):
        self._server_url = server_url
        self._token = token
        self._version = version
        self._on_tool_request = on_tool_request
        self._vault = vault
        self._heartbeat_interval = heartbeat_interval
        self._max_reconnect_delay = max_reconnect_delay
        self._ws = None
        self._connected = False
        self._should_run = True
        self._reconnect_delay = 1
        self._signal_bridges: Dict[str, 'SignalBridge'] = {}  # agent_id → bridge

    async def run(self):
        """Main connection loop with auto-reconnect."""
        while self._should_run:
            try:
                await self._connect_and_handle()
            except Exception as e:
                logger.error(f"Connection error: {e}")

            if not self._should_run:
                break

            # Exponential backoff
            logger.info(f"Reconnecting in {self._reconnect_delay}s...")
            await asyncio.sleep(self._reconnect_delay)
            self._reconnect_delay = min(self._reconnect_delay * 2, self._max_reconnect_delay)

    async def _connect_and_handle(self):
        """Single connection lifecycle."""
        ws_url = f"{self._server_url}/ws/workspace/{self._token}/"
        logger.info(f"Connecting to {ws_url}...")

        try:
            async with websockets.connect(
                ws_url,
                ping_interval=20,
                ping_timeout=10,
                max_size=10_485_760,  # 10MB max message
                close_timeout=5,
            ) as ws:
                self._ws = ws
                logger.info("WebSocket connected, authenticating...")

                # 1. Auth handshake
                await ws.send(make_auth_message(self._token, self._version))
                auth_response = await asyncio.wait_for(ws.recv(), timeout=10)
                auth_msg = ControlMessage.from_json(auth_response)

                if auth_msg.type == MessageType.AUTH_FAIL:
                    reason = auth_msg.data.get("reason", "Unknown")
                    logger.error(f"Authentication failed: {reason}")
                    self._should_run = False
                    return

                if auth_msg.type == MessageType.AUTH_OK:
                    logger.info(f"✅ Authenticated. User: {auth_msg.data.get('user_id', '?')}")
                    self._connected = True
                    self._reconnect_delay = 1  # Reset backoff

                    # Check for update notification
                    if auth_msg.data.get("update_available"):
                        latest = auth_msg.data.get("latest_version", "?")
                        logger.warning(f"⚠️ Update available: v{latest} (current: v{self._version})")

                # 2. Run message loop + heartbeat concurrently
                heartbeat_task = asyncio.create_task(self._heartbeat_loop(ws))
                try:
                    await self._message_loop(ws)
                finally:
                    heartbeat_task.cancel()
                    # Stop all signal bridges on disconnect
                    await self._stop_all_bridges()
                    self._connected = False

        except ConnectionClosed as e:
            logger.warning(f"Connection closed: {e.code} {e.reason}")
        except asyncio.TimeoutError:
            logger.error("Connection timed out")
        except OSError as e:
            logger.error(f"Network error: {e}")

    async def _message_loop(self, ws):
        """Receive and dispatch messages."""
        async for raw in ws:
            try:
                message = parse_message(raw)

                if isinstance(message, ToolRequest):
                    # Execute tool in thread pool to avoid blocking
                    response = await asyncio.get_event_loop().run_in_executor(
                        None, self._on_tool_request, message
                    )
                    await ws.send(response.to_jsonrpc())

                elif isinstance(message, ControlMessage):
                    await self._handle_control(ws, message)

            except json.JSONDecodeError as e:
                logger.error(f"Invalid JSON: {e}")
            except Exception as e:
                logger.exception(f"Error handling message: {e}")

    async def _handle_control(self, ws, msg: ControlMessage):
        """Handle non-RPC control messages."""
        if msg.type == MessageType.HEARTBEAT:
            await ws.send(make_heartbeat_ack())

        elif msg.type == MessageType.CREDENTIAL_SYNC:
            self._handle_credential_sync(msg.data)

        elif msg.type == MessageType.CREDENTIAL_DELETE:
            slug = msg.data.get("service_slug", "")
            self._vault.delete(slug)
            logger.info(f"Credential deleted: {slug}")
            await ws.send(json.dumps({
                "type": MessageType.CREDENTIAL_ACK.value,
                "service_slug": slug,
                "action": "deleted",
            }))

        elif msg.type == MessageType.POLICY_UPDATE:
            logger.info(f"Policy update received: {msg.data}")

        elif msg.type == MessageType.UPDATE_AVAILABLE:
            latest = msg.data.get("version", "?")
            logger.warning(f"⚠️ Update available: v{latest}")

        elif msg.type == MessageType.DISCONNECT:
            logger.info("Server requested disconnect")
            self._should_run = False

        elif msg.type == MessageType.START_SIGNAL_BRIDGE:
            await self._start_signal_bridge(msg.data)

        elif msg.type == MessageType.STOP_SIGNAL_BRIDGE:
            bridge_id = msg.data.get('bridge_id', '') if msg.data else ''
            await self._stop_signal_bridge(bridge_id)

        elif msg.type == MessageType.SYNC_SIGNAL_BRIDGES:
            await self._sync_signal_bridges(msg.data.get('bridges', []))

        elif msg.type == MessageType.SIGNAL_THROTTLE:
            bridge_id = msg.data.get('bridge_id', '')
            if bridge_id and bridge_id in self._signal_bridges:
                self._signal_bridges[bridge_id].set_throttle(msg.data.get('delay_ms', 0))

    async def _start_signal_bridge(self, config: dict):
        """Start or restart a single signal bridge."""
        from .signal_bridge import SignalBridge

        redis_url = config.get('redis_url', '')
        stream_key = config.get('stream_key', '')
        consumer_group = config.get('consumer_group', '')
        agent_id = str(config.get('agent_id', config.get('bridge_id', 'unknown')))

        if not redis_url or not stream_key:
            logger.error(f"[BRIDGE] Missing redis_url or stream_key for agent {agent_id}")
            return

        # Stop existing bridge for this agent if running
        if agent_id in self._signal_bridges:
            await self._signal_bridges[agent_id].stop()

        bridge = SignalBridge(
            redis_url=redis_url,
            stream_key=stream_key,
            consumer_group=consumer_group or f"darrxscale_{agent_id}",
            agent_id=agent_id,
            send_fn=self._ws.send,
        )
        self._signal_bridges[agent_id] = bridge
        await bridge.start()
        logger.info(f"[BRIDGE] Started bridge for agent {agent_id} → {stream_key}")

    async def _stop_signal_bridge(self, bridge_id: str = ''):
        """Stop a specific bridge, or all bridges if no ID given."""
        if bridge_id:
            bridge = self._signal_bridges.pop(bridge_id, None)
            if bridge:
                await bridge.stop()
                logger.info(f"[BRIDGE] Stopped bridge for agent {bridge_id}")
        else:
            # Stop all
            for bid, bridge in list(self._signal_bridges.items()):
                await bridge.stop()
            self._signal_bridges.clear()
            logger.info("[BRIDGE] Stopped all bridges")

    async def _sync_signal_bridges(self, bridges: list):
        """
        Diff running bridges vs desired state and start/stop as needed.
        Called on auth (initial sync) and on config updates.
        """
        desired_ids = set()
        for cfg in bridges:
            agent_id = str(cfg.get('agent_id', ''))
            if not agent_id:
                continue
            desired_ids.add(agent_id)

            # Check if bridge already running with same config
            existing = self._signal_bridges.get(agent_id)
            if existing and existing.is_running:
                # Same stream? skip restart
                if (existing._redis_url == cfg.get('redis_url') and
                        existing._stream_key == cfg.get('stream_key')):
                    continue

            # Start (or restart with new config)
            await self._start_signal_bridge(cfg)

        # Stop bridges for agents no longer in desired list
        to_remove = set(self._signal_bridges.keys()) - desired_ids
        for bid in to_remove:
            await self._stop_signal_bridge(bid)

        logger.info(
            f"[BRIDGE] Sync complete: {len(desired_ids)} desired, "
            f"{len(self._signal_bridges)} running, {len(to_remove)} stopped"
        )

    def _handle_credential_sync(self, data: dict):
        """Store a credential synced from the frontend."""
        try:
            cred = StoredCredential(
                service_slug=data["service_slug"],
                auth_type=data.get("auth_type", "bearer"),
                header_name=data.get("header_name", "Authorization"),
                header_value_pattern=data.get("header_value_pattern", "Bearer {value}"),
                secret_value=data["value"],
                query_param_name=data.get("query_param_name", ""),
                username=data.get("username", ""),
            )
            self._vault.store(cred)
            logger.info(f"Credential synced: {cred.service_slug}")
        except KeyError as e:
            logger.error(f"Invalid credential_sync: missing {e}")

    async def _heartbeat_loop(self, ws):
        """Send heartbeats to keep the connection alive. Includes bridge status every 3rd beat."""
        beat_count = 0
        while True:
            await asyncio.sleep(self._heartbeat_interval)
            try:
                beat_count += 1
                heartbeat = {"type": "heartbeat"}
                # Piggyback bridge status every 3rd heartbeat (~90s)
                if self._signal_bridges and beat_count % 3 == 0:
                    heartbeat["bridge_status"] = self._get_bridge_statuses()
                await ws.send(json.dumps(heartbeat))
            except Exception:
                break

    def _get_bridge_statuses(self) -> list:
        """Collect status from all running bridges."""
        return [bridge.get_status() for bridge in self._signal_bridges.values()]

    async def stop(self):
        """Gracefully disconnect — stop all bridges, then close WebSocket."""
        self._should_run = False
        # Stop all signal bridges first
        await self._stop_all_bridges()
        if self._ws:
            await self._ws.close()

    async def _stop_all_bridges(self):
        """Stop all running signal bridges and clean up resources."""
        if not self._signal_bridges:
            return
        logger.info(f"[BRIDGE] Stopping {len(self._signal_bridges)} bridge(s)...")
        for bid, bridge in list(self._signal_bridges.items()):
            try:
                await bridge.stop()
            except Exception as e:
                logger.warning(f"[BRIDGE] Error stopping bridge {bid}: {e}")
        self._signal_bridges.clear()
        logger.info("[BRIDGE] All bridges stopped")

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