import logging
import queue
import socket
import struct
import sys
from threading import (
    Lock,
    Thread,
)
from time import (
    sleep,
    time,
)
from typing import (
    Optional,
)

from ._MBException import MBException
from ._MBMessage import MBMessage
from ._MBService import MBService
from ._MBTimeout import MBTimeout
from ._crc import calc_crc
from ._format import (
    FM_PARAM,
    FM_STRING,
    FM_STR_ARRAY_PIPE,
    FM_UNDEFINED,
)
from ._status import (
    SE_MB_ERR,
    SE_MB_RCVBODY,
    SE_MB_RCVHDR,
    SE_MB_RCVTRAILER,
    SE_NOTFOUND,
    SE_SUCCESS,
    SE_TIMEOUT,
)
from ._token import (
    TK_CMP_GETDEPENDS,
    TK_CMP_GETSERVICES,
    TK_CMP_REGSERVICE,
    TK_CMP_START,
    TK_CMP_TERM_NOW,
    TK_EVT_EVENT,
    TK_EVT_SUBSCRIBE,
    TK_HV_FIND_SERVICE,
    TK_HV_START_COMP,
    TK_SYS_ACK,
    TK_SYS_NAK,
    TK_SYS_QUEUE_FULL,
    TK_EVT_UNSUBSCRIBE,
)

TCP_MSG_MAGIC_EXT = 0x21A2132F
TCP_MSG_TRAILER = 0x81010077
TCP_MSG_TRAILER_OW = 0x18989933

logger = logging.getLogger(__name__)


class MBContext(object):
    def __init__(
        self,
        my_service: Optional[MBService],
        name: str,
    ):
        self.my_service = my_service
        self.name = name

        if self.my_service is not None:
            self.msg_queue: queue.Queue[MBMessage] = queue.Queue(500)
            self.global_queue_lock = Lock()

            self.receive_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            linger_data = struct.pack('ii', 0, 0)
            self.receive_sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger_data)
            self.receive_sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            self.receive_sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self.receive_sock.bind((my_service.address, my_service.port))
            self.receive_sock.listen(100)

            self.running = True
            self.receive_thread = Thread(
                target=self._receiver,
                name="MBContextReceiverThread",
                daemon=True,
            )
            self.receive_thread.start()

    def MB_GetMessage(self, filtermin=0, filtermax=0, timeout:Optional[int]=None) -> MBMessage:
        if filtermin == 0 and filtermax == 0:
            msg = self.msg_queue.get(timeout=timeout)
            logger.debug("Returning msg with token [0x%08X]", msg.token)
            return msg

        while True:
            self.global_queue_lock.acquire(blocking=True)
            all_messages = []
            try:
                while True:
                    all_messages.append(self.msg_queue.get(block=False))
            except queue.Empty:
                pass

            if len(all_messages) == 0:
                self.global_queue_lock.release()
                all_messages.append(self.msg_queue.get(block=True))
                self.global_queue_lock.acquire(blocking=True)
                try:
                    all_messages.append(self.msg_queue.get(block=False))
                except queue.Empty:
                    pass

            ret = None
            for msg in all_messages:
                if ret is None and filtermin <= msg.token <= filtermax:
                    ret = msg
                else:
                    self.msg_queue.put(msg)

            self.global_queue_lock.release()

            if ret is not None:
                logger.debug("Returning msg with token [0x%08X]", ret.token)
                return ret

    def MB_SendFullMessage(
        self,
        destination: MBService,
        token: int,
        format: int = 0,
        data: bytes = None,
        headers: bytes = None,
        timeout: int = -1,
        ttl: int = -1,
    ):

        return self._send_message(
            destination=destination,
            token=token,
            format=format,
            data=data,
            headers=headers,
            timeout=timeout,
            ttl=ttl,
            one_way=False,
        )

    def MB_SendMessage(
        self,
        destination: MBService,
        token: int,
        format: int = 0,
        data: Optional[bytes] = None,
        timeout: int = -1,
    ) -> MBMessage:

        return self.MB_SendFullMessage(
            destination=destination,
            token=token,
            format=format,
            data=data,
            timeout=timeout,
        )

    def MB_SendFullOneWayMessage(
        self,
        destination: MBService,
        token: int,
        format: int = 0,
        data: bytes = None,
        headers: bytes = None,
        timeout: int = -1,
        ttl: int = -1,
    ):
        self._send_message(
            destination=destination,
            token=token,
            format=format,
            data=data,
            headers=headers,
            timeout=timeout,
            ttl=ttl,
            one_way=True,
        )

    def MB_SendOneWayMessage(
        self,
        destination: MBService,
        token: int,
        format: int = 0,
        data: bytes = None,
        timeout: int = -1,
    ):
        self._send_message(
            destination=destination,
            token=token,
            format=format,
            data=data,
            headers=None,
            timeout=timeout,
            ttl=-1,
            one_way=True,
        )

    def MB_ReplyMessage(
        self,
        message: MBMessage,
        format: int = FM_UNDEFINED,
        data: bytes = b"",
    ) -> None:

        if message.trailer == TCP_MSG_TRAILER_OW:
            logger.error("MBMessage is one-way, must not reply")
            return

        sock = getattr(message, "_socket", None)
        if sock is None:
            logger.error("MBMessage without socket, already replied")
            return

        try:
            self._send_data(
                sock=sock,
                dest_service_name=None,
                token=message.token,
                format=format,
                data=data,
                headers=b"",
                ttl=-1,
                one_way=False,
                end_time=int(time() + 5),
            )
            sock.close()
        except (Exception,):
            logger.exception("Error sending reply")

        delattr(message, "_socket")

    def MB_EvtSubscribe(
        self,
        hv_service: MBService,
        subject: str,
    ) -> None:

        reply = self.MB_SendMessage(
            destination=hv_service,
            token=TK_EVT_SUBSCRIBE,
            format=FM_PARAM,
            data=f"{subject}\0{self.my_service.address}\0{self.my_service.port}".encode(),
        )

        if reply.token != TK_SYS_ACK:
            raise MBException("Error calling MB_EvtSubscribe.", SE_MB_ERR)

    def MB_EvtUnsubscribe(
        self,
        hv_service: MBService,
        subject: str,
    ) -> None:

        reply = self.MB_SendMessage(
            destination=hv_service,
            token=TK_EVT_UNSUBSCRIBE,
            format=FM_PARAM,
            data=f"{subject}\0{self.my_service.address}\0{self.my_service.port}".encode(),
        )

        if reply.token != TK_SYS_ACK:
            raise MBException("Error calling MB_EvtSubscribe.", SE_MB_ERR)

    def MB_EvtSend(
        self,
        hv_service: MBService,
        subject: str,
        type: str,
        xml: str,
        synchronous: bool = False,
        sourceid: int = 0,
        queue: str = None,
        timeout: int = -1,
    ):
        return self.MB_EvtSendFull(
            hv_service=hv_service,
            headers=None,
            subject=subject,
            type=type,
            xml=xml,
            synchronous=synchronous,
            sourceid=sourceid,
            queue=queue,
            timeout=timeout,
            ttl=-1,
        )

    def MB_EvtSendFull(
        self,
        hv_service: MBService,
        headers: Optional[bytes],
        subject: str,
        type: str,
        xml: str,
        synchronous: bool = False,
        sourceid: int = 0,
        queue: str = None,
        timeout: int = -1,
        ttl: int = -1,
    ):
        sync = "true" if synchronous else "false"
        timeout_in_seconds = timeout / 1000000
        evt_data = f"{xml}\0{subject}\0{type}\0{sync}\0{sourceid}\0{timeout_in_seconds}"
        if queue is not None:
            evt_data += f"\0{queue}"
        if not synchronous:
            self.MB_SendFullOneWayMessage(
                destination=hv_service,
                token=TK_EVT_EVENT,
                format=FM_PARAM,
                data=evt_data.encode(),
                headers=headers,
                timeout=timeout,
                ttl=ttl,
            )
            return MBMessage(token=TK_SYS_ACK)
        else:
            return self.MB_SendMessage(
                destination=hv_service,
                token=TK_EVT_EVENT,
                format=FM_PARAM,
                data=evt_data.encode(),
                timeout=timeout,
            )

    def MB_LocateService(
        self,
        hv_service: MBService,
        service_name: str,
        maxretries: int = 5,
    ):
        reply = self.MB_SendFullMessage(
            destination=hv_service,
            token=TK_HV_FIND_SERVICE,
            format=FM_STRING,
            data=service_name.encode(),
        )

        if reply.token == TK_HV_START_COMP and maxretries > 0:
            sleep(0.5)
            reply = self.MB_LocateService(
                hv_service=hv_service,
                service_name=service_name,
                maxretries=maxretries - 1,
            )

        if reply.token != TK_SYS_ACK:
            raise MBException("Error calling MB_LocateService.", SE_NOTFOUND)

        parts = reply.data.decode().split("\0")
        return MBService(
            address=parts[0],
            port=int(parts[1]),
        )

    def MB_RegisterService(
        self,
        my_service: MBService,
        hv_service: MBService,
        service_type: str,
        service_name: str,
    ) -> None:

        reply = self.MB_SendFullMessage(
            destination=hv_service,
            token=TK_CMP_REGSERVICE,
            format=FM_PARAM,
            data=f"{service_type}\0{service_name}\0{my_service.address}\0{my_service.port}".encode(),
        )

        if reply.token != TK_SYS_ACK:
            raise MBException("Error calling MB_RegisterService.", SE_MB_ERR)

    def MB_WaitStart(
        self,
        exported_services: Optional[str] = None,
        required_services: Optional[str] = None,
    ) -> int:
        while True:
            msg = self.MB_GetMessage(filtermin=TK_CMP_TERM_NOW, filtermax=TK_CMP_START)
            if msg.token == TK_CMP_TERM_NOW:
                msg.token = TK_SYS_ACK
                self.MB_ReplyMessage(
                    message=msg,
                    format=FM_UNDEFINED,
                    data=b"",
                )
                raise MBException("Component is terminating", SE_MB_ERR)

            elif msg.token == TK_CMP_GETSERVICES:
                if exported_services is not None and exported_services != "":
                    msg.token = TK_SYS_ACK
                    self.MB_ReplyMessage(
                        message=msg,
                        format=FM_STR_ARRAY_PIPE,
                        data=exported_services.encode(),
                    )
                else:
                    msg.token = TK_SYS_NAK
                    self.MB_ReplyMessage(
                        message=msg,
                        format=FM_UNDEFINED,
                        data=b"",
                    )

            elif msg.token == TK_CMP_GETDEPENDS:
                if required_services is not None and required_services != "":
                    msg.token = TK_SYS_ACK
                    self.MB_ReplyMessage(
                        message=msg,
                        format=FM_STR_ARRAY_PIPE,
                        data=required_services.encode(),
                    )
                else:
                    msg.token = TK_SYS_NAK
                    self.MB_ReplyMessage(
                        message=msg,
                        format=FM_UNDEFINED,
                        data=b"",
                    )

            elif msg.token == TK_CMP_START:
                msg.token = TK_SYS_ACK
                self.MB_ReplyMessage(
                    message=msg,
                    format=FM_UNDEFINED,
                    data=b"",
                )
                return SE_SUCCESS

    def MB_Finalize(self):
        if self.my_service is None:
            return

        self.running = False
        if sys.platform != "win32":
            self.receive_sock.shutdown(socket.SHUT_RDWR)

        self.receive_sock.close()
        self.receive_thread.join()

    def _receive_message(
        self,
        sock: socket.socket,
        from_service: MBService,
        end_time: int,
    ) -> MBMessage:

        header_size = 32
        header = receive(sock, header_size, end_time)
        if len(header) != header_size:
            raise MBException("Error receiving header.", SE_MB_RCVHDR)

        magic = int.from_bytes(header[0:4], byteorder="little")
        if magic != TCP_MSG_MAGIC_EXT:
            raise MBException("Received magic is wrong", SE_MB_RCVHDR)

        data_size = int.from_bytes(header[4:8], byteorder="little")
        token = int.from_bytes(header[8:12], byteorder="little")
        format = int.from_bytes(header[12:16], byteorder="little")
        headers_size = int.from_bytes(header[16:20], byteorder="little")
        ttl = int.from_bytes(header[24:32], byteorder="little", signed=True)
        align_data = header[20:24]
        if int.from_bytes(align_data, byteorder="little") != 0x0:
            raise MBException("Invalid alignment data.", SE_MB_RCVHDR)

        dest_service_size = receive(sock, 4, end_time)
        if len(dest_service_size) != 4:
            raise MBException("Error receiving destination service size.", SE_MB_RCVHDR)

        dest_service_size = int.from_bytes(dest_service_size, byteorder="little")
        dest_service_name = None
        if dest_service_size > 0:
            dest_service_name = receive(sock, dest_service_size, end_time)
            if len(dest_service_name) != dest_service_size:
                raise MBException("Error receiving destination service size.", SE_MB_RCVHDR)

            dest_service_name = dest_service_name.decode()

        encrypted_data = b""
        if data_size > 0:
            encrypted_data = receive(sock, data_size, end_time)
            if len(encrypted_data) != data_size:
                raise MBException("Error receiving destination service size.", SE_MB_RCVBODY)

        received_data = encrypt(data=encrypted_data)

        headers = b""
        if headers_size > 0:
            headers = receive(sock, headers_size, end_time)
            if len(headers) != headers_size:
                raise MBException("Error receiving destination service size.", SE_MB_RCVBODY)

        trailer_bytes = receive(sock, 4, end_time)
        if len(trailer_bytes) != 4:
            raise MBException("Error receiving message trailer.", SE_MB_RCVTRAILER)
        trailer = int.from_bytes(trailer_bytes, byteorder="little")
        if trailer not in (TCP_MSG_TRAILER, TCP_MSG_TRAILER_OW):
            raise MBException("Invalid message trailer.", SE_MB_RCVTRAILER)

        crc_bytes = receive(sock, 4, end_time)
        if len(crc_bytes) != 4:
            raise MBException("Error receiving message trailer.", SE_MB_RCVTRAILER)

        received_crc = int.from_bytes(crc_bytes, byteorder="little")
        crc_data = TCP_MSG_MAGIC_EXT.to_bytes(4, byteorder="little")
        crc_data += data_size.to_bytes(4, byteorder="little")
        crc_data += token.to_bytes(4, byteorder="little")
        crc_data += format.to_bytes(4, byteorder="little")
        calculated_crc = calc_crc(crc_data, 0xffffffff)
        if data_size > 0:
            calculated_crc = calc_crc(received_data, calculated_crc)

        if received_crc != calculated_crc:
            raise MBException("Message CRC does not match.", SE_MB_RCVTRAILER)

        return MBMessage(
            token=token,
            format=format,
            ttl=ttl,
            destname=dest_service_name,
            data=received_data,
            headers=headers,
            trailer=trailer,
            remotehost=from_service.address,
        )


    def _send_message(
        self,
        destination: MBService,
        token: int,
        format: int = 0,
        data: bytes = None,
        headers: bytes = None,
        timeout: int = -1,
        ttl: int = -1,
        one_way: bool = False,
    ) -> Optional[MBMessage]:

        logger.debug("Sending token [0x%08X] to [%s]", token, destination)
        if data is None:
            data = b""

        if headers is None:
            headers = b""

        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        socket.gethostbyaddr(destination.address)
        try:
            linger_data = struct.pack('ii', 0, 0)
            sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger_data)
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            if timeout > 0:
                sock.settimeout(timeout / 1000000)
            else:
                sock.settimeout(5)

            start = time()
            if timeout > 0:
                end_time = start + timeout / 1000000
            else:
                end_time = start + 2 ** 16

            sock.connect((destination.address, destination.port))
            logger.debug("Successfully connected")

            now = time()
            remaining = end_time - now
            if remaining <= 0:
                raise MBTimeout("Timed-out on MB_Connect.", SE_TIMEOUT)

            self._send_data(
                sock=sock,
                dest_service_name=destination.name,
                token=token,
                format=format,
                data=data,
                headers=headers,
                ttl=ttl,
                one_way=one_way,
                end_time=end_time,
            )
            logger.debug("Data sent")

            reply = self._receive_message(
                sock=sock,
                from_service=destination,
                end_time=end_time,
            )
            logger.debug("Reply received with token [0x%08X]", reply.token)
            return reply
        except MBTimeout:
            raise
        except socket.timeout:
            raise MBTimeout("Timed-out on MB_Connect.", SE_TIMEOUT)
        except Exception as ex:
            raise MBException("Unknown error sending message.", SE_MB_ERR) from ex
        finally:
            sock.close()

    def _send_data(
        self,
        sock: socket.socket,
        token: int,
        format: int,
        data: bytes,
        headers: bytes,
        ttl: int,
        one_way: bool,
        dest_service_name: Optional[str],
        end_time: int,
    ):
        data_to_send = self._build_data_to_send(
            token=token,
            format=format,
            data=data,
            headers=headers,
            ttl=ttl,
            one_way=one_way,
            dest_service_name=dest_service_name,
        )

        now = time()
        remaining = end_time - now
        sock.settimeout(remaining)
        sock.sendall(data_to_send)

    def _build_data_to_send(
        self,
        token: int,
        format: int,
        data: bytes,
        headers: bytes,
        ttl: int,
        one_way: bool,
        dest_service_name: Optional[str],
    ):
        header_data = bytearray()
        header_data += TCP_MSG_MAGIC_EXT.to_bytes(4, byteorder="little")
        header_data += len(data).to_bytes(4, byteorder="little")
        header_data += token.to_bytes(4, byteorder="little")
        header_data += format.to_bytes(4, byteorder="little")

        crc = calc_crc(header_data, 0xffffffff)
        if len(data) > 0:
            crc = calc_crc(data, crc)

        send_data = header_data
        send_data += len(headers).to_bytes(4, byteorder="little")
        send_data += 0x0.to_bytes(4, byteorder="little")
        send_data += ttl.to_bytes(8, byteorder="little", signed=True)
        service_name_bytes = b""
        if dest_service_name is not None:
            service_name_bytes = dest_service_name.encode()
        send_data += len(service_name_bytes).to_bytes(4, byteorder="little")
        send_data += service_name_bytes
        send_data += encrypt(data=data)
        send_data += headers
        send_data += (TCP_MSG_TRAILER_OW if one_way else TCP_MSG_TRAILER).to_bytes(4, byteorder="little")
        send_data += crc.to_bytes(4, byteorder="little")

        return send_data

    def _receiver(self):
        while self.running:
            try:
                sock, address = self.receive_sock.accept()
                logger.debug("Accepted connection from %s:%d", address[0], address[1])
                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
                linger_data = struct.pack('ii', 0, 0)
                sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, linger_data)
                from_service = MBService(address[0], address[1])
                msg = self._receive_message(
                    sock=sock,
                    from_service=from_service,
                    end_time=int(time() + 5),
                )
                logger.debug("Received token [0x%08X]", msg.token)
                if msg.trailer == TCP_MSG_TRAILER_OW:
                    self._send_data(
                        sock=sock,
                        dest_service_name=None,
                        token=msg.token,
                        format=FM_UNDEFINED,
                        data=b"",
                        headers=b"",
                        ttl=-1,
                        one_way=False,
                        end_time=int(time() + 5),
                    )
                    sock.close()
                    pass
                else:
                    setattr(msg, "_socket", sock)

                try:
                    with self.global_queue_lock:
                        self.msg_queue.put(
                            item=msg,
                            block=False,
                        )
                except queue.Full:
                    logger.exception("Queue is full")
                    try:
                        self._send_data(
                            sock=sock,
                            dest_service_name=from_service.name,
                            token=TK_SYS_QUEUE_FULL,
                            format=FM_UNDEFINED,
                            data=b"",
                            headers=b"",
                            ttl=-1,
                            one_way=False,
                            end_time=int(time() + 5),
                        )
                        sock.close()
                    except (Exception,):
                        logger.exception("Error sending TK_SYS_QUEUE_FULL")

            except (Exception,):
                if self.receive_sock.fileno() == -1 and self.running:
                    logger.error("Socket closed but we are still running")
                elif self.receive_sock.fileno() != -1:
                    logger.exception("Error on receiver")


def encrypt(data: bytes):
    encrypted_data = bytearray()
    for b in data:
        encrypted_data += (b ^ 0x9B).to_bytes(1, byteorder="little")

    return encrypted_data

def receive(sock: socket.socket, size:int, end_time: int):
    remaining_size = size
    total_received = b""

    while True:
        now = time()
        remaining_time = int(end_time - now)
        sock.settimeout(remaining_time)
        try:
            received = sock.recv(remaining_size)
            if received == b"":
                return total_received

            total_received += received
            remaining_size -= len(received)
            if remaining_size <= 0:
                return total_received
        except socket.timeout:
            raise MBTimeout("Timed-out on MB_Connect.", SE_TIMEOUT)
