import socket
from collections.abc import AsyncGenerator
from typing import Optional, Union

from pymobiledevice3.exceptions import NotificationTimeoutError
from pymobiledevice3.lockdown_service_provider import LockdownServiceProvider
from pymobiledevice3.remote.remote_service_discovery import RemoteServiceDiscoveryService
from pymobiledevice3.services.lockdown_service import LockdownService


class NotificationProxyService(LockdownService):
    SERVICE_NAME = "com.apple.mobile.notification_proxy"
    RSD_SERVICE_NAME = "com.apple.mobile.notification_proxy.shim.remote"

    INSECURE_SERVICE_NAME = "com.apple.mobile.insecure_notification_proxy"
    RSD_INSECURE_SERVICE_NAME = "com.apple.mobile.insecure_notification_proxy.shim.remote"

    def __init__(self, lockdown: LockdownServiceProvider, insecure=False, timeout: Optional[Union[float, int]] = None):
        if isinstance(lockdown, RemoteServiceDiscoveryService):
            secure_service_name = self.RSD_SERVICE_NAME
            insecure_service_name = self.RSD_INSECURE_SERVICE_NAME
        else:
            secure_service_name = self.SERVICE_NAME
            insecure_service_name = self.INSECURE_SERVICE_NAME

        if insecure:
            super().__init__(lockdown, insecure_service_name)
        else:
            super().__init__(lockdown, secure_service_name)

        if timeout is not None:
            self.service.socket.settimeout(timeout)

    async def notify_post(self, name: str) -> None:
        """Send notification to the device's notification_proxy."""
        await self.service.send_plist({"Command": "PostNotification", "Name": name})

    async def notify_register_dispatch(self, name: str) -> None:
        """Tells the device to send a notification on the specified event."""
        self.logger.info(f"Observing {name}")
        await self.service.send_plist({"Command": "ObserveNotification", "Name": name})

    async def receive_notification(self) -> AsyncGenerator[dict, None]:
        while True:
            try:
                yield await self.service.recv_plist()
            except socket.timeout as e:
                raise NotificationTimeoutError from e
