Source code for alab_management.user_input

import time
from datetime import datetime
from enum import Enum
from typing import Any, cast

from bson import ObjectId

from alab_management.alarm import Alarm
from alab_management.experiment_view.experiment_view import ExperimentView
from alab_management.task_view import TaskCancelledError, TaskView
from alab_management.utils.data_objects import get_collection

from .config import AlabOSConfig

CANCEL_RESPONSE = "Cancelled"
SUPERSEDED_RESPONSE = "Superseded"




class UserRequestStatus(Enum):
    """Enum for user response."""

    FULLFILLED = "fulfilled"
    PENDING = "pending"
    ERROR = "error"





class UserInputView:
    """Sample view manages the samples and their positions."""

    def __init__(self):
        self._input_collection = get_collection("user_input")
        self._task_view = TaskView()
        self._experiment_view = ExperimentView()
        alarm_config = AlabOSConfig().get("alarm", {})
        self._alarm = Alarm(**alarm_config)



    def insert_request(
        self,
        prompt: str,
        options: list[str],
        task_id: ObjectId | None = None,
        maintenance: bool = False,
        category: str = "Unknown Category",
        request_context_extra: dict[str, Any] | None = None,
    ) -> ObjectId:
        """
        Insert a request into the database.

        Returns the request ObjectID
        """
        context = {"maintenance": maintenance}
        if task_id is None and not maintenance:
            raise ValueError("task_id is required for non-maintenance requests!")
        if task_id is not None:
            self._task_view.get_task(
                task_id=task_id
            )  # will throw error if task id does not exist
            experiment_id = self._experiment_view.get_experiment_by_task_id(task_id)[
                "_id"
            ]
            context.update(
                {
                    "experiment_id": experiment_id,
                    "task_id": task_id,
                }
            )
        if request_context_extra:
            context.update(request_context_extra)

        # Keep the User Inputs UI to one actionable row per duplicate: when a fresh
        # request replaces an older pending one with the same device + prompt (typical
        # orphaned arm "set to remote" piles after reconnect/restart), fulfill the
        # stale rows so only the new waiter is shown.
        self._supersede_duplicate_pending_requests(
            prompt=prompt,
            context=context,
        )

        request_id = ObjectId()
        self._input_collection.insert_one(
            {
                "_id": request_id,
                "prompt": prompt,
                "options": [str(opt) for opt in options],
                "status": UserRequestStatus.PENDING.value,
                "request_context": context,
                "last_updated": datetime.now(),
            }
        )
        if maintenance is True:
            category = "Maintenance"
        self._alarm.alert(f"User input requested: {prompt}", category)
        return request_id


    def _supersede_duplicate_pending_requests(
        self, *, prompt: str, context: dict[str, Any]
    ) -> int:
        """Fulfill pending requests that a new insert would duplicate.

        Matches on identical ``prompt`` plus ``request_context.device`` when a device
        is present (robot-arm / instrument maintenance prompts). Returns how many
        stale requests were superseded.
        """
        device = context.get("device")
        if not device:
            return 0

        now = datetime.now()
        result = self._input_collection.update_many(
            {
                "status": UserRequestStatus.PENDING.value,
                "prompt": prompt,
                "request_context.device": device,
            },
            {
                "$set": {
                    "status": UserRequestStatus.FULLFILLED.value,
                    "response": SUPERSEDED_RESPONSE,
                    "note": (
                        "Replaced by a newer request for the same device and prompt."
                    ),
                    "last_updated": now,
                }
            },
        )
        return int(result.modified_count)



    def get_request(self, request_id: ObjectId) -> dict[str, Any]:
        """
        Get a request.

        Returns a request.
        """
        request = self._input_collection.find_one({"_id": request_id})

        if request is None:
            raise ValueError(f"User input request id {request_id} does not exist!")
        return cast(dict[str, Any], request)




    def update_request_status(self, request_id: ObjectId, response: str, note: str):
        """Update the status of a request."""
        self.get_request(request_id)  # will error is request does not exist
        self._input_collection.update_one(
            {"_id": request_id},
            {
                "$set": {
                    "response": response,
                    "note": note,
                    "status": UserRequestStatus.FULLFILLED.value,
                    "last_updated": datetime.now(),
                }
            },
        )




    def retrieve_user_input(
        self, request_id: ObjectId, *, abort_if_task_cancelled: bool = True
    ) -> str:
        """
        Retrive response from user for a given request. Blocks until request is marked as completed.

        Returns the user response, which is one of a list of options.

        ``abort_if_task_cancelled`` is the normal case: a task waiting for an operator
        should stop if someone cancelled it. Cleanup prompts pass ``False`` because they
        exist *because* the task is being torn down — aborting them leaves the worker
        unable to start.
        """
        request = None
        try:
            while True:
                request = self.get_request(request_id)
                if abort_if_task_cancelled:
                    self._raise_if_owning_task_cancelled(request)
                if UserRequestStatus(request["status"]) != UserRequestStatus.PENDING:
                    break
                time.sleep(0.5)
        except TaskCancelledError:
            raise
        except:  # noqa: E722
            self._input_collection.update_one(
                {"_id": request_id}, {"$set": {"status": UserRequestStatus.ERROR.name}}
            )
            raise
        if abort_if_task_cancelled:
            self._raise_if_owning_task_cancelled(request)
        return request["response"]


    def _raise_if_owning_task_cancelled(self, request: dict[str, Any]) -> None:
        """Abort a blocked user-input wait when the owning task is being cancelled."""
        task_id = (request.get("request_context") or {}).get("task_id")
        if task_id is None or not self._task_view.is_canceling(task_id):
            return
        if UserRequestStatus(request["status"]) == UserRequestStatus.PENDING:
            self.update_request_status(
                request_id=request["_id"],
                response=CANCEL_RESPONSE,
                note="Dismissed because the task was cancelled.",
            )
        raise TaskCancelledError("Cancelled while waiting for user input.")



    def dismiss_pending_requests(
        self,
        *,
        experiment_id: ObjectId | None = None,
        task_id: ObjectId | None = None,
        response: str = CANCEL_RESPONSE,
        note: str = "Dismissed because the experiment or task was cancelled.",
    ) -> int:
        """Fulfill pending prompts so a cancelled task is not stuck waiting for a click."""
        if experiment_id is None and task_id is None:
            raise ValueError("experiment_id or task_id is required")
        query: dict[str, Any] = {"status": UserRequestStatus.PENDING.value}
        if experiment_id is not None:
            query["request_context.experiment_id"] = experiment_id
        if task_id is not None:
            query["request_context.task_id"] = task_id
        result = self._input_collection.update_many(
            query,
            {
                "$set": {
                    "status": UserRequestStatus.FULLFILLED.value,
                    "response": response,
                    "note": note,
                    "last_updated": datetime.now(),
                }
            },
        )
        return result.modified_count




    def clean_up_user_input_collection(self):
        """Drop the sample position collection."""
        self._input_collection.drop()




    def get_all_pending_requests(self) -> list:
        """
        Get all pending requests.

        Collapses duplicate device+prompt rows (keeps newest) so the dashboard does
        not show a pile of stale identical maintenance prompts.
        """
        self._collapse_duplicate_pending_requests()
        return cast(
            list[dict[str, Any]],
            list(
                self._input_collection.find(
                    {"status": UserRequestStatus.PENDING.value}
                )
            ),
        )


    def _collapse_duplicate_pending_requests(self) -> int:
        """Supersede older pending requests that share device + prompt with a newer one."""
        pending = list(
            self._input_collection.find({"status": UserRequestStatus.PENDING.value})
        )
        newest_by_key: dict[tuple[str, str], dict[str, Any]] = {}
        superseded_ids: list[ObjectId] = []
        for request in pending:
            context = request.get("request_context") or {}
            device = context.get("device")
            if not device:
                continue
            key = (str(device), str(request.get("prompt") or ""))
            current = newest_by_key.get(key)
            if current is None:
                newest_by_key[key] = request
                continue
            current_updated = current.get("last_updated") or datetime.min
            request_updated = request.get("last_updated") or datetime.min
            if request_updated >= current_updated:
                superseded_ids.append(current["_id"])
                newest_by_key[key] = request
            else:
                superseded_ids.append(request["_id"])

        if not superseded_ids:
            return 0
        now = datetime.now()
        result = self._input_collection.update_many(
            {
                "_id": {"$in": superseded_ids},
                "status": UserRequestStatus.PENDING.value,
            },
            {
                "$set": {
                    "status": UserRequestStatus.FULLFILLED.value,
                    "response": SUPERSEDED_RESPONSE,
                    "note": (
                        "Replaced by a newer request for the same device and prompt."
                    ),
                    "last_updated": now,
                }
            },
        )
        return int(result.modified_count)



    def get_pending_request_by_context(
        self, context_filters: dict[str, Any]
    ) -> dict[str, Any] | None:
        """Return one pending request whose ``request_context`` matches all filters."""
        query: dict[str, Any] = {"status": UserRequestStatus.PENDING.value}
        for key, value in context_filters.items():
            query[f"request_context.{key}"] = value
        return self._input_collection.find_one(query)




    def retrieve_user_input_with_note(
        self, request_id: ObjectId, *, abort_if_task_cancelled: bool = True
    ) -> tuple[str, str]:
        """
        Retrive response from user for a given request. Blocks until request is marked as completed.

        Returns the user response, which is one of a list of options
        """
        request = None
        try:
            while True:
                request = self.get_request(request_id)
                if abort_if_task_cancelled:
                    self._raise_if_owning_task_cancelled(request)
                if UserRequestStatus(request["status"]) != UserRequestStatus.PENDING:
                    break
                time.sleep(0.5)
        except TaskCancelledError:
            raise
        except:  # noqa: E722
            self._input_collection.update_one(
                {"_id": request_id}, {"$set": {"status": UserRequestStatus.ERROR.name}}
            )
            raise
        if abort_if_task_cancelled:
            self._raise_if_owning_task_cancelled(request)
        return request["response"], request["note"]






def request_user_input(
    task_id: ObjectId | None,
    prompt: str,
    options: list[str],
    maintenance: bool = False,
    category: str = "Unknown Category",
    request_context_extra: dict[str, Any] | None = None,
    abort_if_task_cancelled: bool = True,
) -> str:
    """
    Request user input through the dashboard. Blocks until response is given.

    Args:
        task_id (ObjectId): task id requesting user input
        prompt (str): prompt to give user
        options (List[str]): response options to give user
        maintenance (bool): if true, mark this as a request for overall system maintenance
        request_context_extra (Optional[dict]): extra context stored on the request, used to
          associate it with whatever raised it (e.g. ``{"device": "DASH_arm_JEEVES"}``)

    Returns
    -------
        response (str): user response as string
    """
    user_input_view = UserInputView()
    while True:
        request_id = user_input_view.insert_request(
            task_id=task_id,
            prompt=prompt,
            options=options,
            maintenance=maintenance,
            category=category,
            request_context_extra=request_context_extra,
        )
        response = user_input_view.retrieve_user_input(
            request_id=request_id,
            abort_if_task_cancelled=abort_if_task_cancelled,
        )
        if response != SUPERSEDED_RESPONSE:
            return response

        # Another insert replaced this request; wait on the surviving pending row
        # instead of returning "Superseded" to the caller.
        device = (request_context_extra or {}).get("device")
        if not device:
            continue
        replacement = user_input_view._input_collection.find_one(
            {
                "status": UserRequestStatus.PENDING.value,
                "prompt": prompt,
                "request_context.device": device,
            }
        )
        if replacement is None:
            continue
        response = user_input_view.retrieve_user_input(
            request_id=replacement["_id"],
            abort_if_task_cancelled=abort_if_task_cancelled,
        )
        if response != SUPERSEDED_RESPONSE:
            return response





def request_maintenance_input(
    prompt: str,
    options: list[str],
    request_context_extra: dict[str, Any] | None = None,
):
    """
    Request user input for maintenance through the dashboard. Blocks until response is given.

    Args:
        prompt (str): prompt to give user
        options (List[str]): response options to give user
        request_context_extra (Optional[dict]): extra context stored on the request, used to
          associate it with whatever raised it (e.g. ``{"device": "DASH_arm_JEEVES"}``)

    Returns
    -------
        response (str): user response as string
    """
    return request_user_input(
        task_id=None,
        prompt=prompt,
        options=options,
        maintenance=True,
        category="Maintenance",
        request_context_extra=request_context_extra,
    )





def request_user_input_with_note(
    task_id: ObjectId | None,
    prompt: str,
    options: list[str],
    maintenance: bool = False,
    category: str = "Unknown Category",
) -> tuple[str, str]:
    """
    Request user input through the dashboard. Blocks until response. Returns response and note.

    Args:
        task_id (ObjectId): task id requesting user input
        prompt (str): prompt to give user
        options (List[str]): response options to give user
        maintenance (bool): if true, mark this as a request for overall system maintenance

    Returns
    -------
        response (str): user response as string
        note (str): note from the user
    """
    user_input_view = UserInputView()
    request_id = user_input_view.insert_request(
        task_id=task_id,
        prompt=prompt,
        options=options,
        maintenance=maintenance,
        category=category,
    )
    return user_input_view.retrieve_user_input_with_note(request_id=request_id)