"""
WebSocket endpoint for real-time execution status updates across all executions.

Allows the kanban board to receive instant status changes without polling.
This is a lightweight endpoint that only streams status changes, not full
execution content (messages, tool calls, etc.).

Features:
- Long-running connection (no timeout)
- Authentication via JWT tokens (same as websocket_client.py)
- Subscribes to organization-wide status updates
- Heartbeat/keepalive mechanism

Architecture:
    Browser → WebSocket → Control Plane API → Redis Pub/Sub → Status Events

Events sent:
- execution_status: { execution_id, status, updated_at }
- execution_created: { execution_id, status, entity_type, entity_name, prompt_preview }
- execution_deleted: { execution_id }
"""

from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, status
from typing import Optional, Dict, Any, Set
import structlog
import json
import asyncio
from datetime import datetime, timezone

from control_plane_api.app.lib.redis_client import get_redis_client

logger = structlog.get_logger()

router = APIRouter()


class StatusConnectionManager:
    """
    Manages WebSocket connections for organization-wide status updates.

    Features:
    - Per-organization connection tracking
    - Connection statistics
    - Cleanup on disconnect
    """

    def __init__(self):
        # connection_id -> WebSocket
        self._connections: Dict[str, WebSocket] = {}

        # organization_id -> Set[connection_id]
        self._org_connections: Dict[str, Set[str]] = {}

        # Connection statistics
        self._stats = {
            "total_connections": 0,
            "active_connections": 0,
            "messages_sent": 0,
            "errors": 0,
        }

        # Connection counter for unique IDs
        self._connection_counter = 0

    def _generate_connection_id(self) -> str:
        """Generate unique connection ID."""
        self._connection_counter += 1
        return f"status_conn_{self._connection_counter}_{int(datetime.now(timezone.utc).timestamp() * 1000)}"

    async def connect(
        self,
        organization_id: str,
        websocket: WebSocket,
    ) -> str:
        """Register a new status WebSocket connection."""
        connection_id = self._generate_connection_id()

        # Track connection
        self._connections[connection_id] = websocket

        if organization_id not in self._org_connections:
            self._org_connections[organization_id] = set()
        self._org_connections[organization_id].add(connection_id)

        # Update stats
        self._stats["total_connections"] += 1
        self._stats["active_connections"] = len(self._connections)

        logger.info(
            "status_websocket_connected",
            connection_id=connection_id,
            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
            active_connections=self._stats["active_connections"],
        )

        return connection_id

    async def disconnect(self, connection_id: str, organization_id: str) -> None:
        """Unregister a status WebSocket connection."""
        if connection_id in self._connections:
            del self._connections[connection_id]

        if organization_id in self._org_connections:
            self._org_connections[organization_id].discard(connection_id)
            if not self._org_connections[organization_id]:
                del self._org_connections[organization_id]

        self._stats["active_connections"] = len(self._connections)

        logger.info(
            "status_websocket_disconnected",
            connection_id=connection_id,
            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
            active_connections=self._stats["active_connections"],
        )

    def get_stats(self) -> Dict[str, Any]:
        """Get connection statistics."""
        return {
            **self._stats,
            "connections_by_org": {
                org_id[:8] if len(org_id) > 8 else org_id: len(conns)
                for org_id, conns in self._org_connections.items()
            },
        }


# Global connection manager
status_manager = StatusConnectionManager()


async def send_json(websocket: WebSocket, event_type: str, data: Dict[str, Any]) -> None:
    """
    Send JSON message via WebSocket.

    Args:
        websocket: WebSocket connection
        event_type: Event type (e.g., 'execution_status', 'connected')
        data: Event data payload
    """
    try:
        message = {
            "type": event_type,
            **data
        }
        await websocket.send_text(json.dumps(message))
        status_manager._stats["messages_sent"] += 1
    except Exception as e:
        logger.error("failed_to_send_status_message", error=str(e), event_type=event_type)
        status_manager._stats["errors"] += 1
        raise


async def handle_auth(websocket: WebSocket, token: str) -> Optional[str]:
    """
    Handle authentication message.

    Supports multiple authentication methods:
    1. JWT tokens with org_id claim (Auth0 tokens)
    2. Kubiya API keys (validated via Kubiya API)

    Args:
        websocket: WebSocket connection
        token: Authentication token (JWT or API key)

    Returns:
        organization_id if authentication successful, None otherwise
    """
    try:
        # Import auth utilities
        from control_plane_api.app.middleware.auth import (
            decode_jwt_token,
            validate_kubiya_api_key,
        )

        organization_id = None
        user_id = None

        # First, try to decode as JWT
        decoded = decode_jwt_token(token)

        if decoded:
            # Check for org_id in various claim locations
            organization_id = (
                decoded.get('https://kubiya.ai/org_id') or
                decoded.get('org_id') or
                decoded.get('organization_id') or
                decoded.get('organization')  # Kubiya API keys use 'organization' field
            )
            user_id = decoded.get('sub') or decoded.get('user_id') or decoded.get('email')

        # If JWT decode failed or org_id not found, try Kubiya API validation
        if not organization_id:
            logger.debug("status_ws_jwt_org_id_not_found_trying_api_validation")

            # Try UserKey format first (for CLI API keys)
            user_data = await validate_kubiya_api_key(f"UserKey {token}", use_userkey=True)

            if not user_data:
                # Try Bearer format
                user_data = await validate_kubiya_api_key(f"Bearer {token}", use_userkey=False)

            if user_data:
                organization_id = user_data.get('org')
                user_id = user_data.get('uuid') or user_data.get('email')
                logger.info("status_websocket_auth_via_kubiya_api", org=organization_id)

        if not organization_id:
            logger.error("status_ws_org_id_missing", decoded_claims=list(decoded.keys()) if decoded else [])
            await send_json(websocket, "auth_error", {
                "error": "Organization ID not found in token",
                "code": "ORG_ID_MISSING",
            })
            return None

        logger.info(
            "status_websocket_authenticated",
            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
            user_id=user_id[:8] if user_id and len(user_id) > 8 else user_id,
        )

        # Send auth success
        await send_json(websocket, "auth_success", {
            "organization_id": organization_id,
            "user_id": user_id,
            "authenticated_at": datetime.now(timezone.utc).isoformat(),
        })

        return organization_id

    except Exception as e:
        logger.error("status_authentication_failed", error=str(e), error_type=type(e).__name__)
        await send_json(websocket, "auth_error", {
            "error": "Authentication failed",
            "code": "AUTH_FAILED",
        })
        return None


@router.websocket("/ws/executions/status")
async def websocket_execution_status_stream(
    websocket: WebSocket,
):
    """
    WebSocket endpoint for organization-wide execution status updates.

    Streams lightweight status events for ALL executions in the organization,
    enabling real-time kanban board updates without polling.

    Protocol:
        1. Client connects
        2. Client sends auth message: {"type": "auth", "token": "jwt..."}
        3. Server authenticates and sends 'auth_success'
        4. Server sends 'connected' event
        5. Server subscribes to Redis pub/sub for organization status events
        6. Server streams events until client disconnects
        7. Client sends ping, server responds with pong (keepalive)

    Events sent to client:
        - auth_success: Authentication successful
        - auth_error: Authentication failed
        - connected: Connection established
        - execution_status: { execution_id, status, updated_at }
        - pong: Response to ping (keepalive)
        - error: Error occurred

    Client messages:
        - {"type": "auth", "token": "..."}: Authenticate
        - {"type": "ping"}: Keepalive ping
    """
    organization_id: Optional[str] = None
    connection_id: Optional[str] = None
    pubsub = None

    try:
        # Accept connection
        await websocket.accept()

        logger.info("status_websocket_connection_started")

        # Wait for authentication message (timeout after 5 seconds)
        try:
            auth_message = await asyncio.wait_for(websocket.receive_text(), timeout=5.0)
            auth_data = json.loads(auth_message)

            if auth_data.get("type") == "auth" and "token" in auth_data:
                organization_id = await handle_auth(websocket, auth_data["token"])
                if not organization_id:
                    await websocket.close(code=4001, reason="Authentication failed")
                    return

                # Register connection
                connection_id = await status_manager.connect(organization_id, websocket)

        except asyncio.TimeoutError:
            logger.error("status_ws_auth_timeout")
            await websocket.close(code=4002, reason="Authentication timeout")
            return
        except json.JSONDecodeError:
            logger.error("status_ws_invalid_auth_message")
            await websocket.close(code=4003, reason="Invalid authentication message")
            return

        # Send connected event
        await send_json(websocket, "connected", {
            "connection_id": connection_id,
            "organization_id": organization_id,
            "connected_at": datetime.now(timezone.utc).isoformat(),
        })

        # Get Redis client and subscribe to status channel
        redis_client = get_redis_client()
        if not redis_client:
            logger.error("status_ws_redis_not_available")
            await send_json(websocket, "error", {"error": "Redis not available"})
            await websocket.close(code=1011, reason="Redis not available")
            return

        # Subscribe to organization-wide status channel
        status_channel = f"organization:{organization_id}:execution_status"
        pubsub = redis_client.pubsub()
        await pubsub.subscribe(status_channel)

        logger.info(
            "status_websocket_subscribed",
            channel=status_channel,
            organization_id=organization_id[:8] if len(organization_id) > 8 else organization_id,
        )

        # Listen for both Redis events and client messages
        async def listen_redis():
            """Listen for Redis pub/sub status events."""
            try:
                async for message in pubsub.listen():
                    if message["type"] == "message":
                        try:
                            data = json.loads(message["data"])
                            event_type = data.get("event_type", "execution_status")

                            # Forward the event to WebSocket
                            await send_json(websocket, event_type, {
                                "execution_id": data.get("execution_id"),
                                "status": data.get("status"),
                                "updated_at": data.get("updated_at") or datetime.now(timezone.utc).isoformat(),
                                "entity_type": data.get("entity_type"),
                                "entity_name": data.get("entity_name"),
                            })

                            logger.debug(
                                "status_event_forwarded",
                                execution_id=data.get("execution_id", "")[:8],
                                status=data.get("status"),
                            )

                        except json.JSONDecodeError as e:
                            logger.warning("status_ws_invalid_redis_message", error=str(e))

            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.error("status_redis_listener_error", error=str(e))

        async def listen_client():
            """Listen for client messages (ping, etc.)."""
            try:
                while True:
                    message = await websocket.receive_text()
                    data = json.loads(message)

                    if data.get("type") == "ping":
                        await send_json(websocket, "pong", {
                            "timestamp": int(datetime.now(timezone.utc).timestamp() * 1000)
                        })

            except WebSocketDisconnect:
                pass
            except asyncio.CancelledError:
                pass
            except Exception as e:
                logger.error("status_client_listener_error", error=str(e))

        # Run both listeners concurrently
        redis_task = asyncio.create_task(listen_redis())
        client_task = asyncio.create_task(listen_client())

        # Wait for either task to complete (client disconnect or error)
        done, pending = await asyncio.wait(
            {redis_task, client_task},
            return_when=asyncio.FIRST_COMPLETED
        )

        # Cancel pending tasks
        for task in pending:
            task.cancel()
            try:
                await task
            except asyncio.CancelledError:
                pass

    except WebSocketDisconnect:
        logger.info("status_websocket_client_disconnected", connection_id=connection_id)

    except Exception as e:
        logger.error(
            "status_websocket_error",
            error=str(e),
            error_type=type(e).__name__,
            connection_id=connection_id,
        )
        try:
            await send_json(websocket, "error", {"error": str(e)})
        except:
            pass

    finally:
        # Cleanup
        if pubsub:
            try:
                await pubsub.unsubscribe()
                await pubsub.close()
            except:
                pass

        if connection_id and organization_id:
            await status_manager.disconnect(connection_id, organization_id)

        logger.info("status_websocket_closed", connection_id=connection_id)


@router.get("/ws/executions/status/stats")
async def websocket_status_stats():
    """
    Get WebSocket status connection statistics.

    Returns statistics about active connections and message throughput.
    """
    return status_manager.get_stats()
