# ##### BEGIN GPL LICENSE BLOCK #####
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU General Public License as published by
#   the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU General Public License for more details.
#
#   You should have received a copy of the GNU General Public License
#   along with this program.  If not, see <https://www.gnu.org/licenses/>.
#
# ##### END GPL LICENSE BLOCK #####


import logging
import threading
import time
import zmq
import pickle
import copy
import uuid
import string
import random
import marshal

from .data import (ReplicatedCommand,
                   RepDeleteCommand,
                   RepAuthCommand,
                   RepSnapshotCommand,
                   ReplicatedDatablock,
                   RepRightCommand,
                   ReplicatedCommandFactory,
                   RepDisconnectCommand,
                   RepUpdateClientsState,
                   RepUpdateUserMetadata,
                   RepServerSnapshotCommand,
                   RepKickCommand)
from .graph import ReplicationGraph
from .constants import (STATE_ACTIVE, STATE_INITIAL,
                        STATE_SYNCING,
                        STATE_AUTH, STATE_WAITING, STATE_SRV_SYNC,
                        STATE_LOBBY, COMMITED, RP_COMMON,
                        RP_STRICT, RP_COMMON)
from .utils import current_milli_time
from .service import Service
from .exception import StateError
# from replication import __version__

AUTH_TIMEOUT = 1000


class ClientNetService(Service):
    def __init__(
            self,
            store_reference=None,
            factory=None,
            items_to_push=None,
            tasks=None,
            id=None,
            address='127.0.0.1',
            port=5560,
            ipc_port=5560,
            password=None
    ):
        Service.__init__(
            self,
            name="NetworkIO",
            ipc_port=ipc_port
        )

        self._tasks = tasks

        # Replication
        self._factory = factory
        self._store_reference = store_reference
        self._id = "None"
        self._items_to_push = items_to_push

        # Networking
        self._state = STATE_INITIAL
        self._srv_snapshot_size = 0
        self._srv_snapshot_progress = 0
        self._id = id
        tmp_uuid = uuid.uuid4()
        # uuid needed to avoid reconnexion problems on router sockets
        self._uuid = tmp_uuid.bytes
        self._uuid_str = str(tmp_uuid)

        logging.info(f"connecting on {address}:{port}")
        self._command = self.context.socket(zmq.DEALER)
        self._command.setsockopt(zmq.IDENTITY, self._uuid)
        self._command.connect(f"tcp://{address}:{port}")
        self._command.linger = 0

        self._subscriber = self.context.socket(zmq.DEALER)
        self._subscriber.setsockopt(zmq.IDENTITY, self._uuid)
        self._subscriber.connect(f"tcp://{address}:{port+1}")
        self._subscriber.linger = 0
        self._subscriber.setsockopt(zmq.RATE, 1000000)
        self._subscriber.setsockopt(zmq.RCVBUF, 2000000)

        self._publish = self.context.socket(zmq.PUSH)
        self._publish.setsockopt(zmq.IDENTITY, self._uuid)
        self._publish.connect(f"tcp://{address}:{port+2}")
        self._publish.linger = 0
        self._publish.setsockopt(zmq.RATE, 1000000)
        self._publish.setsockopt(zmq.SNDBUF, 2000000)

        self.poller.register(self._command, zmq.POLLIN)
        self.poller.register(self._subscriber, zmq.POLLIN)
        self._command_factory = ReplicatedCommandFactory()

        auth_type = 'ADMIN' if password else 'CLIENT'
        auth_request = RepAuthCommand(
            owner=self._id, data={
                "AUTH_TYPE": auth_type,
                "AUTH_ID": self._id,
                "PWD": password,
            })
        auth_request.push(self._command)

        self.state = STATE_AUTH
        self._connection_start_time = current_milli_time()

        # logging.info(__version__)
        self.start()

    def _assert_state(self, state):
        if self._state not in state:
            raise StateError("Client in wrong state")

    def _push_waiting_items(self):
        while not self._items_to_push.empty() and self._state != STATE_INITIAL:
            node = self._items_to_push.get(False)

            if isinstance(node, (RepDeleteCommand, RepKickCommand, RepRightCommand,
                                 RepDisconnectCommand, RepUpdateUserMetadata)):
                node.push(self._command)
            else:
                node.push(self._publish)

                if self.state == STATE_SRV_SYNC:
                    # START TO PUSH TO THE SERVER
                    self._srv_snapshot_progress += 1

                    self.notify_state(
                        current=self._srv_snapshot_progress,
                        total=self._srv_snapshot_size
                    )

                    if self._srv_snapshot_progress == self._srv_snapshot_size:
                        snapshot_end = RepServerSnapshotCommand(
                            owner='server',
                            data='SNAPSHOT_END')
                        snapshot_end.push(self._command)
                        self.state = STATE_ACTIVE
                        self.notify_event('CONNECTED')

    def notify_state(self, current=-1, total=-1):
        """ Notify the ochestrator about a change of state
        
        :param current: current state progression counter
        :type current: int
        :param total: current state total counter 
        :type total: int
        """
        state = {
            'STATE': self.state,
            'CURRENT': current,
            'TOTAL': total
        }
        self.notify(['STATE/CONNECTION', state])

    def notify_event(self, event_path):
        """ Notify the orchestrator of an event
        
        :param event_path: event identifer
        :type event_path: str
        """
        self.notify([f"EVENT/{event_path}"])

    def request_snapshot_init(self):
        """ Ask to the server for repository init. """

        logging.info('Requestion snapshots init')

        self._snapshot_progress = 0
        self._snapshot_total = 0

        snapshot_request = RepSnapshotCommand(
            owner=self._id,
            data={'STATE': "REQUEST_INIT"})

        snapshot_request.push(self._command)

    def request_server_repository_init(self):
        """ Ask to the server for repository init.

            :param command: incoming command
            :type command: RepAuthCommand
        """
        logging.info('Request server init')

        self._srv_snapshot_size = len(self._store_reference)
        keys = [k for k, v in self._store_reference.items()]
        assert(self._srv_snapshot_size > 0)

        snapshot_cmd = RepServerSnapshotCommand(
            owner='server',
            data={'STATE': 'INIT',
                  'SIZE': self._srv_snapshot_size,
                  'NODES': keys})
        snapshot_cmd.push(self._command)

        self.state = STATE_SRV_SYNC

    def handle_authentification(self, command):
        """ Manage client authentification 

            :param command: incoming command
            :type command: RepAuthCommand
        """
        self._assert_state([STATE_AUTH])

        connection_status = command.data

        if 'LOBBY' in connection_status:
            self.state = STATE_LOBBY
            self.notify_event('LOBBY')
        if 'RUNNING' in connection_status:
            self.state = STATE_LOBBY
            self.request_snapshot_init()
        if 'FAILURE' in connection_status:
            self.notify_event('CONNECTION_REFUSED')

    def handle_client_snapshot(self, command):
        """ Manage incoming snapshot commands

            :param command: incoming command
            :type command: RepSnapshotCommand
        """
        self._assert_state([STATE_SYNCING, STATE_LOBBY])

        snapshot_state = command.data['STATE']
        if snapshot_state == 'INIT':
            logging.info("client init")
            self._snapshot_progress = 0
            self._snapshot_catalog = command.data.get('CATALOG')
            self._snapshot_total = len(self._snapshot_catalog)

            self.notify_state(
                current=self._snapshot_progress,
                total=self._snapshot_total
            )

            self._current_snapshot = self._snapshot_catalog.pop()
            self.get_snapshot(self._current_snapshot)

            self.state = STATE_SYNCING

    def handle_server_repository_init(self, command):
        """ Manage server initialization commands

            :param command: incoming command
            :type command: RepServerSnapshotCommand
        """
        self._assert_state([STATE_SRV_SYNC])

        cli_snapshot_state = command.data.get('STATE')

        if cli_snapshot_state == 'ACCEPTED':
            for index, node in enumerate(self._store_reference.values()):
                node.commit()
                node._serialize()
                node.push(self._publish)
                self.notify_state(
                    current=index,
                    total=len(self._store_reference)
                )

            snapshot_cmd = RepServerSnapshotCommand(
                owner='server',
                data={'STATE': 'END'})
            snapshot_cmd.push(self._command)

        elif cli_snapshot_state == 'DONE':
            self.state = STATE_ACTIVE
            self.notify_event('CONNECTED')
        elif cli_snapshot_state == 'REJECTED':
            logging.error("client snapshot refused by the server.")
            self._state = STATE_LOBBY

    def get_snapshot(self, id):
        """ Ask a specific snapshot to the server
        
            :param id: uuid of the data
            :type id: str
        """
        logging.debug(f"get {id}")
        snapshot_request = RepSnapshotCommand(
            owner=self._id,
            data={
                'STATE': "GET",
                'ID': id})

        snapshot_request.push(self._command)

    def main(self, sockets):
        # DATA OUT
        if self.state in [STATE_ACTIVE, STATE_SRV_SYNC]:
            self._push_waiting_items()

        # COMMANDS I/O
        if self._command in sockets:
            command = ReplicatedCommand.fetch(
                self._command, self._command_factory)

            # AUTHENTIFICATION
            if isinstance(command, RepAuthCommand):
                self.handle_authentification(command)

            # DISCONNECT CONFIRMATION
            if isinstance(command, RepDisconnectCommand):
                self.notify(['EVENT/DISCONNECT', command.data])

            # CLIENTS INFO UPDATE
            if isinstance(command, RepUpdateClientsState):
                self.notify(['STATE/CLIENTS', command.data])

            # SERVER-->CLIENT SNAPSHOT
            if isinstance(command, RepSnapshotCommand):
                self.handle_client_snapshot(command)

            # CLIENT -> SERVER SNAPSHOT
            if isinstance(command, RepServerSnapshotCommand):
                self.handle_server_repository_init(command)

            # GRAPH OPERATION (DELETE, CHANGE_RIGHT)
            if type(command) in [RepDeleteCommand, RepRightCommand]:
                command.execute(self._store_reference)

        # DATA IN
        if self._subscriber in sockets:
            datablock = ReplicatedDatablock.fetch(
                self._subscriber, self._factory)
            datablock.store(self._store_reference)

            # SERVER-->CLIENT SNAPSHOT update
            if self.state == STATE_SYNCING:
                assert(datablock.uuid == self._current_snapshot)

                self._snapshot_progress += 1
                self.notify_state(
                    current=self._snapshot_progress,
                    total=self._snapshot_total
                )

                if not self._snapshot_catalog:
                    self._state = STATE_ACTIVE
                    self.notify_event('CONNECTED')
                else:
                    self._current_snapshot = self._snapshot_catalog.pop()
                    self.get_snapshot(self._current_snapshot)

        # Various timeout checks
        # auth
        if self.state == STATE_AUTH:
            if (current_milli_time()-self._connection_start_time) > AUTH_TIMEOUT:
                self.notify_event('CONNECTION_FAILED')

    def stop(self):
        # Exit
        self._command.close()
        self._subscriber.close()
        self._publish.close()

        self._state = STATE_INITIAL

    @property
    def state(self):
        return self._state

    @state.setter
    def state(self, state):
        self._state = state
        self.notify_state()


class ServerNetService(threading.Thread):
    def __init__(self):

        # Threading
        threading.Thread.__init__(self)
        self.name = "Server_Network"
        self._exit_event = threading.Event()

        # Networking
        self._graph = ReplicationGraph()
        self._context = zmq.Context.instance()
        self._command = None
        self._publisher = None
        self._pull = None
        self._state = 0
        self.clients = {}
        self._ttl = None

    def listen(
            self,
            port=5560,
            password='admin',
            attached=False,
            timeout=2000):
        self._password = password
        self._port = port
        self._attached = attached

        # Update request
        self._command = self._context.socket(zmq.ROUTER)
        self._command.setsockopt(zmq.IDENTITY, b'SERVER_COMMAND')
        self._command.bind(f"tcp://*:{port}")
        self._command.linger = 0

        # Update all clients
        self._publisher = self._context.socket(zmq.ROUTER)
        self._publisher.setsockopt(zmq.IDENTITY, b'SERVER_PUSH')
        self._publisher.bind(f"tcp://*:{port+1}")
        self._publisher.linger = 0
        self._publisher.setsockopt(zmq.RATE, 1000000)
        self._publisher.setsockopt(zmq.SNDBUF, 2000000)

        # Update collector
        self._pull = self._context.socket(zmq.PULL)
        self._pull.setsockopt(zmq.IDENTITY, b'SERVER_PULL')
        self._pull.bind(f"tcp://*:{port+2}")
        self._pull.linger = 0
        self._pull.setsockopt(zmq.RATE, 1000000)
        self._pull.setsockopt(zmq.RCVBUF, 2000000)

        # TTL communication
        self._ttl_pipe = self._context.socket(zmq.DEALER)
        self._ttl_pipe.bind("inproc://server_ttl")
        self._ttl_pipe.linger = 0

        self._ttl = ServerTTL(
            port=port+3,
            timeout=timeout
        )

        self._client_snapshot_size = 0
        self._client_snapshot_progress = 0

        self.start()

    def disconnect_client(self, client):
        if client in self.clients:
            leaving_client = self.clients[client]
            logging.info(f"Disconnecting {leaving_client['id']} from server")
            cleanup_commands = []
            for key, node in self._graph.items():
                if node.owner == leaving_client:
                    cleanup_commands.append(
                        RepRightCommand(
                            owner='server',
                            data={
                                'uuid': node.uuid,
                                'owner': RP_COMMON
                            }
                        ))
            for rr_cmd in cleanup_commands:
                for cli in self.clients.keys():
                    if cli != client and not self._command._closed:
                        self._command.send(cli, zmq.SNDMORE)
                        rr_cmd.push(self._command)
                        rr_cmd.execute(self._graph)

            del self.clients[client]
            
            if len(self.clients) == 0 and self._attached:
                self.stop()

    def _clients_update(self):
        user_dict = {}
        for user, user_data in self.clients.items():
            user_dict[user_data['id']] = user_data

        clients_states = RepUpdateClientsState(
            owner='server',
            data=user_dict
        )

        # Push it to every clients
        for cli_name in self.clients:
            self._command.send(cli_name, zmq.SNDMORE)
            clients_states.push(self._command)

    def _register_client(self, auth_type, uid, id, password):
        """ Register a client on the server

        return:
        FAILED
        LOBBY
        RUNNING
        """

        for cli in self.clients.values():
            if id == cli['id']:
                logging.debug("client logged in")
                return 'FAILED'

        if auth_type == 'ADMIN' and password != self._password:
            return 'FAILED'

        logging.info(f"{id} is online")

        self.clients[uid] = {
            'id': id,
            'admin': auth_type == 'ADMIN',
            'latency': 999,
            'status': STATE_WAITING,
            'metadata': {},
        }
        if self._state in [STATE_WAITING, STATE_SRV_SYNC]:
            return 'LOBBY'
        else:
            return 'RUNNING'

    def kick(self, user):
        """ kick the given user

        :arg user: username of the kicked client
        :type user: str
        """
        for k, v in self.clients.items():
            if v['id'] == user:
                disconnect = RepDisconnectCommand(
                    owner='server',
                    data={
                        'reason': 'kicked by admin',
                    }
                )
                self._command.send(k, zmq.SNDMORE)
                disconnect.push(self._command)

                self._ttl_pipe.send_multipart([b'STOP_WATCHING', k])
                self.disconnect_client(k)

                logging.warning(f"{user} kicked from the session.")
                break
        else:
            logging.error(f"Can't kick {user}, user not found.")

    def send_client_snapshot_init(self, client):
        catalog = [str(k) for k in self._graph.keys()]
        snapshot_state = RepSnapshotCommand(
            owner='server',
            data={
                'STATE': 'INIT',
                'CATALOG': catalog})
        logging.info("srv client init")
        self._command.send(client, zmq.SNDMORE)
        snapshot_state.push(self._command)

    def handle_client_snapshot(self, command):
        """ Handle client snapshot commands """

        snapshot_state = command.data.get('STATE')

        if snapshot_state == 'REQUEST_INIT':
            self.send_client_snapshot_init(command.sender)
        elif snapshot_state == 'GET':
            node = self._graph.get(command.data['ID'])
            node.push(self._publisher, identity=command.sender)

    def run(self):
        logging.info(f"Listening on {self._port}.")
        poller = zmq.Poller()

        poller.register(self._command, zmq.POLLIN)
        poller.register(self._pull, zmq.POLLIN)
        poller.register(self._ttl_pipe, zmq.POLLIN)

        command_factory = ReplicatedCommandFactory()

        self._state = STATE_WAITING

        while not self._exit_event.is_set():
            # Non blocking poller
            socks = dict(poller.poll(1000))

            # COMMAND HANDLING
            if self._command in socks:
                command = ReplicatedCommand.server_fetch(
                    self._command, command_factory)

                # CLIENT DISCONNECTION
                if isinstance(command, RepDisconnectCommand):
                    self._ttl_pipe.send_multipart(
                        [b'STOP_WATCHING', command.sender])
                    self.disconnect_client(command.sender)
                    confirmation = RepDisconnectCommand(
                        owner='server',
                        data={
                            'reason': 'ended by user',
                        }
                    )
                    self._command.send(command.sender, zmq.SNDMORE)
                    confirmation.push(self._command)

                # AUHTENTIFICATION
                if isinstance(command, RepAuthCommand):
                    auth_type = command.data.get('AUTH_TYPE')
                    auth_origin = command.sender
                    auth_id = command.data.get('AUTH_ID')
                    auth_pass = command.data.get('PWD', None)

                    auth_status = self._register_client(auth_type,
                                                        auth_origin,
                                                        auth_id,
                                                        auth_pass)

                    auth_response = RepAuthCommand(
                        owner="server",
                        data=auth_status)

                    self._command.send(command.sender, zmq.SNDMORE)
                    auth_response.push(self._command)

                # SERVER-> CLIENT SNAPSHOT
                if isinstance(command, RepSnapshotCommand):
                    self.handle_client_snapshot(command)

                # CLIENT-> SERVER SNAPSHOT
                if isinstance(command, RepServerSnapshotCommand):
                    cli_snapshot_state = command.data.get('STATE')
                    cli_snapshot_lenght = command.data.get('SIZE')
                    cli_snapshot_dict = command.data.get('NODES')

                    if cli_snapshot_state == 'INIT':
                        if self._state == STATE_SRV_SYNC:  # REJECT
                            snapshot_status = "REJECTED"
                        if self._state == STATE_WAITING:  # ACCEPT and CONFIGURE
                            snapshot_status = "ACCEPTED"
                            self._client_snapshot_size = cli_snapshot_lenght
                            self._client_snapshot_dict = cli_snapshot_dict

                        snapshot_cmd = RepServerSnapshotCommand(
                            owner='server',
                            data={'STATE': snapshot_status})
                        self._command.send(command.sender, zmq.SNDMORE)
                        snapshot_cmd.push(self._command)
                        self._state = STATE_SRV_SYNC
                    if cli_snapshot_state == 'END':
                        logging.info("Done")

                # CLIENT METADATA
                if isinstance(command, RepUpdateUserMetadata):
                    user = self.clients.get(command.sender)

                    if user:
                        user['metadata'].update(command.data)
                        self._clients_update()

                # KICK
                if isinstance(command, RepKickCommand):
                    self.kick(command.data['user'])

                # OTHERS
                if type(command) in [RepDeleteCommand, RepRightCommand]:
                    command.execute(self._graph)
                    for cli_name in self.clients:
                        if cli_name != command.sender:
                            self._command.send(cli_name, zmq.SNDMORE)
                            command.push(self._command)

            # TTL HANDLING
            if self._ttl_pipe in socks:
                notification = self._ttl_pipe.recv_multipart()

                if notification[0] == b'STATE':
                    clients_states = pickle.loads(notification[1])

                    # Prepare update
                    for id, state in clients_states.items():
                        cli = self.clients.get(id)
                        if cli:
                            cli['latency'] = state['latency']
                        else:
                            self._ttl_pipe.send_multipart(
                                [b'STOP_WATCHING', id])

                    self._clients_update()

                if notification[0] == b'LOST':
                    self.disconnect_client(notification[1])

            # DATA HANDLING
            if self._pull in socks:
                # Regular update  routing (Clients / Server / Clients)
                datablock = ReplicatedDatablock.fetch(self._pull)
                datablock.store(self._graph)

                if self._state == STATE_ACTIVE:
                    # Update all clients
                    for cli_name in self.clients.keys():
                        if cli_name != datablock.sender:
                            datablock.push(self._publisher, identity=cli_name)
                elif self._state == STATE_SRV_SYNC:
                    self._client_snapshot_progress += 1
                    logging.info(
                        f"Receiving snapshot {self._client_snapshot_progress}/{self._client_snapshot_size}")
                    self._client_snapshot_dict.remove(datablock.uuid)
                    if len(self._client_snapshot_dict) == 0:
                        # Launch snapshot for other waiting clients
                        snapshot_cmd = RepServerSnapshotCommand(
                            owner='server',
                            data={'STATE': 'DONE'})

                        self._command.send(command.sender, zmq.SNDMORE)
                        snapshot_cmd.push(self._command)

                        for client in self.clients.keys():
                            if client != command.sender:
                                self.send_client_snapshot_init(client)

                        self._state = STATE_ACTIVE

        
        while self._ttl._state != STATE_INITIAL:
            time.sleep(1)
        
        self._command.close()
        self._pull.close()
        self._publisher.close()
        self._ttl_pipe.close()


    def stop(self):
        self._ttl.stop()
        self._exit_event.set() 
        self._state = 0



class ServerTTL(threading.Thread):
    def __init__(
            self,
            context=zmq.Context(),
            port=5562,
            timeout=1000):
        # Threading
        threading.Thread.__init__(self)
        self.name = "server_TTL"
        self.daemon = False
        self._id = id
        self._exit_event = threading.Event()

        # Networking
        self._context = context = zmq.Context.instance()
        self._command = self._context.socket(zmq.ROUTER)
        self._command.bind(f"tcp://*:{port}")
        self._command.linger = 0
        self._pipe = self._context.socket(zmq.DEALER)
        self._pipe.connect("inproc://server_ttl")
        self._pipe.linger = 0

        self._timeout = timeout
        self._state = STATE_INITIAL
        self._clients_state = {}

        self.start()

    def run(self):
        self._state = STATE_ACTIVE
        poller = zmq.Poller()

        poller.register(self._command, zmq.POLLIN)
        poller.register(self._pipe, zmq.POLLIN)
        last_update_time = current_milli_time()
        while not self._exit_event.is_set():
            socks = dict(poller.poll(1))
            current_time = current_milli_time()

            if self._command in socks:
                identity, frame = self._command.recv_multipart(0)

                if frame == b'INIT':
                    self._clients_state[identity] = {}
                    self._clients_state[identity]['latency'] = 999
                    self._clients_state[identity]['last_update'] = current_time

                client = self._clients_state.get(identity)

                if client is None:
                    continue

                self._command.send(identity, zmq.SNDMORE)
                self._command.send(b"PING")

                # Update delay information
                client['latency'] = (current_time-(client['last_update']))-50
                client['last_update'] = current_time

            if self._pipe in socks:
                notification = self._pipe.recv_multipart()

                if notification[0] == b'STOP_WATCHING':
                    self.stop_monitor(notification[1])

            if current_time-last_update_time > 1000:
                last_update_time = current_time
                client_to_remove = []

                # Check clients status
                for client in self._clients_state.keys():
                    ping = (current_time -
                            (self._clients_state[client]['last_update']))

                    if ping > self._timeout:
                        logging.error(
                            f"Lost connexion to {client} ({ping} ms)")
                        client_to_remove.append(client)
                        self._pipe.send_multipart([b'LOST', client])

                for client in client_to_remove:
                    self.stop_monitor(client)

                self._pipe.send_multipart(
                    [b'STATE', pickle.dumps(self._clients_state)])
        self._command.close()
        self._pipe.close()
        self._state = STATE_INITIAL
        

    def stop_monitor(self, client):
        if client in self._clients_state:
            logging.debug(f"Removing client {client} from watchlist")
            del self._clients_state[client]

    @property
    def state(self):
        return self._state

    @property
    def clients_state(self):
        return copy.copy(self._clients_state)

    def stop(self):
        self._exit_event.set()
