# SPDX-FileCopyrightText: 2025 Weibo, Inc.
#
# SPDX-License-Identifier: Apache-2.0

"""
Device namespace for Socket.IO.

This module implements the /local-executor namespace for local device connections.
It handles device authentication, registration, heartbeat, and task execution.

Authentication:
- Supports both JWT Token and API Key authentication
- API Key: Token starting with 'wg-' prefix (personal keys only)
- JWT Token: Standard JWT token with user info

Events:
- connect: Device authenticates with user JWT token or API Key
- device:register: Device registers itself with device_id and name
- device:heartbeat: Device sends heartbeat every 30s
- device:status: Device reports status (idle/busy)
- task:execute: Backend pushes task to device
- response.*: OpenAI Responses API streaming events from executor
- disconnect: Cleanup on device disconnection
"""

import asyncio
import logging
import uuid
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime
from typing import Any, Dict, Generator, Optional

import socketio
from sqlalchemy.orm import Session

from app.api.ws.decorators import trace_websocket_event
from app.core.auth_utils import is_api_key, verify_api_key
from app.core.events import TaskCompletedEvent, get_event_bus
from app.db.session import SessionLocal
from app.models.subtask import Subtask, SubtaskStatus
from app.models.task import TaskResource
from app.models.user import User
from app.schemas.device import (
    DeviceHeartbeatPayload,
    DeviceOfflineEvent,
    DeviceOnlineEvent,
    DeviceRegisterPayload,
    DeviceSlotUpdateEvent,
    DeviceStatusEvent,
    DeviceStatusPayload,
)
from app.services.chat.access import get_token_expiry, verify_jwt_token
from app.services.chat.storage.db import get_db_session, run_sync_in_executor
from app.services.chat.webpage_ws_chat_emitter import get_extended_emitter
from app.services.device_service import device_service
from app.services.execution.dispatcher import ResponsesAPIEventParser
from app.services.execution.emitters.status_updating import StatusUpdatingEmitter
from app.services.execution.emitters.websocket import WebSocketResultEmitter
from shared.models import EventType
from shared.telemetry.context import set_request_context, set_user_context

logger = logging.getLogger(__name__)


@contextmanager
def _db_session() -> Generator[Session, None, None]:
    """Context manager for database session with auto-commit and auto-close."""
    db = SessionLocal()
    try:
        yield db
        db.commit()
    except Exception:
        db.rollback()
        raise
    finally:
        db.close()


@dataclass
class FailedSubtaskInfo:
    """Information about a failed subtask for WebSocket emission."""

    task_id: int
    subtask_id: int
    message_id: Optional[int]
    user_id: int


def _handle_device_disconnect(user_id: int, device_id: str) -> list[FailedSubtaskInfo]:
    """
    Handle device disconnection in database.

    With CRD model, device record stays in kinds table (just becomes offline via Redis TTL).
    Fails running subtasks and updates parent task status.
    Returns list of failed subtasks for WebSocket emission.
    """
    failed_subtasks = []
    try:
        with _db_session() as db:
            # Note: Device CRD remains in kinds table, it's just offline (Redis TTL expired)
            # No need to call mark_device_offline on MySQL

            # Find and fail running subtasks
            executor_name = f"device-{device_id}"
            running_subtasks = (
                db.query(Subtask)
                .filter(
                    Subtask.executor_name == executor_name,
                    Subtask.status == SubtaskStatus.RUNNING,
                )
                .all()
            )

            # Track unique task IDs to update parent task status
            task_ids_to_fail = set()

            for subtask in running_subtasks:
                subtask.status = SubtaskStatus.FAILED
                subtask.error_message = "Device disconnected unexpectedly"
                subtask.completed_at = datetime.now()
                task_ids_to_fail.add(subtask.task_id)
                failed_subtasks.append(
                    FailedSubtaskInfo(
                        task_id=subtask.task_id,
                        subtask_id=subtask.id,
                        message_id=subtask.message_id,
                        user_id=user_id,
                    )
                )
                logger.warning(
                    f"[Device WS] Marked subtask {subtask.id} as FAILED due to device disconnect"
                )

            # Update parent task status to FAILED
            if task_ids_to_fail:
                from app.schemas.task import TaskUpdate
                from app.services.adapters.task_kinds import task_kinds_service

                for task_id in task_ids_to_fail:
                    try:
                        task_kinds_service.update_task(
                            db=db,
                            task_id=task_id,
                            obj_in=TaskUpdate(status="FAILED"),
                            user_id=user_id,
                        )
                        logger.warning(
                            f"[Device WS] Marked task {task_id} as FAILED due to device disconnect"
                        )
                    except Exception as e:
                        logger.error(
                            f"[Device WS] Failed to update task {task_id} status: {e}"
                        )

    except Exception as e:
        logger.error(f"[Device WS] Error handling device disconnect: {e}")

    return failed_subtasks


def _register_device(
    user_id: int,
    device_id: str,
    name: str,
    client_ip: Optional[str] = None,
    device_type: Optional[str] = None,
    bind_shell: Optional[str] = None,
) -> tuple[bool, Optional[str], Optional[str]]:
    """
    Register or update device CRD in database.

    Args:
        user_id: Device owner user ID
        device_id: Device unique identifier (stored in Kind.name)
        name: Device display name
        client_ip: Device's client IP address
        device_type: Device type ('local' or 'cloud')
        bind_shell: Shell runtime binding ('claudecode' or 'openclaw')

    Returns (success, persisted_display_name, error_message).
    """
    try:
        with _db_session() as db:
            device_kind = device_service.upsert_device_crd(
                db=db,
                user_id=user_id,
                device_id=device_id,
                name=name,
                client_ip=client_ip,
                device_type=device_type,
                bind_shell=bind_shell,
            )
            persisted_display_name = (
                device_kind.json.get("spec", {}).get("displayName") or name
            )
        return True, persisted_display_name, None
    except Exception as e:
        logger.error(f"[Device WS] Error registering device: {e}")
        return False, None, str(e)


def _update_device_heartbeat(user_id: int, device_id: str) -> None:
    """
    Update device heartbeat timestamp.

    With CRD model, heartbeat is primarily tracked in Redis.
    No MySQL update needed.
    """
    # Heartbeat is managed via Redis in the async handler
    pass


def _update_device_status(user_id: int, device_id: str, status: str) -> None:
    """
    Update device status.

    With CRD model, status is primarily tracked in Redis.
    No MySQL update needed.
    """
    # Status is managed via Redis in the async handler
    pass


def _match_cloud_device_sync(
    user_id: int, client_ip: str, executor_device_id: str
) -> Optional[tuple[str, bool, Optional[dict]]]:
    """
    Synchronous helper to match cloud device by device_id.

    Returns:
        Tuple of (sandbox_id, needs_migration, device_data) if matched, None otherwise.
        - sandbox_id: The matched sandbox ID
        - needs_migration: True if legacy device needs migration
        - device_data: Dict with device info for migration (device_id, etc.)
    """
    from sqlalchemy import and_

    from app.models.kind import Kind
    from app.schemas.device import DeviceType

    with get_db_session() as db:
        cloud_devices = (
            db.query(Kind)
            .filter(
                and_(
                    Kind.user_id == user_id,
                    Kind.kind == "Device",
                    Kind.namespace == "default",
                    Kind.is_active == True,
                )
            )
            .all()
        )

        logger.info(
            f"[Device WS] Cloud device matching: "
            f"user_id={user_id}, client_ip={client_ip}, "
            f"cloud_device_count={len(cloud_devices)}, "
            f"executor_device_id={executor_device_id}"
        )

        for device in cloud_devices:
            spec = device.json.get("spec", {})
            if spec.get("deviceType") != DeviceType.CLOUD.value:
                continue

            cloud_config = spec.get("cloudConfig", {})
            sandbox_id = cloud_config.get("sandboxId", device.name)
            server_device_id = cloud_config.get("deviceId")

            # New logic: verify server-generated device_id matches
            if server_device_id:
                if server_device_id == executor_device_id:
                    logger.info(
                        f"[Device WS] Cloud device matched by device_id: "
                        f"sandbox_id={sandbox_id}, "
                        f"device_id={executor_device_id}"
                    )
                    return (sandbox_id, False, None)
                else:
                    # Device ID mismatch - skip this device
                    logger.debug(
                        f"[Device WS] Cloud device ID mismatch: "
                        f"expected={server_device_id}, "
                        f"got={executor_device_id}"
                    )
                    continue
            else:
                # Backward compatibility: old device without deviceId field
                # Use legacy matching (device.name still equals sandbox_id)
                if device.name == sandbox_id:
                    logger.info(
                        f"[Device WS] Cloud device matched (legacy mode): "
                        f"sandbox_id={sandbox_id}, "
                        f"new_device_id={executor_device_id}"
                    )
                    return (
                        sandbox_id,
                        True,
                        {"device_id": device.id},
                    )

    return None


def _update_cloud_device_id_sync(
    user_id: int,
    device_db_id: int,
    executor_device_id: str,
    sandbox_id: str,
) -> str:
    """
    Synchronous helper to update cloud device ID in CRD for backward compatibility.

    Args:
        user_id: User ID
        device_db_id: Device database ID (Kind.id)
        executor_device_id: New device ID from executor
        sandbox_id: Sandbox ID

    Returns:
        Sandbox ID
    """
    import copy

    from sqlalchemy.orm.attributes import flag_modified

    from app.models.kind import Kind

    with get_db_session() as db:
        device = db.query(Kind).filter(Kind.id == device_db_id).first()
        if not device:
            logger.error(
                f"[Device WS] Device not found for migration: id={device_db_id}"
            )
            return sandbox_id

        device_json = copy.deepcopy(device.json)
        old_device_id = device.name
        device_json["metadata"]["name"] = executor_device_id
        device_json["spec"]["deviceId"] = executor_device_id

        # Update cloudConfig with deviceId for future matching
        if "cloudConfig" in device_json["spec"]:
            device_json["spec"]["cloudConfig"]["deviceId"] = executor_device_id

        device.name = executor_device_id
        device.json = device_json
        flag_modified(device, "json")

        logger.info(
            f"[Device WS] Migrated legacy cloud device to new format: "
            f"old_id={old_device_id}, new_id={executor_device_id}, "
            f"sandbox_id={sandbox_id}"
        )

    return sandbox_id


def _verify_api_key_sync(token: str) -> Optional[tuple[int, str]]:
    """
    Synchronous helper to verify API key.

    Returns:
        Tuple of (user_id, user_name) if valid, None otherwise.
    """
    with get_db_session() as db:
        user = verify_api_key(db, token)
        if user:
            return (user.id, user.user_name)
    return None


def _get_device_slot_usage_sync(user_id: int, device_id: str) -> dict:
    """
    Synchronous helper to get device slot usage.

    Returns:
        Dict with slot usage info.
    """
    with get_db_session() as db:
        return device_service.get_device_slot_usage(db, user_id, device_id)


class DeviceNamespace(socketio.AsyncNamespace):
    """
    Socket.IO namespace for local executor connections.

    Handles:
    - Authentication on connect (using user JWT token)
    - Device registration and management
    - Heartbeat monitoring
    - Task execution routing
    """

    def __init__(self, namespace: str = "/local-executor"):
        """Initialize the device namespace."""
        super().__init__(namespace)

        # Map colon-separated event names to handler methods
        self._event_handlers: Dict[str, str] = {
            "device:register": "on_device_register",
            "device:heartbeat": "on_device_heartbeat",
            "device:status": "on_device_status",
            "device:upgrade_status": "on_device_upgrade_status",
        }

        # Shared event parser for OpenAI Responses API events
        self._event_parser = ResponsesAPIEventParser()

        # Known OpenAI Responses API event prefixes
        self._responses_api_prefixes = ("response.", "error")

        # Per-subtask locks to ensure events are processed in order
        # This prevents race conditions when multiple response.output_text.delta
        # events arrive concurrently for the same subtask
        self._subtask_locks: Dict[int, asyncio.Lock] = {}

    def _get_subtask_lock(self, subtask_id: int) -> asyncio.Lock:
        """Get or create a lock for the given subtask.

        Args:
            subtask_id: Subtask ID

        Returns:
            asyncio.Lock for the subtask
        """
        if subtask_id not in self._subtask_locks:
            self._subtask_locks[subtask_id] = asyncio.Lock()
        return self._subtask_locks[subtask_id]

    def _cleanup_subtask_lock(self, subtask_id: int) -> None:
        """Clean up the lock for a completed subtask.

        Args:
            subtask_id: Subtask ID
        """
        self._subtask_locks.pop(subtask_id, None)

    @trace_websocket_event(
        exclude_events={"connect"},
        extract_event_data=True,
    )
    async def trigger_event(self, event: str, sid: str, *args):
        """
        Override trigger_event to handle colon-separated event names.

        Args:
            event: Event name (e.g., 'device:register')
            sid: Socket ID
            *args: Event arguments

        Returns:
            Result from the event handler
        """
        return await self._execute_handler(event, sid, *args)

    async def _execute_handler(self, event: str, sid: str, *args):
        """Execute the event handler for the given event."""
        if event in self._event_handlers:
            handler_name = self._event_handlers[event]
            handler = getattr(self, handler_name, None)
            if handler:
                logger.debug(
                    f"[Device WS] Routing event '{event}' to handler '{handler_name}'"
                )
                return await handler(sid, *args)

        # Handle OpenAI Responses API events (e.g., response.output_text.delta)
        if event.startswith(self._responses_api_prefixes):
            return await self._handle_responses_api_event(sid, event, *args)

        return await super().trigger_event(event, sid, *args)

    def _get_client_ip(self, environ: dict) -> Optional[str]:
        """Extract client IP from WSGI environ.

        Checks X-Forwarded-For header first (for proxied connections),
        then falls back to REMOTE_ADDR.

        Args:
            environ: WSGI environ dict

        Returns:
            Client IP address or None
        """
        # Check X-Forwarded-For header (common for proxied connections)
        forwarded_for = environ.get("HTTP_X_FORWARDED_FOR")
        if forwarded_for:
            # X-Forwarded-For can contain multiple IPs, take the first one
            return forwarded_for.split(",")[0].strip()

        # Check X-Real-IP header
        real_ip = environ.get("HTTP_X_REAL_IP")
        if real_ip:
            return real_ip

        # Fall back to REMOTE_ADDR
        return environ.get("REMOTE_ADDR")

    async def _match_cloud_device(
        self, user_id: int, client_ip: str, executor_device_id: str
    ) -> Optional[str]:
        """Match cloud device by verifying server-generated device_id.

        When a cloud device executor connects, it should use the server-generated
        device_id (passed via DEVICE_ID environment variable). This method verifies
        that the executor's device_id matches the one stored in cloudConfig.deviceId.

        For backward compatibility: if cloudConfig.deviceId is not set (old device),
        falls back to the legacy matching logic.

        Args:
            user_id: User ID
            client_ip: WebSocket client IP address (kept for logging)
            executor_device_id: Device ID from executor (should match server-generated)

        Returns:
            Cloud device ID (sandbox_id) if matched, None otherwise
        """
        try:
            # Run database query in executor to avoid blocking event loop
            result = await run_sync_in_executor(
                _match_cloud_device_sync, user_id, client_ip, executor_device_id
            )

            if result is None:
                return None

            sandbox_id, needs_migration, device_data = result

            # If legacy device needs migration, do it in executor
            if needs_migration and device_data:
                await run_sync_in_executor(
                    _update_cloud_device_id_sync,
                    user_id,
                    device_data["device_id"],
                    executor_device_id,
                    sandbox_id,
                )

            return sandbox_id

        except Exception as e:
            logger.error(f"[Device WS] Error matching cloud device: {e}")

        return None

    async def on_connect(self, sid: str, environ: dict, auth: Optional[dict] = None):
        """
        Handle device connection.

        Verifies JWT token or API Key and prepares for device registration.

        Authentication:
        - If token starts with 'wg-', treated as API Key (personal keys only)
        - Otherwise, treated as JWT Token

        Args:
            sid: Socket ID
            environ: WSGI environ dict
            auth: Authentication data (expected: {"token": "..."})
                  Token can be either JWT Token or API Key (starting with 'wg-')

        Raises:
            ConnectionRefusedError: If authentication fails
        """
        from app.core.shutdown import shutdown_manager

        request_id = str(uuid.uuid4())[:8]
        set_request_context(request_id)

        logger.info(f"[Device WS] Connection attempt sid={sid}")

        # Reject new connections during graceful shutdown
        if shutdown_manager.is_shutting_down:
            logger.warning(
                f"[Device WS] Rejecting connection during shutdown sid={sid}"
            )
            raise ConnectionRefusedError("Server is shutting down")

        # Check auth token
        if not auth or not isinstance(auth, dict):
            logger.warning(f"[Device WS] Missing auth data sid={sid}")
            raise ConnectionRefusedError("Missing authentication token")

        token = auth.get("token")
        if not token:
            logger.warning(f"[Device WS] Missing token in auth sid={sid}")
            raise ConnectionRefusedError("Missing authentication token")

        # Determine auth type and verify token
        user = None
        auth_type = ""
        token_exp = None

        if is_api_key(token):
            # API Key authentication - run in executor to avoid blocking event loop
            auth_type = "api_key"
            user_info = await run_sync_in_executor(_verify_api_key_sync, token)
            if not user_info:
                key_preview = token[:10] + "..." if len(token) > 10 else token
                logger.warning(
                    f"[Device WS] Invalid API key sid={sid}, key={key_preview}"
                )
                raise ConnectionRefusedError("Invalid or expired API key")
            user_id, user_name = user_info
            # API Key has no expiry (token_exp stays None)
            token_exp = None
        else:
            # JWT Token authentication
            auth_type = "jwt"
            user = verify_jwt_token(token)
            if not user:
                logger.warning(f"[Device WS] Invalid JWT token sid={sid}")
                raise ConnectionRefusedError("Invalid or expired token")
            user_id = user.id
            user_name = user.user_name
            # Extract token expiry for JWT
            token_exp = get_token_expiry(token)

        # Get client IP from environ
        client_ip = self._get_client_ip(environ)

        # Save user info to session (device_id will be added on register)
        await self.save_session(
            sid,
            {
                "user_id": user_id,
                "user_name": user_name,
                "request_id": request_id,
                "token_exp": token_exp,
                "auth_token": token,
                "auth_type": auth_type,
                "device_id": None,  # Set on device:register
                "registered": False,
                "client_ip": client_ip,  # For cloud device matching
            },
        )

        set_user_context(user_id=str(user_id), user_name=user_name)

        # Join user room for device-related notifications
        user_room = f"user:{user_id}"
        await self.enter_room(sid, user_room)

        logger.info(
            f"[Device WS] Connected user={user_id} ({user_name}) via {auth_type} "
            f"sid={sid}, awaiting registration"
        )

    async def on_disconnect(self, sid: str):
        """
        Handle device disconnection.

        Cleans up Redis online status, updates MySQL, and marks running tasks as failed.

        Args:
            sid: Socket ID
        """
        try:
            session = await self.get_session(sid)
            user_id = session.get("user_id")
            device_id = session.get("device_id")
            request_id = session.get("request_id")

            if request_id:
                set_request_context(request_id)
            if user_id:
                set_user_context(user_id=str(user_id))

            logger.info(
                f"[Device WS] Disconnected user={user_id}, device={device_id}, sid={sid}"
            )

            if user_id and device_id:
                # Remove from Redis online status
                await device_service.set_device_offline(user_id, device_id)

                # Database operation: run in executor to avoid blocking event loop
                # Returns list of failed subtasks for WebSocket emission
                failed_subtasks = await run_sync_in_executor(
                    _handle_device_disconnect, user_id, device_id
                )

                # WebSocket emissions happen AFTER database connection is released
                extended_emitter = get_extended_emitter()
                # Track unique task IDs to emit task:status only once per task
                emitted_task_ids = set()
                for info in failed_subtasks:
                    await extended_emitter.emit_chat_error(
                        task_id=info.task_id,
                        subtask_id=info.subtask_id,
                        error="Device disconnected",
                        message_id=info.message_id,
                    )

                # Broadcast device offline event
                await self._broadcast_device_offline(user_id, device_id)

        except Exception as e:
            logger.error(f"[Device WS] Error in disconnect handler: {e}")

    # ============================================================
    # Device Registration and Heartbeat Events
    # ============================================================

    async def on_device_register(self, sid: str, data: dict) -> dict:
        """
        Handle device:register event.

        Registers the device in MySQL and sets online status in Redis.

        Args:
            sid: Socket ID
            data: {"device_id": str, "name": str}

        Returns:
            {"success": True, "device_id": str} or {"error": str}
        """
        try:
            payload = DeviceRegisterPayload(**data)
        except Exception as e:
            logger.warning(f"[Device WS] Invalid register payload: {e}")
            return {"error": f"Invalid payload: {e}"}

        session = await self.get_session(sid)
        user_id = session.get("user_id")

        if not user_id:
            return {"error": "Not authenticated"}

        logger.info(
            f"[Device WS] device:register user={user_id}, device_id={payload.device_id}, "
            f"name={payload.name}, executor_version={payload.executor_version}, client_ip={payload.client_ip}"
        )

        # Check if this is a cloud device registration (by IP matching)
        # Prefer self-reported IP from executor, fall back to WebSocket client IP
        client_ip = payload.client_ip or session.get("client_ip")
        is_cloud_device = False
        if client_ip:
            cloud_device_id = await self._match_cloud_device(
                user_id, client_ip, payload.device_id
            )
            if cloud_device_id:
                is_cloud_device = True
                logger.info(
                    f"[Device WS] Matched cloud device: executor_device_id={payload.device_id}, "
                    f"cloud_device_id={cloud_device_id}"
                )

        # Database operation: skip if cloud device already updated in IP matching
        # Pass client_ip to _register_device for tracking
        # Run in executor to avoid blocking event loop
        if not is_cloud_device:
            success, persisted_display_name, error = await run_sync_in_executor(
                _register_device,
                user_id,
                payload.device_id,
                payload.name,
                payload.client_ip,
                payload.device_type.value,
                payload.bind_shell.value,
            )
            if not success:
                return {"error": f"Registration failed: {error}"}
        else:
            persisted_display_name = payload.name

        effective_device_name = persisted_display_name or payload.name

        # Redis and session operations happen AFTER database connection is released
        await device_service.set_device_online(
            user_id=user_id,
            device_id=payload.device_id,
            socket_id=sid,
            name=effective_device_name,
            executor_version=payload.executor_version,
        )

        # Update session with device_id and device_name
        session["device_id"] = payload.device_id
        session["device_name"] = effective_device_name
        session["registered"] = True
        await self.save_session(sid, session)

        # Join device-specific room
        device_room = f"device:{user_id}:{payload.device_id}"
        await self.enter_room(sid, device_room)

        # Broadcast device online event to user room (via chat namespace)
        await self._broadcast_device_online(
            user_id, payload.device_id, effective_device_name
        )

        logger.info(
            f"[Device WS] Device registered: user={user_id}, device={payload.device_id}"
        )

        return {"success": True, "device_id": payload.device_id}

    async def on_device_heartbeat(self, sid: str, data: dict) -> dict:
        """
        Handle device:heartbeat event.

        Refreshes the device's online status in Redis and updates MySQL.

        Args:
            sid: Socket ID
            data: {"device_id": str}

        Returns:
            {"success": True} or {"error": str}
        """
        try:
            payload = DeviceHeartbeatPayload(**data)
        except Exception as e:
            return {"error": f"Invalid payload: {e}"}

        session = await self.get_session(sid)
        user_id = session.get("user_id")
        session_device_id = session.get("device_id")

        if not user_id:
            return {"error": "Not authenticated"}

        if session_device_id != payload.device_id:
            return {"error": "Device ID mismatch"}

        # Refresh Redis TTL and update running_task_ids
        success = await device_service.refresh_device_heartbeat(
            user_id,
            payload.device_id,
            payload.running_task_ids,
            payload.executor_version,
        )

        if not success:
            # Redis key expired, recreate it to recover from ghost-offline state
            device_name = session.get("device_name", f"device-{payload.device_id[:8]}")
            logger.warning(
                f"[Device WS] Heartbeat recovery: recreating Redis key for "
                f"user={user_id}, device={payload.device_id}"
            )
            await device_service.set_device_online(
                user_id=user_id,
                device_id=payload.device_id,
                socket_id=sid,
                name=device_name,
                executor_version=payload.executor_version,
            )
            # Re-broadcast device online event
            await self._broadcast_device_online(user_id, payload.device_id, device_name)

        # Database operation: quick in, quick out
        _update_device_heartbeat(user_id, payload.device_id)

        # Broadcast slot update to user
        await self._broadcast_device_slot_update(user_id, payload.device_id)

        logger.debug(
            f"[Device WS] Heartbeat received: user={user_id}, device={payload.device_id}, "
            f"running_tasks={len(payload.running_task_ids)}"
        )

        return {"success": True}

    async def on_device_status(self, sid: str, data: dict) -> dict:
        """
        Handle device:status event.

        Updates the device status (online/busy) in Redis.

        Args:
            sid: Socket ID
            data: {"device_id": str, "status": str}

        Returns:
            {"success": True} or {"error": str}
        """
        try:
            payload = DeviceStatusPayload(**data)
        except Exception as e:
            return {"error": f"Invalid payload: {e}"}

        session = await self.get_session(sid)
        user_id = session.get("user_id")
        session_device_id = session.get("device_id")

        if not user_id:
            return {"error": "Not authenticated"}

        if session_device_id != payload.device_id:
            return {"error": "Device ID mismatch"}

        # Update status in Redis
        await device_service.update_device_status_in_redis(
            user_id, payload.device_id, payload.status.value
        )

        # Broadcast status change to user room
        await self._broadcast_device_status(
            user_id, payload.device_id, payload.status.value
        )

        logger.info(
            f"[Device WS] Status updated: user={user_id}, device={payload.device_id}, status={payload.status}"
        )

        return {"success": True}

    async def on_device_upgrade_status(self, sid: str, data: dict) -> dict:
        """
        Handle device:upgrade_status event from executor.

        Receives upgrade status updates from the executor and broadcasts
        them to the user's room via the chat namespace.

        Args:
            sid: Socket ID
            data: Upgrade status data containing device_id, status, message, etc.

        Returns:
            {"success": True} or {"error": str}
        """
        try:
            from app.schemas.device import DeviceUpgradeStatusEvent

            payload = DeviceUpgradeStatusEvent(**data)
        except Exception as e:
            logger.warning(f"[Device WS] Invalid upgrade_status payload: {e}")
            return {"error": f"Invalid payload: {e}"}

        session = await self.get_session(sid)
        user_id = session.get("user_id")
        session_device_id = session.get("device_id")

        if not user_id:
            return {"error": "Not authenticated"}

        if session_device_id != payload.device_id:
            return {"error": "Device ID mismatch"}

        logger.info(
            f"[Device WS] Upgrade status: user={user_id}, device={payload.device_id}, "
            f"status={payload.status}, message={payload.message}"
        )

        # Broadcast to user room via chat namespace
        await self._broadcast_device_upgrade_status(user_id, payload)

        # If terminal state (success/error/skipped), update device metadata
        if payload.status in ["success", "error", "skipped"]:
            logger.info(
                f"[Device WS] Upgrade terminal state reached: "
                f"status={payload.status}, device={payload.device_id}"
            )

        return {"success": True}

    # ============================================================
    # OpenAI Responses API Event Handler
    # ============================================================

    async def _handle_responses_api_event(
        self, sid: str, event_type: str, *args
    ) -> dict:
        """Handle OpenAI Responses API events from local executor.

        Reuses the same processing chain as /api/internal/callback:
        ResponsesAPIEventParser -> StatusUpdatingEmitter -> WebSocketResultEmitter

        IMPORTANT: Uses per-subtask locking to ensure events are processed in order.
        This prevents race conditions when multiple response.output_text.delta events
        arrive concurrently for the same subtask, which would cause text content to
        be appended to Redis in the wrong order.

        Args:
            sid: Socket ID
            event_type: OpenAI Responses API event type (e.g., response.output_text.delta)
            *args: Event arguments (first arg is the event data dict)

        Returns:
            {"success": True} or {"error": str}
        """
        session = await self.get_session(sid)
        user_id = session.get("user_id")
        device_id = session.get("device_id")

        if not user_id or not device_id:
            return {"error": "Not authenticated or not registered"}

        if not args:
            return {"error": "Missing event data"}

        data = args[0]
        if not isinstance(data, dict):
            return {"error": "Invalid event data format"}

        task_id = data.get("task_id")
        subtask_id = data.get("subtask_id")
        event_data = data.get("data", {})
        message_id = data.get("message_id")

        if not task_id or not subtask_id:
            return {"error": "Missing task_id or subtask_id"}

        logger.debug(
            f"[Device WS] Responses API event: type={event_type}, "
            f"task_id={task_id}, subtask_id={subtask_id}"
        )

        # Get lock for this subtask to ensure events are processed in order
        # This prevents race conditions when multiple events arrive concurrently
        lock = self._get_subtask_lock(subtask_id)

        # Track whether this is a terminal event for lock cleanup
        is_terminal = False

        try:
            # Acquire lock to ensure sequential processing of events for this subtask
            async with lock:
                # Parse using shared ResponsesAPIEventParser
                event = self._event_parser.parse(
                    task_id=task_id,
                    subtask_id=subtask_id,
                    message_id=message_id,
                    event_type=event_type,
                    data=event_data,
                )

                if event is None:
                    # Lifecycle events (response.created, etc.) are skipped
                    return {"success": True}

                # Emit via StatusUpdatingEmitter -> WebSocketResultEmitter
                # Pass user_id for task:status notification
                ws_emitter = WebSocketResultEmitter(
                    task_id=task_id,
                    subtask_id=subtask_id,
                    user_id=user_id,
                )
                emitter = StatusUpdatingEmitter(
                    wrapped=ws_emitter,
                    task_id=task_id,
                    subtask_id=subtask_id,
                )
                await emitter.emit(event)
                await emitter.close()

                # Handle terminal events
                is_terminal = event.type in (
                    EventType.DONE.value,
                    EventType.ERROR.value,
                    EventType.CANCELLED.value,
                )
                if is_terminal:
                    await self._publish_task_completed_event(
                        task_id, subtask_id, user_id, device_id, event
                    )

            # Clean up lock after terminal event (outside the lock context)
            if is_terminal:
                self._cleanup_subtask_lock(subtask_id)

            return {"success": True}

        except Exception as e:
            logger.exception(
                f"[Device WS] Error handling Responses API event: "
                f"type={event_type}, subtask_id={subtask_id}, error={e}"
            )
            # Clean up lock on error to prevent memory leak
            self._cleanup_subtask_lock(subtask_id)
            return {"error": str(e)}

    async def _publish_task_completed_event(
        self,
        task_id: int,
        subtask_id: int,
        user_id: int,
        device_id: str,
        event: Any,
    ) -> None:
        """Publish TaskCompletedEvent and broadcast slot update for terminal events."""
        try:
            if event.type == EventType.DONE.value:
                status = "COMPLETED"
                result = event.result if hasattr(event, "result") else None
                error = None
            elif event.type == EventType.ERROR.value:
                status = "FAILED"
                result = None
                error = event.error if hasattr(event, "error") else "Unknown error"
            else:
                status = "CANCELLED"
                result = None
                error = None

            event_bus = get_event_bus()
            await event_bus.publish(
                TaskCompletedEvent(
                    task_id=task_id,
                    subtask_id=subtask_id,
                    user_id=user_id,
                    status=status,
                    result=result,
                    error=error,
                )
            )

            logger.info(
                f"[Device WS] Published TaskCompletedEvent: "
                f"task_id={task_id}, subtask_id={subtask_id}, status={status}"
            )

            # Broadcast slot update after task completion
            await self._broadcast_device_slot_update(user_id, device_id)

        except Exception as e:
            logger.error(
                f"[Device WS] Failed to publish TaskCompletedEvent: {e}",
                exc_info=True,
            )

    # ============================================================
    # Broadcast Helpers
    # ============================================================

    async def _broadcast_device_online(
        self, user_id: int, device_id: str, name: str
    ) -> None:
        """Broadcast device:online event to user room via chat namespace."""
        from app.core.socketio import get_sio
        from app.schemas.device import DeviceStatusEnum

        sio = get_sio()
        event_data = DeviceOnlineEvent(
            device_id=device_id,
            name=name,
            status=DeviceStatusEnum.ONLINE,
        ).model_dump()

        await sio.emit(
            "device:online",
            event_data,
            room=f"user:{user_id}",
            namespace="/chat",
        )
        logger.debug(f"[Device WS] Broadcast device:online to user:{user_id}")

    async def _broadcast_device_offline(self, user_id: int, device_id: str) -> None:
        """Broadcast device:offline event to user room via chat namespace."""
        from app.core.socketio import get_sio

        sio = get_sio()
        event_data = DeviceOfflineEvent(device_id=device_id).model_dump()

        await sio.emit(
            "device:offline",
            event_data,
            room=f"user:{user_id}",
            namespace="/chat",
        )
        logger.debug(f"[Device WS] Broadcast device:offline to user:{user_id}")

    async def _broadcast_device_status(
        self, user_id: int, device_id: str, status: str
    ) -> None:
        """Broadcast device:status event to user room via chat namespace."""
        from app.core.socketio import get_sio
        from app.schemas.device import DeviceStatusEnum

        sio = get_sio()
        # Convert string status to enum
        status_enum = DeviceStatusEnum(status)
        event_data = DeviceStatusEvent(
            device_id=device_id, status=status_enum
        ).model_dump()

        await sio.emit(
            "device:status",
            event_data,
            room=f"user:{user_id}",
            namespace="/chat",
        )
        logger.debug(
            f"[Device WS] Broadcast device:status to user:{user_id}, status={status}"
        )

    async def _broadcast_device_slot_update(self, user_id: int, device_id: str) -> None:
        """
        Broadcast device:slot_update event to user room via chat namespace.

        Queries current slot usage and emits the update.
        Uses run_sync_in_executor to avoid blocking the event loop.
        """
        from app.core.socketio import get_sio
        from app.schemas.device import DeviceRunningTask

        try:
            # Run database query in executor to avoid blocking event loop
            slot_info = await run_sync_in_executor(
                _get_device_slot_usage_sync, user_id, device_id
            )

            sio = get_sio()
            event_data = DeviceSlotUpdateEvent(
                device_id=device_id,
                slot_used=slot_info["used"],
                slot_max=slot_info["max"],
                running_tasks=[
                    DeviceRunningTask(**task) for task in slot_info["running_tasks"]
                ],
            ).model_dump()

            await sio.emit(
                "device:slot_update",
                event_data,
                room=f"user:{user_id}",
                namespace="/chat",
            )
            logger.debug(
                f"[Device WS] Broadcast device:slot_update to user:{user_id}, "
                f"slot_used={slot_info['used']}"
            )
        except Exception as e:
            logger.error(f"[Device WS] Error broadcasting slot update: {e}")

    async def emit_upgrade_command(self, socket_id: str, params: dict) -> bool:
        """
        Emit device:upgrade command to a specific device.

        This method is called from the internal API to trigger a remote upgrade.

        Args:
            socket_id: The Socket.IO session ID of the target device
            params: Upgrade parameters (force, auto_confirm, verbose, etc.)

        Returns:
            True if the command was emitted successfully, False otherwise
        """
        try:
            await self.emit(
                "device:upgrade",
                params,
                room=socket_id,
            )
            logger.info(f"[Device WS] Sent upgrade command to socket {socket_id}")
            return True
        except Exception as e:
            logger.error(f"[Device WS] Failed to send upgrade command: {e}")
            return False

    async def _broadcast_device_upgrade_status(
        self, user_id: int, payload: "DeviceUpgradeStatusEvent"
    ) -> None:
        """
        Broadcast device:upgrade_status event to owner and admin rooms via chat namespace.

        Args:
            user_id: Device owner's user ID
            payload: DeviceUpgradeStatusEvent payload
        """
        from app.core.socketio import get_sio

        try:
            sio = get_sio()
            event_data = payload.model_dump()
            target_user_ids = {user_id}

            with _db_session() as db:
                admin_user_ids = (
                    db.query(User.id)
                    .filter(User.role == "admin", User.is_active == True)
                    .all()
                )
                target_user_ids.update(admin_id for (admin_id,) in admin_user_ids)

            for target_user_id in target_user_ids:
                await sio.emit(
                    "device:upgrade_status",
                    event_data,
                    room=f"user:{target_user_id}",
                    namespace="/chat",
                )
            logger.debug(
                f"[Device WS] Broadcast device:upgrade_status to users={sorted(target_user_ids)}, "
                f"device={payload.device_id}, status={payload.status}"
            )
        except Exception as e:
            logger.error(f"[Device WS] Error broadcasting upgrade status: {e}")


# Global singleton instance (initialized by register_device_namespace)
device_namespace: Optional[DeviceNamespace] = None


# Factory function to create the namespace
def create_device_namespace() -> DeviceNamespace:
    """Create and return a DeviceNamespace instance."""
    return DeviceNamespace()


def register_device_namespace(sio: socketio.AsyncServer) -> None:
    """
    Register the device namespace with the Socket.IO server.

    Args:
        sio: Socket.IO server instance
    """
    global device_namespace
    device_namespace = DeviceNamespace("/local-executor")
    sio.register_namespace(device_namespace)
    logger.info("Device namespace registered at /local-executor")
