"""Classes to manage ADB connections.

* :py:class:`ADBPythonAsync` utilizes a Python implementation of the ADB protocol.
* :py:class:`ADBServerAsync` utilizes an ADB server to communicate with the device.

"""


import asyncio
from contextlib import asynccontextmanager
import logging

from adb_shell.adb_device import AdbDeviceUsb
from adb_shell.adb_device_async import AdbDeviceTcpAsync
from adb_shell.auth.sign_pythonrsa import PythonRSASigner
from adb_shell.constants import DEFAULT_PUSH_MODE, DEFAULT_READ_TIMEOUT_S
import aiofiles
import async_timeout
from ppadb.client import Client

from ..constants import (
    DEFAULT_ADB_TIMEOUT_S,
    DEFAULT_AUTH_TIMEOUT_S,
    DEFAULT_LOCK_TIMEOUT_S,
    DEFAULT_TRANSPORT_TIMEOUT_S,
)
from ..exceptions import LockNotAcquiredException

_LOGGER = logging.getLogger(__name__)


class AdbDeviceUsbAsync:
    """An async wrapper for the adb-shell ``AdbDeviceUsb`` class."""

    def __init__(self, serial=None, port_path=None, default_transport_timeout_s=None, banner=None):
        self._adb = AdbDeviceUsb(serial, port_path, default_transport_timeout_s, banner)

    @property
    def available(self):
        """Whether or not an ADB connection to the device has been established."""
        return self._adb.available

    async def close(self):
        """Close the connection via the provided transport's ``close()`` method."""
        await asyncio.get_running_loop().run_in_executor(None, self._adb.close)

    async def connect(
        self,
        rsa_keys=None,
        transport_timeout_s=None,
        auth_timeout_s=DEFAULT_AUTH_TIMEOUT_S,
        read_timeout_s=DEFAULT_READ_TIMEOUT_S,
        auth_callback=None,
    ):
        """Establish an ADB connection to the device."""
        return await asyncio.get_running_loop().run_in_executor(
            None, self._adb.connect, rsa_keys, transport_timeout_s, auth_timeout_s, read_timeout_s, auth_callback
        )

    async def pull(
        self,
        device_path,
        local_path,
        progress_callback=None,
        transport_timeout_s=None,
        read_timeout_s=DEFAULT_READ_TIMEOUT_S,
    ):
        """Pull a file from the device."""
        await asyncio.get_running_loop().run_in_executor(
            None, self._adb.pull, device_path, local_path, progress_callback, transport_timeout_s, read_timeout_s
        )

    async def push(
        self,
        local_path,
        device_path,
        st_mode=DEFAULT_PUSH_MODE,
        mtime=0,
        progress_callback=None,
        transport_timeout_s=None,
        read_timeout_s=DEFAULT_READ_TIMEOUT_S,
    ):
        """Push a file or directory to the device."""
        await asyncio.get_running_loop().run_in_executor(
            None,
            self._adb.push,
            local_path,
            device_path,
            st_mode,
            mtime,
            progress_callback,
            transport_timeout_s,
            read_timeout_s,
        )

    async def shell(
        self, command, transport_timeout_s=None, read_timeout_s=DEFAULT_READ_TIMEOUT_S, timeout_s=None, decode=True
    ):
        """Send an ADB shell command to the device."""
        return await asyncio.get_running_loop().run_in_executor(
            None, self._adb.shell, command, transport_timeout_s, read_timeout_s, timeout_s, decode
        )


class DeviceAsync:
    """An async wrapper for the pure-python-adb ``Device`` class."""

    def __init__(self, device):
        self._device = device

    async def pull(self, device_path, local_path):
        """Download a file."""
        return await asyncio.get_running_loop().run_in_executor(None, self._device.pull, device_path, local_path)

    async def push(self, local_path, device_path):
        """Upload a file."""
        return await asyncio.get_running_loop().run_in_executor(None, self._device.push, local_path, device_path)

    async def screencap(self):
        """Take a screencap."""
        return await asyncio.get_running_loop().run_in_executor(None, self._device.screencap)

    async def shell(self, cmd):
        """Send a shell command."""
        return await asyncio.get_running_loop().run_in_executor(None, self._device.shell, cmd)


# pylint: disable=too-few-public-methods
class ClientAsync:
    """An async wrapper for the pure-python-adb ``Client`` class."""

    def __init__(self, host, port):
        self._client = Client(host, port)

    async def device(self, serial):
        """Get a ``DeviceAsync`` instance."""
        dev = await asyncio.get_running_loop().run_in_executor(None, self._client.device, serial)
        if dev:
            return DeviceAsync(dev)
        return None


@asynccontextmanager
async def _acquire(lock, timeout=DEFAULT_LOCK_TIMEOUT_S):
    """Handle acquisition and release of an ``asyncio.Lock`` object with a timeout.

    Parameters
    ----------
    lock : asyncio.Lock
        The lock that we will try to acquire
    timeout : float
        The timeout in seconds

    Yields
    ------
    acquired : bool
        Whether or not the lock was acquired

    Raises
    ------
    LockNotAcquiredException
        Raised if the lock was not acquired

    """
    try:
        acquired = False
        try:
            async with async_timeout.timeout(timeout):
                acquired = await lock.acquire()
            if not acquired:
                raise LockNotAcquiredException
            yield acquired

        except asyncio.TimeoutError as exc:
            raise LockNotAcquiredException from exc

    finally:
        if acquired:
            lock.release()


class ADBPythonAsync(object):
    """A manager for ADB connections that uses a Python implementation of the ADB protocol.

    Parameters
    ----------
    host : str
        The address of the device; may be an IP address or a host name
    port : int
        The device port to which we are connecting (default is 5555)
    adbkey : str
        The path to the ``adbkey`` file for ADB authentication
    signer : PythonRSASigner, None
        The signer for the ADB keys, as loaded by :meth:`ADBPythonAsync.load_adbkey`

    """

    def __init__(self, host, port, adbkey="", signer=None):
        self.host = host
        self.port = int(port)
        self.adbkey = adbkey

        if host:
            self._adb = AdbDeviceTcpAsync(
                host=self.host, port=self.port, default_transport_timeout_s=DEFAULT_ADB_TIMEOUT_S
            )
        else:
            self._adb = AdbDeviceUsbAsync(default_transport_timeout_s=DEFAULT_ADB_TIMEOUT_S)

        self._signer = signer

        # use a lock to make sure that ADB commands don't overlap
        self._adb_lock = asyncio.Lock()

    @property
    def available(self):
        """Check whether the ADB connection is intact.

        Returns
        -------
        bool
            Whether or not the ADB connection is intact

        """
        return self._adb.available

    async def close(self):
        """Close the ADB socket connection."""
        await self._adb.close()

    async def connect(
        self,
        log_errors=True,
        auth_timeout_s=DEFAULT_AUTH_TIMEOUT_S,
        transport_timeout_s=DEFAULT_TRANSPORT_TIMEOUT_S,
    ):
        """Connect to an Android TV / Fire TV device.

        Parameters
        ----------
        log_errors : bool
            Whether errors should be logged
        auth_timeout_s : float
            Authentication timeout (in seconds)
        transport_timeout_s : float
            Transport timeout (in seconds)

        Returns
        -------
        bool
            Whether or not the connection was successfully established and the device is available

        """
        try:
            async with _acquire(self._adb_lock):
                # Catch exceptions
                try:
                    # Connect with authentication
                    if self.adbkey:
                        if not self._signer:
                            self._signer = await self.load_adbkey(self.adbkey)

                        await self._adb.connect(
                            rsa_keys=[self._signer],
                            transport_timeout_s=transport_timeout_s,
                            auth_timeout_s=auth_timeout_s,
                        )

                    # Connect without authentication
                    else:
                        await self._adb.connect(transport_timeout_s=transport_timeout_s, auth_timeout_s=auth_timeout_s)

                    # ADB connection successfully established
                    _LOGGER.debug("ADB connection to %s:%d successfully established", self.host, self.port)
                    return True

                except OSError as exc:
                    if log_errors:
                        if exc.strerror is None:
                            exc.strerror = "Timed out trying to connect to ADB device."
                        _LOGGER.warning(
                            "Couldn't connect to %s:%d.  %s: %s",
                            self.host,
                            self.port,
                            exc.__class__.__name__,
                            exc.strerror,
                        )

                    # ADB connection attempt failed
                    await self.close()
                    return False

                except Exception as exc:  # pylint: disable=broad-except
                    if log_errors:
                        _LOGGER.warning(
                            "Couldn't connect to %s:%d.  %s: %s", self.host, self.port, exc.__class__.__name__, exc
                        )

                    # ADB connection attempt failed
                    await self.close()
                    return False

        except LockNotAcquiredException:
            _LOGGER.warning("Couldn't connect to %s:%d because adb-shell lock not acquired.", self.host, self.port)
            await self.close()
            return False

    @staticmethod
    async def load_adbkey(adbkey):
        """Load the ADB keys.

        Parameters
        ----------
        adbkey : str
            The path to the ``adbkey`` file for ADB authentication

        Returns
        -------
        PythonRSASigner
            The ``PythonRSASigner`` with the key files loaded

        """
        # private key
        async with aiofiles.open(adbkey) as f:
            priv = await f.read()

        # public key
        try:
            async with aiofiles.open(adbkey + ".pub") as f:
                pub = await f.read()
        except FileNotFoundError:
            pub = ""

        return PythonRSASigner(pub, priv)

    async def pull(self, local_path, device_path):
        """Pull a file from the device using the Python ADB implementation.

        Parameters
        ----------
        local_path : str
            The path where the file will be saved
        device_path : str
            The file on the device that will be pulled

        """
        if not self.available:
            _LOGGER.debug(
                "ADB command not sent to %s:%d because adb-shell connection is not established: pull(%s, %s)",
                self.host,
                self.port,
                local_path,
                device_path,
            )
            return

        async with _acquire(self._adb_lock):
            _LOGGER.debug(
                "Sending command to %s:%d via adb-shell: pull(%s, %s)", self.host, self.port, local_path, device_path
            )
            await self._adb.pull(device_path, local_path)
            return

    async def push(self, local_path, device_path):
        """Push a file to the device using the Python ADB implementation.

        Parameters
        ----------
        local_path : str
            The file that will be pushed to the device
        device_path : str
            The path where the file will be saved on the device

        """
        if not self.available:
            _LOGGER.debug(
                "ADB command not sent to %s:%d because adb-shell connection is not established: push(%s, %s)",
                self.host,
                self.port,
                local_path,
                device_path,
            )
            return

        async with _acquire(self._adb_lock):
            _LOGGER.debug(
                "Sending command to %s:%d via adb-shell: push(%s, %s)", self.host, self.port, local_path, device_path
            )
            await self._adb.push(local_path, device_path)
            return

    async def screencap(self):
        """Take a screenshot using the Python ADB implementation.

        Returns
        -------
        bytes
            The screencap as a binary .png image

        """
        if not self.available:
            _LOGGER.debug(
                "ADB screencap not taken from %s:%d because adb-shell connection is not established",
                self.host,
                self.port,
            )
            return None

        async with _acquire(self._adb_lock):
            _LOGGER.debug("Taking screencap from %s:%d via adb-shell", self.host, self.port)
            result = await self._adb.shell("screencap -p", decode=False)
            if result and result[5:6] == b"\r":
                return result.replace(b"\r\n", b"\n")
            return result

    async def shell(self, cmd):
        """Send an ADB command using the Python ADB implementation.

        Parameters
        ----------
        cmd : str
            The ADB command to be sent

        Returns
        -------
        str, None
            The response from the device, if there is a response

        """
        if not self.available:
            _LOGGER.debug(
                "ADB command not sent to %s:%d because adb-shell connection is not established: %s",
                self.host,
                self.port,
                cmd,
            )
            return None

        async with _acquire(self._adb_lock):
            _LOGGER.debug("Sending command to %s:%d via adb-shell: %s", self.host, self.port, cmd)
            return await self._adb.shell(cmd)


class ADBServerAsync(object):
    """A manager for ADB connections that uses an ADB server.

    Parameters
    ----------
    host : str
        The address of the device; may be an IP address or a host name
    port : int
        The device port to which we are connecting (default is 5555)
    adb_server_ip : str
        The IP address of the ADB server
    adb_server_port : int
        The port for the ADB server

    """

    def __init__(self, host, port=5555, adb_server_ip="", adb_server_port=5037):
        self.host = host
        self.port = int(port)
        self.adb_server_ip = adb_server_ip
        self.adb_server_port = adb_server_port
        self._adb_client = None
        self._adb_device = None

        # keep track of whether the ADB connection is intact
        self._available = False

        # use a lock to make sure that ADB commands don't overlap
        self._adb_lock = asyncio.Lock()

    @property
    def available(self):
        """Check whether the ADB connection is intact.

        Returns
        -------
        bool
            Whether or not the ADB connection is intact

        """
        if not self._adb_client or not self._adb_device:
            return False

        return self._available

    async def close(self):
        """Close the ADB server socket connection.

        Currently, this doesn't do anything except set ``self._available = False``.

        """
        self._available = False

    async def connect(self, log_errors=True):
        """Connect to an Android TV / Fire TV device.

        Parameters
        ----------
        log_errors : bool
            Whether errors should be logged

        Returns
        -------
        bool
            Whether or not the connection was successfully established and the device is available

        """
        try:
            async with _acquire(self._adb_lock):
                # Catch exceptions
                try:
                    self._adb_client = ClientAsync(host=self.adb_server_ip, port=self.adb_server_port)
                    self._adb_device = await self._adb_client.device("{}:{}".format(self.host, self.port))

                    # ADB connection successfully established
                    if self._adb_device:
                        _LOGGER.debug(
                            "ADB connection to %s:%d via ADB server %s:%d successfully established",
                            self.host,
                            self.port,
                            self.adb_server_ip,
                            self.adb_server_port,
                        )
                        self._available = True
                        return True

                    # ADB connection attempt failed (without an exception)
                    if log_errors:
                        _LOGGER.warning(
                            "Couldn't connect to %s:%d via ADB server %s:%d because the server is not connected to the device",
                            self.host,
                            self.port,
                            self.adb_server_ip,
                            self.adb_server_port,
                        )

                    await self.close()
                    self._available = False
                    return False

                # ADB connection attempt failed
                except Exception as exc:  # noqa pylint: disable=broad-except
                    if log_errors:
                        _LOGGER.warning(
                            "Couldn't connect to %s:%d via ADB server %s:%d, error: %s",
                            self.host,
                            self.port,
                            self.adb_server_ip,
                            self.adb_server_port,
                            exc,
                        )

                    await self.close()
                    self._available = False
                    return False

        except LockNotAcquiredException:
            _LOGGER.warning(
                "Couldn't connect to %s:%d via ADB server %s:%d because pure-python-adb lock not acquired.",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
            )
            await self.close()
            self._available = False
            return False

    async def pull(self, local_path, device_path):
        """Pull a file from the device using an ADB server.

        Parameters
        ----------
        local_path : str
            The path where the file will be saved
        device_path : str
            The file on the device that will be pulled

        """
        if not self.available:
            _LOGGER.debug(
                "ADB command not sent to %s:%d via ADB server %s:%d because pure-python-adb connection is not established: pull(%s, %s)",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
                local_path,
                device_path,
            )
            return

        async with _acquire(self._adb_lock):
            _LOGGER.debug(
                "Sending command to %s:%d via ADB server %s:%d: pull(%s, %s)",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
                local_path,
                device_path,
            )
            await self._adb_device.pull(device_path, local_path)
            return

    async def push(self, local_path, device_path):
        """Push a file to the device using an ADB server.

        Parameters
        ----------
        local_path : str
            The file that will be pushed to the device
        device_path : str
            The path where the file will be saved on the device

        """
        if not self.available:
            _LOGGER.debug(
                "ADB command not sent to %s:%d via ADB server %s:%d because pure-python-adb connection is not established: push(%s, %s)",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
                local_path,
                device_path,
            )
            return

        async with _acquire(self._adb_lock):
            _LOGGER.debug(
                "Sending command to %s:%d via ADB server %s:%d: push(%s, %s)",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
                local_path,
                device_path,
            )
            await self._adb_device.push(local_path, device_path)
            return

    async def screencap(self):
        """Take a screenshot using an ADB server.

        Returns
        -------
        bytes, None
            The screencap as a binary .png image, or ``None`` if there was an ``IndexError`` exception

        """
        if not self.available:
            _LOGGER.debug(
                "ADB screencap not taken from %s:%d via ADB server %s:%d because pure-python-adb connection is not established",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
            )
            return None

        async with _acquire(self._adb_lock):
            _LOGGER.debug(
                "Taking screencap from %s:%d via ADB server %s:%d",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
            )
            return await self._adb_device.screencap()

    async def shell(self, cmd):
        """Send an ADB command using an ADB server.

        Parameters
        ----------
        cmd : str
            The ADB command to be sent

        Returns
        -------
        str, None
            The response from the device, if there is a response

        """
        if not self.available:
            _LOGGER.debug(
                "ADB command not sent to %s:%d via ADB server %s:%d because pure-python-adb connection is not established: %s",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
                cmd,
            )
            return None

        async with _acquire(self._adb_lock):
            _LOGGER.debug(
                "Sending command to %s:%d via ADB server %s:%d: %s",
                self.host,
                self.port,
                self.adb_server_ip,
                self.adb_server_port,
                cmd,
            )
            return await self._adb_device.shell(cmd)