"""WebSocket Progress Server

Real-time progress streaming for workflow execution.
Enables live UI updates in VS Code and other clients.

Copyright 2025 Smart AI Memory, LLC
Licensed under the Apache License, Version 2.0
"""

from __future__ import annotations

import asyncio
import json
import logging
import signal
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any

from .progress import ProgressCallback, ProgressTracker, ProgressUpdate

# Try to import websockets, provide helpful error if not available
WebSocketServerProtocol: Any = None
WEBSOCKETS_AVAILABLE = False
try:
    import websockets

    # Try new name first (websockets 11+), fall back to old name
    try:
        from websockets.server import ServerProtocol

        WebSocketServerProtocol = ServerProtocol
    except ImportError:
        from websockets.server import WebSocketServerProtocol as _WS  # type: ignore[attr-defined]

        WebSocketServerProtocol = _WS
    WEBSOCKETS_AVAILABLE = True
except ImportError:
    pass

logger = logging.getLogger(__name__)


@dataclass
class ProgressServerConfig:
    """Configuration for the progress WebSocket server."""

    host: str = "localhost"
    port: int = 8766
    ping_interval: float = 20.0
    ping_timeout: float = 20.0
    max_connections: int = 100


class ProgressServer:
    """WebSocket server for broadcasting workflow progress.

    Clients connect and receive real-time progress updates for all
    running workflows. Supports multiple concurrent connections.
    """

    def __init__(self, config: ProgressServerConfig | None = None):
        """Initialize the progress server.

        Args:
            config: Server configuration; uses defaults if not provided.
        """
        if not WEBSOCKETS_AVAILABLE:
            raise ImportError(
                "websockets package is required for ProgressServer. "
                "Install with: pip install websockets",
            )

        self.config = config or ProgressServerConfig()
        self._clients: set[WebSocketServerProtocol] = set()
        self._server: Any = None
        self._running = False
        self._trackers: dict[str, ProgressTracker] = {}

    async def start(self) -> None:
        """Start the WebSocket server."""
        self._running = True

        self._server = await websockets.serve(
            self._handle_connection,
            self.config.host,
            self.config.port,
            ping_interval=self.config.ping_interval,
            ping_timeout=self.config.ping_timeout,
        )

        logger.info(f"Progress server started on ws://{self.config.host}:{self.config.port}")

    async def stop(self) -> None:
        """Stop the WebSocket server."""
        self._running = False

        # Close all client connections
        if self._clients:
            await asyncio.gather(
                *[client.close(1001, "Server shutting down") for client in self._clients],
                return_exceptions=True,
            )
            self._clients.clear()

        # Close server
        if self._server:
            self._server.close()
            await self._server.wait_closed()

        logger.info("Progress server stopped")

    async def _handle_connection(self, websocket: WebSocketServerProtocol) -> None:
        """Handle a new client connection."""
        if len(self._clients) >= self.config.max_connections:
            await websocket.close(1013, "Maximum connections reached")
            return

        self._clients.add(websocket)
        client_id = id(websocket)
        logger.debug(f"Client {client_id} connected. Total clients: {len(self._clients)}")

        try:
            # Send welcome message with current state
            await websocket.send(
                json.dumps(
                    {
                        "type": "connected",
                        "message": "Connected to Empathy progress server",
                        "active_workflows": list(self._trackers.keys()),
                    },
                ),
            )

            # Handle incoming messages (subscriptions, etc.)
            async for message in websocket:
                await self._handle_message(websocket, message)

        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            self._clients.discard(websocket)
            logger.debug(f"Client {client_id} disconnected. Total clients: {len(self._clients)}")

    async def _handle_message(self, websocket: WebSocketServerProtocol, message: str) -> None:
        """Handle incoming message from client."""
        try:
            data = json.loads(message)
            msg_type = data.get("type")

            if msg_type == "ping":
                await websocket.send(json.dumps({"type": "pong"}))

            elif msg_type == "subscribe":
                # Client wants updates for specific workflow
                workflow_id = data.get("workflow_id")
                if workflow_id:
                    # Could track per-client subscriptions here
                    await websocket.send(
                        json.dumps({"type": "subscribed", "workflow_id": workflow_id}),
                    )

            elif msg_type == "get_status":
                # Client wants current status of all workflows
                await websocket.send(
                    json.dumps(
                        {
                            "type": "status",
                            "active_workflows": list(self._trackers.keys()),
                            "client_count": len(self._clients),
                        },
                    ),
                )

        except json.JSONDecodeError:
            await websocket.send(json.dumps({"type": "error", "message": "Invalid JSON"}))

    async def broadcast(self, update: ProgressUpdate) -> None:
        """Broadcast a progress update to all connected clients."""
        if not self._clients:
            return

        message = json.dumps({"type": "progress", **update.to_dict()})

        # Broadcast to all clients
        await asyncio.gather(
            *[self._send_safe(client, message) for client in self._clients],
            return_exceptions=True,
        )

    async def _send_safe(self, client: WebSocketServerProtocol, message: str) -> None:
        """Send message to client with error handling."""
        try:
            await client.send(message)
        except websockets.exceptions.ConnectionClosed:
            self._clients.discard(client)

    def create_tracker(
        self,
        workflow_name: str,
        workflow_id: str,
        stage_names: list[str],
    ) -> ProgressTracker:
        """Create a progress tracker that broadcasts to this server.

        Args:
            workflow_name: Name of the workflow
            workflow_id: Unique ID for this run
            stage_names: List of stage names

        Returns:
            ProgressTracker configured to broadcast updates

        """
        tracker = ProgressTracker(
            workflow_name=workflow_name,
            workflow_id=workflow_id,
            stage_names=stage_names,
        )

        # Store tracker
        self._trackers[workflow_id] = tracker

        # Add async callback to broadcast updates
        async def broadcast_callback(update: ProgressUpdate) -> None:
            """Broadcast a progress update to all connected clients."""
            await self.broadcast(update)

        tracker.add_async_callback(broadcast_callback)

        return tracker

    def remove_tracker(self, workflow_id: str) -> None:
        """Remove a tracker when workflow completes."""
        self._trackers.pop(workflow_id, None)

    def get_callback(self) -> ProgressCallback:
        """Get a synchronous callback that queues broadcasts.

        Useful for integration with sync code that can't await.
        """

        def callback(update: ProgressUpdate) -> None:
            """Queue a broadcast in the event loop for sync callers."""
            # Schedule broadcast in event loop
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    asyncio.create_task(self.broadcast(update))
                else:
                    loop.run_until_complete(self.broadcast(update))
            except RuntimeError:
                # No event loop, skip
                pass

        return callback


# Global server instance for singleton pattern
_server_instance: ProgressServer | None = None


def get_progress_server(config: ProgressServerConfig | None = None) -> ProgressServer:
    """Get or create the global progress server instance."""
    global _server_instance

    if _server_instance is None:
        _server_instance = ProgressServer(config)

    return _server_instance


@asynccontextmanager
async def progress_server_context(config: ProgressServerConfig | None = None):
    """Context manager for running the progress server.

    Usage:
        async with progress_server_context() as server:
            tracker = server.create_tracker(...)
            # Run workflow with tracker
    """
    server = get_progress_server(config)

    try:
        await server.start()
        yield server
    finally:
        await server.stop()


async def run_server(host: str = "localhost", port: int = 8766) -> None:
    """Run the progress server standalone.

    Can be run as: python -m attune.workflows.progress_server
    """
    config = ProgressServerConfig(host=host, port=port)
    server = ProgressServer(config)

    # Handle shutdown signals
    loop = asyncio.get_event_loop()
    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, lambda: asyncio.create_task(server.stop()))

    try:
        await server.start()
        print(f"Progress server running on ws://{host}:{port}")
        print("Press Ctrl+C to stop")

        # Keep running until stopped
        while server._running:
            await asyncio.sleep(1)

    except KeyboardInterrupt:
        pass
    finally:
        await server.stop()


# CLI entry point
if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Empathy Progress WebSocket Server")
    parser.add_argument("--host", default="localhost", help="Host to bind to")
    parser.add_argument("--port", type=int, default=8766, help="Port to listen on")

    args = parser.parse_args()

    asyncio.run(run_server(host=args.host, port=args.port))
