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

# A bit of warning suppressing for the retry module
logging.getLogger("retry").setLevel(logging.ERROR)

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 `opengen`.
    """

    def __init__(self, optimizer_path=None, ip=None, port=None):
        """
        Constructs instance of `OptimizerTcpManager`

        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.


        :param optimizer_path: path to auto-generated optimizer (just to be clear: this is
            the folder that contains ``optimizer.yml``)
        :type optimizer_path: str

        :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)
        :type ip: str

        :param port: see ip
        :type port: int

        :return: New instance of :class:`~opengen.tcp.optimizer_tcp_manager.OptimizerTcpManager`
        """
        self.__optimizer_path = optimizer_path
        if optimizer_path is not None:
            # create attribute (including IP and port)
            self.__optimizer_details = None
            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 = version("opengen")
            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)
        try:
            s.connect((ip, port))
        except ConnectionRefusedError as err:
            s.close()
            raise err
        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']
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM, 0) as s:
            result = 0 == s.connect_ex((ip, port))

        return result

    @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

        :param p: vector of parameters
        :type p: list of `float`

        :param initial_guess: initial guess vector
        :type initial_guess: list of `float`

        :param initial_y: initial vector of Lagrange multipliers
        :type initial_y: list of `float`

        :param initial_penalty: initial penalty parameter
        :type initial_penalty: float

        :param buffer_len: buffer length used to read the server response,
             defaults to 4096
        :type buffer_len: int

        :param max_data_size: maximum data size that is expected to be received 
            from the TCP server, defaults to 1048576
        :type max_data_size: int

        :return: SolverResponse object
        :rtype: :class:`~opengen.tcp.solver_response.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))
