"""
Copyright (c) 2012 Maciej Wasilak <http://sixpinetrees.blogspot.com/>
              2017 Nordic Semiconductor ASA

CoAP protocol implementation.
"""
import datetime
import logging
import os
import random
import sys
from threading import Timer
from ipaddress import ip_address

from piccata.constants import *
from piccata.message import Message
from piccata.types import Endpoint


class _CoapMessageLayer(object):
    """Lower layer of the CoAP protocol.

    Responsible for message retransmisison, message deduplication and CON/ACK pairig.
    Valid requests/responses are forwarded to the transaction layer.
    """

    def __init__(self, transport):
        """ Initialize _CoapMessageLayer object.

        Args:
            transport (transport.TransportBase): A transport that shall be used by the message layer.
        """
        self._transport = transport
        self._transaction_layer = None

        self._message_id = random.randint(0, 65535)

        self._recent_local_ids = {}  # recently received messages with IDs generated locally (identified by message ID and remote)
        self._recent_remote_ids = {}  # recently received messages with IDs generated by remote endpoints (identified by message ID and remote)
        self._active_exchanges = {}  # active exchanges i.e. sent CON messages (identified by message ID and remote)

    def _deduplicate_message(self, message):
        """Check incoming message if it's a duplicate.

        Duplicate is a message with the same Message ID (mid) and sender (remote),
        as message received within last EXCHANGE_LIFETIME seconds (usually 247 seconds).

        Args:
            message (piccata.message.Message): A message to check.

        Returns:
            bool: The return value. True if duplicate was detected, False otherwise.
        """
        def _add_message_to_recent(cache, key):
            timeout = datetime.datetime.now() + datetime.timedelta(seconds=EXCHANGE_LIFETIME)
            cache[key] = (message, timeout)

        # Forget Message ID uses past their lifetime.
        now = datetime.datetime.now()
        for key, data in self._recent_local_ids.items():
            timeout = data[1]
            if timeout <= now:
                del self._recent_local_ids[key]
            else:
                # dict iteration order is guarateed to be in
                # insertion order. The rest of the messages are
                # all newer.
                break

        # Check for reused Message ID, remembering new messages
        # and issuing retransmissions.
        key = (message.mid, message.remote)
        logging.info("Incoming Message ID: %d" % message.mid)

        if message.mtype in (CON, NON):
            if key in self._recent_remote_ids:
                if message.mtype is CON:
                    if len(self._recent_remote_ids[key]) == 3:
                        logging.info('Duplicate CON received, sending old response again')
                        self.send_message(self._recent_remote_ids[key][2])
                    else:
                        logging.info('Duplicate CON received, no response to send')
                else:
                    logging.info('Duplicate NON received')
                return True
            else:
                logging.info('New unique CON or NON message received')
                _add_message_to_recent(self._recent_remote_ids, key)
                return False
        else:
            if key in self._recent_local_ids:
                logging.info('Duplicate ACK or RST received')
                return True
            else:
                logging.info('New unique ACK or RST message received')
                _add_message_to_recent(self._recent_local_ids, key)
                return False

    def _next_message_id(self):
        """Reserve and return a new message ID.

        Returns:
            A new message ID.
        """
        message_id = self._message_id
        self._message_id = (self._message_id + 1) & 0xFFFF
        return message_id

    def _enqueue_exchange(self, message, timeout, retransmission_counter):
        retransmission_timer = Timer(timeout, self._retransmit, (message.mid, timeout, retransmission_counter))
        retransmission_timer.daemon = True
        retransmission_timer.start()
        self._active_exchanges[message.mid] = (message, retransmission_timer)

    def _add_exchange(self, message):
        """Add an outgoing CON message to the retransmission list.

        CON (Confirmable) messages are automatically retransmitted by protocol until
        ACK or RST message with the same Message ID is received from target host.

        Args:
            message (piccata.message.Message): A message to retransmit.
        """
        timeout = random.uniform(ACK_TIMEOUT, ACK_TIMEOUT * ACK_RANDOM_FACTOR)
        retransmission_counter = 0
        self._enqueue_exchange(message, timeout, retransmission_counter)
        logging.info("Exchange added, Message ID: %d." % message.mid)

    def _remove_exchange(self, mid):
        """Remove a message from retranmission list and cancel the timer for next retransmission.

        Args:
           mid (int): An ID of a message to remove.

        Returns:
            piccata.message.Message: A message removed from the retransmission list. None if no message was found.
        """
        msg, timer = self._active_exchanges.pop(mid, (None, None))
        if timer != None:
            timer.cancel()
        logging.info("Exchange removed, Message ID: %d." % mid)
        return msg

    def _retransmit(self, mid, timeout, retransmission_counter):
        """Retransmit CON message that has not been ACKed or RSTed.

        Args:
            mid (int): An ID of a message to retransmit.
            timeout (int): A last timeout value.
            retransmission_counter (int): A number of times the message was retransmitted.
        """
        message, _ = self._active_exchanges.pop(mid, (None, None))
        if message != None:
            if retransmission_counter < MAX_RETRANSMIT:
                self._transport.send(message.encode(), message.remote)
                retransmission_counter += 1
                timeout *= 2
                self._enqueue_exchange(message, timeout, retransmission_counter)
                logging.info("Retransmission, Message ID: %d." % message.mid)
            else:
                pass
                #TODO: error handling (especially for requests)
        else:
            logging.error("Message no longer exists, Message ID: %d." % mid)

    def register_transaction_layer(self, transaction_layer):
        """Bind a CoAP Transaction layer with a CoAP Message layer.

        Args:
            transaction_layer (piccata.core._CoapTransactionLayer): A transaction layer for handling messages.
        """
        self._transaction_layer = transaction_layer

    def receive(self, data, remote, local):
        """Process raw messages received from transport layer.

        Transport layer will call this funciton on data reception.

        Args:
            data (bytes): Data received.
            remote (piccata.types.Endpoint): An address of the message originator.
            local (piccata.types.Endpoint): A destination address that data was received to.
        """
        logging.info("Received %r from %s:%d" % (data, remote[0], remote[1]))
        message = Message.decode(data, remote)
        if self._deduplicate_message(message):
            return

        if message.mtype in (ACK, RST):
            req = self._remove_exchange(message.mid)
            if message.mtype == RST:
                self._transaction_layer.reset_transaction(req)
                return

        self._transaction_layer.receive_message(message, remote, local)

    def send_message(self, message):
        """Set Message ID, encode and send message. Also if message is Confirmable (CON) add exchange.

        Args:
            message (piccata.message.Message): A message to send.
        """
        logging.info("Sending message to %s:%d" % message.remote)

        # Check if message is present on deduplication list and register response.
        if message.mtype in (ACK, RST):
            recent_key = (message.mid, message.remote)
            if recent_key in self._recent_remote_ids:
                if len(self._recent_remote_ids[recent_key]) != 3:
                    self._recent_remote_ids[recent_key] = self._recent_remote_ids[recent_key] + (message,)

        if message.mid is None:
            message.mid = self._next_message_id()

        raw_message = message.encode()
        self._transport.send(raw_message, message.remote)

        if message.mtype is CON:
            self._add_exchange(message)
        logging.info("Message %r sent successfully" % raw_message)

    def cancel_retransmission(self, mid):
        """Simply cancel further retansmissions.

        Args:
            mid (int): A message ID of a transaction to cancel.
        """
        self._remove_exchange(mid)


class _CoapTransactionLayer(object):
    """Higher layer of the CoAP protocol.

    Responsible for matching Request/Response matching.
    Valid requests are forwarded to registered request_handler.
    Valid responses are forwareded to a callback registered with a respective request.
    """

    def __init__(self, message_layer):
        """Initialize CoAP Transaction layer object.

        Args:
            message_layer (piccata.core._CoapMessageLayer): A _CoapMessageLayer object that shall
                                                        be bound to the transaction layer.
        """
        self._message_layer = message_layer
        self._request_handler = None

        self._outgoing_requests = {}  # unfinished outgoing requests (identified by token and remote)

    def _handle_app_callback(self, callback, result, request, response):
        """Call application callback registered with a request.

        Args:
            callback (function): A callback function registered by a user.
            result (int): A result code of the transaction.
            request (piccata.message.Message): A request that the callback was registered with.
            response (piccata.message.Message): A response received to the respective request.
        """
        cb, args, kw = callback
        args = args or ()
        kw = kw or {}
        cb(result, request, response, *args, **kw)

    def _add_transaction(self, request, callback):
        """Add an active transaction and start a timeout timer.

        Args:
            request (piccata.message.Message): A request that is part of the transaction.
            callback (function): A callback function registered by a user.
        """
        timer = Timer(request.timeout, self._timeout_transaction, (request, ))
        timer.daemon = True
        timer.start()
        self._outgoing_requests[(request.token, request.remote)] = (request, callback, timer)

    def _finish_transaction(self, token, remote, result, response):
        """Finalize the transaction by removing the transaction from list and calling respective callback.

        Args:
            token (bytes): A token related to transaction.
            remote (piccata.types.Endpoint): An endpoint address related to transaction.
            result (int): A result code of the transaction.
            response (piccata.message.Message): A response received to the respective request. May be None if no response was received.
        """
        try:
            request, callback, timer = self._outgoing_requests[(token, remote)]
        except KeyError:
            logging.info("Transaction not found.")
        else:
            if (not remote.addr.is_multicast):
                del self._outgoing_requests[(token, remote)]
                timer.cancel()
            self._handle_app_callback(callback, result, request, response)

    def _timeout_transaction(self, request):
        """Clean the Request after a timeout.

        Args:
            request (piccata.message.Message): A request that has timed out.
        """
        logging.info("Request timed out")
        # In case of transaction layer timeout, remove a possible retransmission on message layer as well
        self._message_layer.cancel_retransmission(request.mid)
        self._finish_transaction(request.token, request.remote, RESULT_TIMEOUT, None)

    def _process_request(self, request):
        """Method used for processing incoming requests.

        Args:
            request (piccata.message.Message): A request to process.
        """
        if request.mtype not in (CON, NON):
            return

        if self._request_handler is not None:
            response = self._request_handler.receive_request(request)
            if response is not None:
                self.send_response(request, response)
        else:
            # Send reset if we do not process requests.
            rst = Message.EmptyRstMessage(request)
            self._message_layer.send_message(rst)

    def _process_response(self, response):
        """Method used for processing incoming responses.

        Args:
            response (piccata.message.Message): A response to process.
        """

        def _reset_unrecognized():
            logging.info("Response not recognized - sending RST.")
            rst = Message.EmptyRstMessage(response)
            self._message_layer.send_message(rst)

        def _ack_if_confirmable():
            if response.mtype is CON:
                ack = Message.EmptyAckMessage(response)
                self._message_layer.send_message(ack)

        logging.info("Received Response, token: %s, host: %s, port: %s" % (response.token.hex(), response.remote[0], response.remote[1]))

        found = False
        for token, remote in self._outgoing_requests.keys():
            if (token == response.token) and (remote == response.remote or remote.addr.is_multicast):
                found = True
                self._finish_transaction(response.token, remote, RESULT_SUCCESS, response)
                _ack_if_confirmable()
                break

        if not found:
            _reset_unrecognized()

    def _process_empty(self, message):
        """Method used for processing empty messages.

        Args:
            message (piccata.message.Message): An empty message to process.
        """
        if message.mtype is CON:
            logging.info('Empty CON message received (CoAP Ping) - replying with RST.')
            rst = Message.EmptyRstMessage(request)
            self._message_layer.send_message(rst)
        # RST response is handled on the message layer (no token information in the RST message.)

    def reset_transaction(self, request):
        """Clean the transaction after reset from message processing layer.

        Args:
            request (piccata.message.Message): A request to reset.
        """
        logging.info("Request reseted from the remote")
        self._finish_transaction(request.token, request.remote, RESULT_RESET, None)

    def cancel_transaction(self, request):
        """Clean request after cancellation from user application.

        Args:
            request (piccata.message.Message): A request to cancel.
        """
        logging.info("Request cancelled")
        # In case application cancels transaction, remove a possible retransmission on message layer as well
        self._message_layer.cancel_retransmission(request.mid)
        self._finish_transaction(request.token, request.remote, RESULT_CANCELLED, None)

    def register_request_handler(self, request_handler):
        """Register an object that will process requests on application level.

        The object shall contatin receive_request function.

        Args:
            request_handler (object): A object for processing requests.
        """
        self._request_handler = request_handler

    def remove_request_handler(self, request_handler):
        """Remove an object for processing requests.

        Args:
            request_handler (object): A object for processing requests.
        """
        if (self._request_handler == request_handler):
            self._request_handler = None

    def receive_message(self, message, remote, local):
        """Process a received message. This function shall be called from the message layer.

        Args:
            message (piccata.message.Message): A received message.
            remote (piccata.types.Endpoint): A source address of the message.
            local (piccata.types.Endpoint): A destination address of the message.
        """
        if message.is_request():
            self._process_request(message)
        elif message.is_response():
            self._process_response(message)
        elif message.code is EMPTY:
            self._process_empty(message)
        else:
            logging.info("Invalid message received.")

    def send_request(self, request, response_callback, response_callback_args, response_callback_kw):
        """Send a request.

        Args:
            request (piccata.message.Message): A request to send.
            response_callback (function): A callback function that shall be called when response is received. May be None.
            response_callback_args (tuple): Optional arguments for the callback. May be None.
            response_callback_kw (dictionary): Optional keyword arguments for the callback. May be None.
        """
        if not request.is_request():
            raise ValueError("Message code is not valid for request")

        if request.mtype is None:
            raise ValueError("Message type not specified")

        assert response_callback == None or callable(response_callback)
        assert request.token is not None

        try:
            self._message_layer.send_message(request)
        except:
            raise
        else:
            if response_callback != None:
                callback = (response_callback, response_callback_args, response_callback_kw)
                self._add_transaction(request, callback)
            logging.info("Sending request - Token: %s, Host: %s, Port: %s" % (request.token.hex(), str(request.remote[0]), request.remote[1]))

    def send_response(self, request, response):
        """Send a response.

        Args:
            request (piccata.message.Message): A request that the response is related to.
            response (piccata.message.Message): A response message.
        """
        logging.info("Preparing response...")

        if not response.is_response():
            raise ValueError("Message code is not valid for a response.")

        response.token = request.token
        logging.info("Token: %s" % ":".join("{:02x}".format(c) for c in response.token))
        response.remote = request.remote

        if response.mtype is None:
            if request.mtype is CON:
                response.mtype = ACK
            else:
                response.mtype = NON

        if response.mid is None:
            if response.mtype in (ACK, RST):
                response.mid = request.mid

        logging.info("Sending response, type = %s (request type = %s)" % (types[response.mtype], types[request.mtype]))
        self._message_layer.send_message(response)

class Coap:
    """The CoAP protocol class.

    This class wraps together Message layer and Transaction layer.
    """

    def __init__(self, transport):
        """Initialize a CoAP protocol instance.

        Args:
            transport (transport.TransportBase): A transport object that Coap shall use for communicationm.
        """
        self._message_layer = _CoapMessageLayer(transport)
        self._transaction_layer = _CoapTransactionLayer(self._message_layer)
        self._message_layer.register_transaction_layer(self._transaction_layer)

    def register_request_handler(self, request_handler):
        """Register an object for handling requests. If no object is registered, a RST message will be sent back.

        The object shall contain receive_request function of the following format:
            receive_request(request)
        where request is an piccata.message.Message object.
        The function shall return a response message of type piccata.message.Message, or None if no response shall be sent.
        This implementation does not automatically ACK confirmable requests. So if response is going to be prepared later,
        but the reuqest shall be acknowledged now, the function shall return an empty ACK message. The postponed response
        can be sent with the respond function later.

        Args:
            request_handler (object): An object that will process requests.
        """
        self._transaction_layer.register_request_handler(request_handler)

    def remove_request_handler(self, request_handler):
        """Unregister an object for handling requests.

        Args:
            request_handler (object): An object that will process requests.
        """
        self._transaction_layer.remove_request_handler(request_handler)

    def receive(self, data, remote, local):
        """A function for receiving messages. Will be called by transport.

        This function passes received data to the message layer.

        Args:
            data (bytes): Data received.
            remote (piccata.types.Endpoint): An address of the message originator.
            local (piccata.types.Endpoint): A destination address that data was received to.
        """
        self._message_layer.receive(data, remote, local)

    def request(self, request, response_callback = None, response_callback_args = None, response_callback_kw = None):
        """Send a request and register a callback for the response.

        If no response is expected, response_callback shall be None.

        The callback function shall have the following format:
            callback(result, request, response, *args, **kwargs)
        where:
            result (int): A code informing of a transaction result.
            request (piccata.message.Message): A request that the callback was registered with.
            response (piccata.message.Message): A response message. May be None if no response was received.
            args (tuple): Opional arguments. May be empty.
            kwargs (dictionary): Optional keyword arguments. May be empty.

        The possible results of the transaction are:
            RESULT_SUCCESS: Response received successfully.
            RESULT_RESET = Request was reseted by the peer.
            RESULT_TIMEOUT = Request timed-out.
            RESULT_CANCELLED = Request was cancelled by the application.

        Args:
            request (piccata.message.Message): A request to be sent.
            response_callback (function): A callback funcition that will be called upon response reception. May be None.
            response_callback_args (tuple): An optional arguments for the callback function. May be None.
            response_callback_kw (dictionary): An optional keyword arguments for the callback function. May be None.
        """
        request.remote = Endpoint(ip_address(request.remote[0]), request.remote[1])
        return self._transaction_layer.send_request(request, response_callback, response_callback_args, response_callback_kw)
        # return Requester(self._transaction_layer, request, response_callback, response_callback_args, response_callback_kw)

    def cancel_request(self, request):
        """Cancel a pending request from the application.

        If a callback was registered and transaction still exists, it will be called.

        Args:
            request (piccata.message.Message): A request to cancel.
        """
        self._transaction_layer.cancel_transaction(request)

    def respond(self, request, response):
        """Send a separate response.

        This function can be used to send a delayed separate CON or NON response in case it could not be generated immediately
        in the request handler.
        Note, that this should not be used for sending ACK or RST responses. These should rather be returned by the request hadler.


        Args:
            request (piccata.message.Message): A request that the response refers to.
            response (piccata.message.Message): A response message.
        """
        request.remote = Endpoint(ip_address(request.remote[0]), request.remote[1])
        response.remote = Endpoint(ip_address(response.remote[0]), response.remote[1])
        self._transaction_layer.send_response(request, response)
