import yaml
import os
import subprocess
import socket
import json
import logging
import time
import math
import pkg_resources
from threading import Thread
from retry import retry
from .solver_response import SolverResponse


class OptimizerTcpManager:
    """Client for TCP interface of parametric optimizers

    This class is used to start and stop a TCP server, which
    has been generated by <code>opengen</code>.
    """

    def __init__(self, optimizer_path=None, ip=None, port=None):
        """
        Constructs instance of <code>OptimizerTcpManager</code>

        There are three ways to use this constructor:

        - OptimizerTcpManager(optimizer_path): creates a TCP manager for a local
          TCP server using the default IP and port of that TCP server (specified
          upon code generation)
        - OptimizerTcpManager(optimizer_path, ip, port): creates a TCP manager
          for a local TCP server, but overrides the default IP and port. This way
          the user can set the address '0.0.0.0', so that the TCP server binds on
          all IPs, or '127.0.0.1' so that it is accessible only locally, or a VPN
          IP address, so that the optimizer is accessible only over a private
          network.
        - OptimizerTcpManager(ip, port): If a path is not provided, then the
          TCP manager can be used to connect to a remote TCP server, as a client,
          but cannot be used to start the server.

        Args:
            :param optimizer_path:
            path to auto-generated optimizer (just to be clear: this is
            the folder that contains <code>optimizer.yml</code>)

            :param ip:
            the user can provide the IP of a remote TCP server (must be up and
            running) so as to establish a remote connection. In that case `path`
            must be equal to `None` (see examples above)

            :param port: see ip

        Returns:
            New instance of <code>OptimizerTcpManager</code>
        """
        self.__optimizer_path = optimizer_path
        if optimizer_path is not None:
            self.__optimizer_details = None  # create attribute (including IP and port)
            self.__load_tcp_details()
            if ip is not None:
                self.__optimizer_details['tcp']['ip'] = ip
            if port is not None:
                self.__optimizer_details['tcp']['port'] = port

            # Check whether the optimizer was built with the current version of opengen
            # We can only check the optimizer version if the optimizer runs locally
            opengen_version = self.__optimizer_details['build']['opengen_version']
            current_opengen_version = pkg_resources.require("opengen")[0].version
            if current_opengen_version != opengen_version:
                logging.warn(
                    'the target optimizer was build with a different version of opengen (%s)' % opengen_version)
                logging.warn('you are running opengen version %s' % current_opengen_version)

        elif ip is not None and port is not None:
            self.__optimizer_details = {"tcp": {"ip": ip, "port": port}}
        else:
            # If the optimizer path has not been provided, both the IP and and the
            # port must be provided, otherwise, raise an exception
            raise Exception("Illegal arguments")

        logging.info("TCP/IP details: %s:%d",
                     self.__optimizer_details['tcp']['ip'],
                     self.__optimizer_details['tcp']['port'])

    def __load_tcp_details(self):
        yaml_file = os.path.join(self.__optimizer_path, "optimizer.yml")
        with open(yaml_file, 'r') as stream:
            self.__optimizer_details = yaml.safe_load(stream)

    @retry(tries=10, delay=1)
    def __obtain_socket_connection(self):
        tcp_data = self.__optimizer_details
        ip = tcp_data['tcp']['ip']
        port = tcp_data['tcp']['port']
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        s.connect((ip, port))
        return s

    def __send_receive_data(self, text_to_send, buffer_size=512, max_data_size=1048576):
        conn_socket = self.__obtain_socket_connection()
        encoded_data = text_to_send.encode()
        conn_socket.sendall(encoded_data)
        conn_socket.shutdown(socket.SHUT_WR)

        max_read_rounds = math.ceil(max_data_size/buffer_size)
        data = b''
        for _i in range(max_read_rounds):
            data_chunk = conn_socket.recv(buffer_size)
            if data_chunk is None:
                break
            data += data_chunk

        conn_socket.close()
        return data.decode()

    def ping(self):
        """Pings the server

        Pings the server to check whether it is up and running
        """
        request = '{"Ping":1}'
        data = self.__send_receive_data(request)
        return json.loads(data)

    def __check_if_server_is_running(self):
        tcp_data = self.__optimizer_details
        ip = tcp_data['tcp']['ip']
        port = tcp_data['tcp']['port']
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0)
        return 0 == s.connect_ex((ip, port))

    @property
    def details(self):
        return self.__optimizer_details

    def start(self):
        """Starts the TCP server

        Note: this method starts a *local* server whose path must have been
        provided - we cannot start a remote server.

        The server starts on a separate thread, so this method does not block
        the execution of the caller's programme.

        """

        # Check if a path has been provided; if not,
        if self.__optimizer_path is None:
            raise Exception("No optimizer path provided - cannot start a remote server")

        # Server start data
        tcp_data = self.__optimizer_details
        ip = tcp_data['tcp']['ip']
        port = tcp_data['tcp']['port']

        # Check if any of the ip/port pairs is occupied
        if self.__check_if_server_is_running():
            msg = "Port %d not available" % port
            raise Exception(msg)

        def threaded_start():
            optimizer_details = self.__optimizer_details
            logging.info("Starting TCP/IP server at %s:%d (in a detached thread)",
                         ip, port)
            command = ['cargo', 'run', '-q']
            command += ["--release"] if optimizer_details['build']['build_mode'] == 'release' else []
            command += ['--', '--port=%d' % port, '--ip=%s' % ip]
            tcp_dir_name = "tcp_iface_" + optimizer_details['meta']['optimizer_name']
            tcp_iface_directory = os.path.join(self.__optimizer_path, tcp_dir_name)
            p = subprocess.Popen(command, cwd=tcp_iface_directory)
            p.wait()

        # start the server in a separate thread
        logging.info("Starting TCP/IP server thread")
        thread = Thread(target=threaded_start)
        thread.start()

        # ping the server until it responds so that we know it's
        # up and running
        logging.info("Waiting for server to start")
        time.sleep(0.1)
        self.ping()

    def kill(self):
        """Kills the server"""
        logging.info("Killing server")
        request = '{"Kill":1}'
        self.__send_receive_data(request)

    def call(self, p, initial_guess=None,
             initial_y=None,
             initial_penalty=None,
             buffer_len=4096,
             max_data_size=1048576) -> SolverResponse:
        """Calls the server

        Consumes the parametric optimizer by providing a parameter vector
        and, optionally, an initial guess

        Args:
             p: vector of parameters (list of float)
             initial_guess: initial guess vector (list of float)
             initial_y: initial vector of Lagrange multipliers (list of float)
             initial_penalty: initial penalty parameter (float)
             buffer_len: buffer length used to read the server response
             (default value: 4096)
             max_data_size: maximum data size that is expected to be
             received from the TCP server (default value: 1048576)

        Returns:
            Instance of SolverResponse

        """
        # Make request
        logging.debug("Sending request to TCP/IP server")
        run_message = '{"Run" : {"parameter": ['
        run_message += ','.join(map(str, p))
        run_message += ']'

        if initial_guess is not None:
            run_message += ', "initial_guess": ['
            run_message += ','.join(map(str, initial_guess))
            run_message += ']'

        if initial_y is not None:
            run_message += ', "initial_lagrange_multipliers": ['
            run_message += ','.join(map(str, initial_y))
            run_message += ']'

        if initial_penalty is not None:
            run_message += ', "initial_penalty": ' + str(float(initial_penalty))

        run_message += '}}'
        data = self.__send_receive_data(run_message, buffer_len, max_data_size)
        return SolverResponse(json.loads(data))
