import dataclasses
import typing
from typing import Optional

from pymobiledevice3.exceptions import DisableMemoryLimitError
from pymobiledevice3.osu.os_utils import get_os_utils
from pymobiledevice3.services.dvt.dvt_secure_socket_proxy import DvtSecureSocketProxyService
from pymobiledevice3.services.dvt.instruments import ChannelService
from pymobiledevice3.services.remote_server import MessageAux

OSUTIL = get_os_utils()


@dataclasses.dataclass
class OutputReceivedEvent:
    pid: int
    date: int
    message: str

    @classmethod
    def create(cls, message) -> "OutputReceivedEvent":
        try:
            date = OSUTIL.parse_timestamp(message[2].value)
        except (ValueError, OSError):
            date = None

        return cls(pid=message[1].value, date=date, message=message[0].value)


class ProcessControl(ChannelService):
    IDENTIFIER = "com.apple.instruments.server.services.processcontrol"

    def __init__(self, dvt: DvtSecureSocketProxyService):
        super().__init__(dvt)

    async def signal(self, pid: int, sig: int):
        """
        Send signal to process
        :param pid: PID of process to send signal.
        :param sig: SIGNAL to send
        """
        channel = await self._channel_ref()
        await channel.sendSignal_toPid_(MessageAux().append_obj(sig).append_obj(pid), expects_reply=True)
        return await channel.receive_plist()

    async def disable_memory_limit_for_pid(self, pid: int) -> None:
        """
        Waive memory limit for a given pid
        :param pid: process id.
        """
        channel = await self._channel_ref()
        await channel.requestDisableMemoryLimitsForPid_(MessageAux().append_int(pid), expects_reply=True)
        if not await channel.receive_plist():
            raise DisableMemoryLimitError()

    async def kill(self, pid: int):
        """
        Kill a process.
        :param pid: PID of process to kill.
        """
        channel = await self._channel_ref()
        await channel.killPid_(MessageAux().append_obj(pid), expects_reply=False)

    async def process_identifier_for_bundle_identifier(self, app_bundle_identifier: str) -> int:
        channel = await self._channel_ref()
        await channel.processIdentifierForBundleIdentifier_(
            MessageAux().append_obj(app_bundle_identifier), expects_reply=True
        )
        return await channel.receive_plist()

    async def launch(
        self,
        bundle_id: str,
        arguments=None,
        kill_existing: bool = True,
        start_suspended: bool = False,
        environment: Optional[dict] = None,
        extra_options: Optional[dict] = None,
    ) -> int:
        """
        Launch a process.
        :param bundle_id: Bundle id of the process.
        :param list arguments: List of argument to pass to process.
        :param kill_existing: Whether to kill an existing instance of this process.
        :param start_suspended: Same as WaitForDebugger.
        :param environment: Environment variables to pass to process.
        :param extra_options: Extra options to pass to process.
        :return: PID of created process.
        """
        arguments = [] if arguments is None else arguments
        environment = {} if environment is None else environment
        options = {
            "StartSuspendedKey": start_suspended,
            "KillExisting": kill_existing,
        }
        if extra_options:
            options.update(extra_options)
        args = (
            MessageAux()
            .append_obj("")
            .append_obj(bundle_id)
            .append_obj(environment)
            .append_obj(arguments)
            .append_obj(options)
        )
        channel = await self._channel_ref()
        await channel.launchSuspendedProcessWithDevicePath_bundleIdentifier_environment_arguments_options_(args)
        result = await channel.receive_plist()
        assert result
        return result

    async def __aiter__(self) -> typing.AsyncGenerator[OutputReceivedEvent, None]:
        while True:
            key, value = await self._receive_key_value()
            if key == "outputReceived:fromProcess:atTime:":
                yield OutputReceivedEvent.create(value)

    async def _receive_key_value(self):
        channel = await self._channel_ref()
        return await channel.receive_key_value()
