"""
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 }
- meta_agent_linked: { execution_id, meta_agent_session_id, timestamp } (dynamic executions)
- access_request_received: { execution_id, requester_id } (access control)
- access_granted: { execution_id, user_id } (access control)
- access_denied: { execution_id, user_id } (access control)
"""

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
from control_plane_api.app.lib.validation.user_id_validation import log_if_invalid_user_id

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')
            log_if_invalid_user_id(
                user_id,
                source="websocket_status_jwt",
                org=organization_id,
                raw_sub=decoded.get('sub') if decoded else None,
                raw_user_id=decoded.get('user_id') if decoded else None,
                fell_back_to_email=user_id == (decoded.get('email') if decoded else None),
            )

        # 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')
                log_if_invalid_user_id(
                    user_id,
                    source="websocket_status_auth",
                    org=organization_id,
                    has_uuid_key="uuid" in user_data if user_data else False,
                    fell_back_to_email=user_id == user_data.get('email') if user_data else False,
                )
                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()
