"""
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, 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

    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()
                    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

    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."""
        while True:
            await asyncio.sleep(self._heartbeat_interval)
            try:
                await ws.send(json.dumps({"type": "heartbeat"}))
            except Exception:
                break

    async def stop(self):
        """Gracefully disconnect."""
        self._should_run = False
        if self._ws:
            await self._ws.close()

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