"""
Copy from https://github.com/doronz88/pymobiledevice3

Add http.client.HTTPConnection
"""
import abc
import plistlib
import socket
import sys
import time
from dataclasses import dataclass
from http.client import HTTPConnection
from typing import List, Mapping, Optional

from construct import Const, CString, Enum, FixedSized, GreedyBytes, Int16ul, Int32ul, Padding, Prefixed, StreamError, \
    Struct, Switch, this

from wdap.usbmux.exceptions import BadCommandError, BadDevError, MuxConnectError, \
    MuxConnectToUsbmuxdError, MuxError, MuxVersionError, NotPairedError

usbmuxd_version = Enum(Int32ul,
                       BINARY=0,
                       PLIST=1,
                       )

usbmuxd_result = Enum(Int32ul,
                      OK=0,
                      BADCOMMAND=1,
                      BADDEV=2,
                      CONNREFUSED=3,
                      BADVERSION=6,
                      )

usbmuxd_msgtype = Enum(Int32ul,
                       RESULT=1,
                       CONNECT=2,
                       LISTEN=3,
                       ADD=4,
                       REMOVE=5,
                       PAIRED=6,
                       PLIST=8,
                       )

usbmuxd_header = Struct(
    'version' / usbmuxd_version,  # protocol version
    'message' / usbmuxd_msgtype,  # message type
    'tag' / Int32ul,  # responses to this query will echo back this tag
)

usbmuxd_request = Prefixed(Int32ul, Struct(
    'header' / usbmuxd_header,
    'data' / Switch(this.header.message, {
        usbmuxd_msgtype.CONNECT: Struct(
            'device_id' / Int32ul,
            'port' / Int16ul,  # TCP port number
            'reserved' / Const(0, Int16ul),
        ),
        usbmuxd_msgtype.PLIST: GreedyBytes,
    }),
), includelength=True)

usbmuxd_device_record = Struct(
    'device_id' / Int32ul,
    'product_id' / Int16ul,
    'serial_number' / FixedSized(256, CString('ascii')),
    Padding(2),
    'location' / Int32ul
)

usbmuxd_response = Prefixed(Int32ul, Struct(
    'header' / usbmuxd_header,
    'data' / Switch(this.header.message, {
        usbmuxd_msgtype.RESULT: Struct(
            'result' / usbmuxd_result,
        ),
        usbmuxd_msgtype.ADD: usbmuxd_device_record,
        usbmuxd_msgtype.REMOVE: Struct(
            'device_id' / Int32ul,
        ),
        usbmuxd_msgtype.PLIST: GreedyBytes,
    }),
), includelength=True)


        

@dataclass
class MuxDevice:
    devid: int
    serial: str
    connection_type: str

    def connect(self, port: int, usbmux_address: Optional[str] = None) -> socket.socket:
        mux = create_mux(usbmux_address=usbmux_address)
        try:
            return mux.connect(self, port)
        except:  # noqa: E722
            mux.close()
            raise

    @property
    def is_usb(self) -> bool:
        return self.connection_type == 'USB'

    @property
    def is_network(self) -> bool:
        return self.connection_type == 'Network'

    def matches_udid(self, udid: str) -> bool:
        return self.serial.replace('-', '') == udid.replace('-', '')

    def make_http_connection(self, port: int) -> HTTPConnection:
        return USBMuxHTTPConnection(self, port)


class SafeStreamSocket:
    """ wrapper to native python socket object to be used with construct as a stream """

    def __init__(self, address, family):
        self._offset = 0
        self.sock = socket.socket(family, socket.SOCK_STREAM)
        self.sock.connect(address)

    def send(self, msg: bytes) -> int:
        self._offset += len(msg)
        self.sock.sendall(msg)
        return len(msg)

    def recv(self, size: int) -> bytes:
        msg = b''
        while len(msg) < size:
            chunk = self.sock.recv(size - len(msg))
            self._offset += len(chunk)
            if not chunk:
                raise MuxError('socket connection broken')
            msg += chunk
        return msg

    def close(self) -> None:
        self.sock.close()

    def settimeout(self, interval: float) -> None:
        self.sock.settimeout(interval)

    def setblocking(self, blocking: bool) -> None:
        self.sock.setblocking(blocking)

    def tell(self) -> int:
        return self._offset

    read = recv
    write = send


class MuxConnection:
    # used on Windows
    ITUNES_HOST = ('127.0.0.1', 27015)

    # used for macOS and Linux
    USBMUXD_PIPE = '/var/run/usbmuxd'

    @staticmethod
    def create_usbmux_socket(usbmux_address: Optional[str] = None) -> SafeStreamSocket:
        try:
            if usbmux_address is not None:
                if ':' in usbmux_address:
                    # assume tcp address
                    hostname, port = usbmux_address.split(':')
                    port = int(port)
                    address = (hostname, port)
                    family = socket.AF_INET
                else:
                    # assume unix domain address
                    address = usbmux_address
                    family = socket.AF_UNIX
            else:
                if sys.platform in ['win32', 'cygwin']:
                    address = MuxConnection.ITUNES_HOST
                    family = socket.AF_INET
                else:
                    address = MuxConnection.USBMUXD_PIPE
                    family = socket.AF_UNIX
            return SafeStreamSocket(address, family)
        except ConnectionRefusedError:
            raise MuxConnectToUsbmuxdError()

    @staticmethod
    def create(usbmux_address: Optional[str] = None):
        # first attempt to connect with possibly the wrong version header (plist protocol)
        sock = MuxConnection.create_usbmux_socket(usbmux_address=usbmux_address)

        message = usbmuxd_request.build({
            'header': {'version': usbmuxd_version.PLIST, 'message': usbmuxd_msgtype.PLIST, 'tag': 1},
            'data': plistlib.dumps({'MessageType': 'ReadBUID'})
        })
        sock.send(message)
        response = usbmuxd_response.parse_stream(sock)

        # if we sent a bad request, we should re-create the socket in the correct version this time
        sock.close()
        sock = MuxConnection.create_usbmux_socket(usbmux_address=usbmux_address)

        if response.header.version == usbmuxd_version.BINARY:
            return BinaryMuxConnection(sock)
        elif response.header.version == usbmuxd_version.PLIST:
            return PlistMuxConnection(sock)

        raise MuxVersionError(f'usbmuxd returned unsupported version: {response.version}')

    def __init__(self, sock: SafeStreamSocket):
        self._sock = sock

        # after initiating the "Connect" packet, this same socket will be used to transfer data into the service
        # residing inside the target device. when this happens, we can no longer send/receive control commands to
        # usbmux on same socket
        self._connected = False

        # message sequence number. used when verifying the response matched the request
        self._tag = 1

        self.devices = []

    @abc.abstractmethod
    def _connect(self, device_id: int, port: int):
        """ initiate a "Connect" request to target port """
        pass

    @abc.abstractmethod
    def get_device_list(self, timeout: float = None):
        """
        request an update to current device list
        """
        pass

    def connect(self, device: MuxDevice, port: int) -> socket.socket:
        """ connect to a relay port on target machine and get a raw python socket object for the connection """
        self._connect(device.devid, socket.htons(port))
        self._connected = True
        return self._sock.sock

    def close(self):
        """ close current socket """
        self._sock.close()

    def _assert_not_connected(self):
        """ verify active state is in state for control messages """
        if self._connected:
            raise MuxError('Mux is connected, cannot issue control packets')

    def _raise_mux_exception(self, result: int, message: str = None):
        exceptions = {
            int(usbmuxd_result.BADCOMMAND): BadCommandError,
            int(usbmuxd_result.BADDEV): BadDevError,
            int(usbmuxd_result.CONNREFUSED): MuxConnectError,
            int(usbmuxd_result.BADVERSION): MuxVersionError,
        }
        exception = exceptions.get(result, MuxError)
        raise exception(message)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


class BinaryMuxConnection(MuxConnection):
    """ old binary protocol """

    def __init__(self, sock: SafeStreamSocket):
        super().__init__(sock)
        self._version = usbmuxd_version.BINARY

    def get_device_list(self, timeout: float = None):
        """ use timeout to wait for the device list to be fully populated """
        self._assert_not_connected()
        end = time.time() + timeout
        self.listen()
        while time.time() < end:
            self._sock.settimeout(end - time.time())
            try:
                self._receive_device_state_update()
            except (BlockingIOError, StreamError):
                continue
            except IOError:
                try:
                    self._sock.setblocking(True)
                    self.close()
                except OSError:
                    pass
                raise MuxError('Exception in listener socket')

    def listen(self):
        """ start listening for events of attached and detached devices """
        self._send_receive(usbmuxd_msgtype.LISTEN)

    def _connect(self, device_id: int, port: int):
        self._send({'header': {'version': self._version,
                               'message': usbmuxd_msgtype.CONNECT,
                               'tag': self._tag},
                    'data': {'device_id': device_id, 'port': port},
                    })
        response = self._receive()
        if response.header.message != usbmuxd_msgtype.RESULT:
            raise MuxError(f'unexpected message type received: {response}')

        if response.data.result != usbmuxd_result.OK:
            raise self._raise_mux_exception(int(response.data.result),
                                            f'failed to connect to device: {device_id} at port: {port}. reason: '
                                            f'{response.data.result}')

    def _send(self, data: Mapping):
        self._assert_not_connected()
        self._sock.send(usbmuxd_request.build(data))
        self._tag += 1

    def _receive(self, expected_tag: int = None):
        self._assert_not_connected()
        response = usbmuxd_response.parse_stream(self._sock)
        if expected_tag and response.header.tag != expected_tag:
            raise MuxError(f'Reply tag mismatch: expected {expected_tag}, got {response.header.tag}')
        return response

    def _send_receive(self, message_type: int):
        self._send({'header': {'version': self._version, 'message': message_type, 'tag': self._tag},
                    'data': b''})
        response = self._receive(self._tag - 1)
        if response.header.message != usbmuxd_msgtype.RESULT:
            raise MuxError(f'unexpected message type received: {response}')

        result = response.data.result
        if result != usbmuxd_result.OK:
            raise self._raise_mux_exception(int(result), f'{message_type} failed: error {result}')

    def _add_device(self, device: MuxDevice):
        self.devices.append(device)

    def _remove_device(self, device_id: int):
        self.devices = [device for device in self.devices if device.devid != device_id]

    def _receive_device_state_update(self):
        response = self._receive()
        if response.header.message == usbmuxd_msgtype.ADD:
            # old protocol only supported USB devices
            self._add_device(MuxDevice(response.data.device_id, response.data.serial_number, 'USB'))
        elif response.header.message == usbmuxd_msgtype.REMOVE:
            self._remove_device(response.data.device_id)
        else:
            raise MuxError(f'Invalid packet type received: {response}')


class PlistMuxConnection(BinaryMuxConnection):
    def __init__(self, sock: SafeStreamSocket):
        super().__init__(sock)
        self._version = usbmuxd_version.PLIST

    def listen(self) -> None:
        self._send_receive({'MessageType': 'Listen'})

    def get_pair_record(self, serial: str) -> Mapping:
        # serials are saved inside usbmuxd without '-'
        self._send({'MessageType': 'ReadPairRecord', 'PairRecordID': serial})
        response = self._receive(self._tag - 1)
        pair_record = response.get('PairRecordData')
        if pair_record is None:
            raise NotPairedError('device should be paired first')
        return plistlib.loads(pair_record)

    def get_device_list(self, timeout: float = None) -> None:
        """ get device list synchronously without waiting the timeout """
        self.devices = []
        self._send({'MessageType': 'ListDevices'})
        for response in self._receive(self._tag - 1)['DeviceList']:
            if response['MessageType'] == 'Attached':
                super()._add_device(MuxDevice(response['DeviceID'], response['Properties']['SerialNumber'],
                                              response['Properties']['ConnectionType']))
            elif response['MessageType'] == 'Detached':
                super()._remove_device(response['DeviceID'])
            else:
                raise MuxError(f'Invalid packet type received: {response}')

    def get_buid(self) -> str:
        """ get SystemBUID """
        self._send({'MessageType': 'ReadBUID'})
        return self._receive(self._tag - 1)['BUID']

    def save_pair_record(self, serial: str, device_id: int, record_data: bytes):
        # serials are saved inside usbmuxd without '-'
        self._send_receive({'MessageType': 'SavePairRecord',
                            'PairRecordID': serial,
                            'PairRecordData': record_data,
                            'DeviceID': device_id})

    def _connect(self, device_id: int, port: int):
        self._send_receive({'MessageType': 'Connect', 'DeviceID': device_id, 'PortNumber': port})

    def _send(self, data: Mapping):
        request = {'ClientVersionString': 'qt4i-usbmuxd', 'ProgName': 'pymobiledevice3', 'kLibUSBMuxVersion': 3}
        request.update(data)
        super()._send({'header': {'version': self._version,
                                  'message': usbmuxd_msgtype.PLIST,
                                  'tag': self._tag},
                       'data': plistlib.dumps(request),
                       })

    def _receive(self, expected_tag: int = None) -> Mapping:
        response = super()._receive(expected_tag=expected_tag)
        if response.header.message != usbmuxd_msgtype.PLIST:
            raise MuxError(f'Received non-plist type {response}')
        return plistlib.loads(response.data)

    def _send_receive(self, data: Mapping):
        self._send(data)
        response = self._receive(self._tag - 1)
        if response['MessageType'] != 'Result':
            raise MuxError(f'got an invalid message: {response}')
        if response['Number'] != 0:
            raise self._raise_mux_exception(response['Number'], f'got an error message: {response}')


def create_mux(usbmux_address: Optional[str] = None) -> MuxConnection:
    return MuxConnection.create(usbmux_address=usbmux_address)


def list_devices(usbmux_address: Optional[str] = None) -> List[MuxDevice]:
    mux = create_mux(usbmux_address=usbmux_address)
    mux.get_device_list(0.1)
    devices = mux.devices
    mux.close()
    return devices


def select_device(udid: str = None, connection_type: str = None, usbmux_address: Optional[str] = None) \
        -> Optional[MuxDevice]:
    """
    select a UsbMux device according to given arguments.
    if more than one device could be selected, always prefer the usb one.
    """
    tmp = None
    for device in list_devices(usbmux_address=usbmux_address):
        if connection_type is not None and device.connection_type != connection_type:
            # if a specific connection_type was desired and not of this one then skip
            continue

        if udid is not None and not device.matches_udid(udid):
            # if a specific udid was desired and not of this one then skip
            continue

        # save best result as a temporary
        tmp = device

        if device.is_usb:
            # always prefer usb connection
            return device

    return tmp


def select_devices_by_connection_type(connection_type: str, usbmux_address: Optional[str] = None) -> List[MuxDevice]:
    """
    select all UsbMux devices by connection type
    """
    tmp = []
    for device in list_devices(usbmux_address=usbmux_address):
        if device.connection_type == connection_type:
            tmp.append(device)

    return tmp



class USBMuxHTTPConnection(HTTPConnection):
    def __init__(self, device: MuxDevice, port=8100):
        super().__init__("localhost", port)
        self.__device = device
        self.__port = port

    def connect(self):
        self.sock = self.__device.connect(self.__port)

    def __enter__(self) -> HTTPConnection:
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()