#!/usr/bin/env python
"""Clients-related part of GRR API client library."""

from typing import Optional

from grr_api_client import context as api_context
from grr_api_client import utils
from grr_response_proto.api import user_pb2


class Notification(object):
  """GRR user notification object with fetched data."""

  def __init__(
      self,
      data: user_pb2.ApiNotification,
      context: api_context.GrrApiContext,
  ):
    self.data: user_pb2.ApiNotification = data
    self._context: api_context.GrrApiContext = context


class GrrUser(object):
  """GRR user object describing the current API user."""

  def __init__(
      self,
      context: api_context.GrrApiContext,
  ):
    self._context: api_context.GrrApiContext = context

  @property
  def username(self) -> str:
    return self._context.username

  def GetPendingNotificationsCount(self) -> int:
    response = self._context.SendRequest(
        "GetPendingUserNotificationsCount", None
    )
    if not isinstance(
        response, user_pb2.ApiGetPendingUserNotificationsCountResult
    ):
      raise TypeError(f"Unexpected response type: {type(response)}")

    return response.count

  def ListPendingNotifications(
      self,
      timestamp: Optional[int] = None,
  ) -> utils.ItemsIterator[Notification]:
    """Lists pending notifications for the user."""
    args = user_pb2.ApiListPendingUserNotificationsArgs()
    if timestamp is not None:
      args.timestamp = timestamp

    items = self._context.SendIteratorRequest(
        "ListPendingUserNotifications", args
    )
    return utils.MapItemsIterator(
        lambda data: Notification(data=data, context=self._context), items
    )