# ##### 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 copy
import gzip
import logging
import os
import pickle
import queue
import subprocess
import sys
import threading
from datetime import datetime
from pathlib import Path
import traceback
try:
    import _pickle as pickle
except ImportError:
    import pickle

import os
import queue
import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path

import zmq

from replication import __version__

from .constants import (ADDED, CLIENT_PING_FREQUENCY, COMMITED,
                        CONNECTION_TIMEOUT, FETCHED, MODIFIED, RP_COMMON,
                        SERVER_SCRIPT_PATH, STATE_ACTIVE, STATE_AUTH,
                        STATE_INITIAL, CONNECTING, STATE_LOBBY,
                        STATE_SRV_SYNC, STATE_SYNCING, STATE_WAITING, UP,
                        TTL_SCRIPT_PATH, STATE_QUITTING)
from .exception import (DataError, NetworkFrameError,
                        NonAuthorizedOperationError, StateError,
                        UnsupportedTypeError)
from .repository import Repository
from .objects import (Node, ReplicationObject, Auth, Delete, Disconnect,
                      Kick, Right, ServerSnapshot, Snapshot,
                      UpdateClientsState, UpdateUserMetadata)
from .utils import (current_milli_time, get_state_str)
from . import porcelain

this = sys.modules[__name__]


class Session(object):
    def __init__(self):
        self._stash = []
        self.repository = None
        self.callbacks = {}
        self._lock_operations = threading.Event()
        self._state = STATE_INITIAL
        self._state_progress = {
            'current': -1,
            'total': -1
        }
        self._server = None

        # Networking
        self._state = STATE_INITIAL
        self._srv_snapshot_size = 0
        self._srv_snapshot_progress = 0

    def _assert_modification_rights(self, node=None):
        if self.repository.graph.get(node).owner not in [self._id, RP_COMMON]:
            raise NonAuthorizedOperationError(
                "Not authorized to modify the node {node}")

    def register(self, name):
        def func_wrapper(func):
            self.callbacks[name] = func
            return func
        return func_wrapper

    def call_registered(self, name=None, **kwargs):
        func = self.callbacks.get(name, None)
        if func is None:
            logging.info("No function registered against - " + str(name))
            return None
        return func(**kwargs)

    def connect(self,
                repository: Repository = None,
                remote: str = "origin",
                timeout=1000,
                password=None):
        """Connect to a session

        :param id: user name
        :type id: string
        :param address: host ip address
        :type address: string
        :param port: host port
        :type port: int
        """
        self.repository = repository
        self.remote = repository.remotes.get(remote)
        self._data_protocol = repository.rdp

        self._id = repository.username
         # uuid needed to avoid reconnexion problems on router sockets
        self._connection_timeout = timeout
        self.context = zmq.Context()
        self._state = CONNECTING

        auth_type = 'ADMIN' if password else 'CLIENT'
        auth_request = Auth(
            owner=self._id, data={
                "AUTH_TYPE": auth_type,
                "AUTH_ID": self._id,
                "PWD": password,
                "VERSION": __version__,
            })
        self.repository.push_cmd(self.remote.command, auth_request)

        self._state = STATE_AUTH
        self._connection_start_time = current_milli_time()
        self._ttl = subprocess.Popen([
                        sys.executable,
                        TTL_SCRIPT_PATH,
                        '-p', str(self.remote.port),
                        '-d', self.remote.address,
                        '-i', str(self.remote.uuid),
                        '-t', str(timeout)]
                    )

    def host(self,
             id="Default",
             repository: Repository = None,
             remote='origin',
             timeout=5000,
             password=None,
             cache_directory='',
             server_log_level='INFO'):
        """Host a session

        :param id: user name
        :type id: string
        :param address: host ip address
        :type address: strings
        :param port: host port
        :type port: int
        """
        active_remote = repository.remotes.get(remote)
        # Create a server and serve
        self._server = subprocess.Popen([
            sys.executable,
            SERVER_SCRIPT_PATH,
            '-p', str(active_remote.port),
            '-t', str(timeout),
            '-pwd', str(password),
            '--attached',
            '--log-level', server_log_level,
            '--log-file', os.path.join(cache_directory, 'multiuser_server.log')]
        )

        self.connect(repository=repository,
                     remote=remote,
                     timeout=timeout,
                     password=password)

    def init(self):
        """ Init the repository data

            commit and push initial graph to the server
        """
        if len(self.repository.object_store) == 0:
            logging.error("Add some data first")
            return

        self.request_server_repository_init()

    def disconnect(self, reason: str = 'None'):
        """Disconnect from session
        """
        self._state = STATE_QUITTING
        self.remote.command.close()
        self.remote.data.close()
        logging.info(f"Exiting session, reason {reason}")
        if self.is_connection_alive():
            self._ttl.kill()
        if self._server:
            self._server.kill()
        self.context.destroy()
        self.call_registered('on_exit', reason=reason)

        self._state = STATE_INITIAL

    # NETWORKING
    def is_connection_alive(self):
        return self._ttl.poll() is None

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

    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 = Snapshot(
            owner=self._id,
            data={'STATE': "REQUEST_INIT"})

        self.repository.push_cmd(self.remote.command, snapshot_request)

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

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

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

        snapshot_cmd = ServerSnapshot(
            owner='server',
            data={'STATE': 'INIT',
                  'SIZE': self._srv_snapshot_size,
                  'NODES': keys})
        self.repository.push_cmd(self.remote.command, snapshot_cmd)

        self._state = STATE_SRV_SYNC

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

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

        connection_status = command.data

        if 'LOBBY' in connection_status:
            self._state = STATE_LOBBY
            if self._server:
                self.request_server_repository_init()
        if 'RUNNING' in connection_status:
            self._state = STATE_LOBBY
            self.request_snapshot_init()
        if 'FAILED' in connection_status:
            self.disconnect(
                reason=f"Failed to connect, authentification refused [{connection_status}]")

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

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

        snapshot_state = command.data['STATE']
        snapshot_data = command.data.get('DATA')
        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._snapshot_late_updates = queue.Queue()

            self._state_progress = {
                '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
        elif snapshot_state == 'SET':
            node  = ReplicationObject.from_raw_chunks(snapshot_data)
            node.state = FETCHED

            if node.uuid == self._current_snapshot:
                self._snapshot_progress += 1
                self._state_progress = {
                    'current': self._snapshot_progress,
                    'total': self._snapshot_total
                }
                self.repository.do_commit(node)

                if not self._snapshot_catalog:
                    # Apply late updates
                    while not self._snapshot_late_updates.empty():
                        late_update = self._snapshot_late_updates.get()
                        logging.info(f"Applying late update: {late_update.uuid}")
                        self.repository.do_commit(late_update)

                    snapshot_request = Snapshot(
                        owner=self._id,
                        data={'STATE': "DONE"})
                    self.repository.push_cmd(self.remote.command, snapshot_request)
                    logging.info("Snapshot done.")
                    self._state = STATE_ACTIVE
                    self.call_registered('on_connection')
                else:
                    self._current_snapshot = self._snapshot_catalog.pop()
                    self.get_snapshot(self._current_snapshot)

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

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

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

        if cli_snapshot_state == 'ACCEPTED':
            for index, node in enumerate(self.repository.object_store.values()):
                porcelain.commit(self.repository, node.uuid)
                node.state = UP
                snapshot_cmd = ServerSnapshot(
                    owner='server',
                    data={
                        'STATE': 'SET',
                        'DATA': node.as_raw_chunks()
                    }
                )

                self.repository.push_cmd(self.remote.command, snapshot_cmd)
                self._state_progress = {
                    'current': index,
                    'total': len(self.repository.object_store)
                }

            snapshot_cmd = ServerSnapshot(
                owner='server',
                data={'STATE': 'END'})
            self.repository.push_cmd(self.remote.command, snapshot_cmd)
        elif cli_snapshot_state == 'DONE':
            self._state = STATE_ACTIVE
            self.call_registered('on_connection')
        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 = Snapshot(
            owner=self._id,
            data={
                'STATE': "GET",
                'ID': id})

        self.repository.push_cmd(self.remote.command, snapshot_request)

    def listen(self, timeout: int = 0):
        """ Non-blocking network listening

            :param timeout: network packet waiting time in millisecond
            :type timeout: int
        """
        if self._state in [STATE_INITIAL, STATE_QUITTING]:
            return

        if not self.is_connection_alive():
            self.disconnect(reason='Server lost')
            return

        sockets = dict(self.remote.poller.poll(timeout))
        # COMMANDS I/O
        if self.remote.command in sockets:
            try:
                command = self.repository.fetch_cmd(self.remote.command)
            except Exception as e:
                logging.error(
                    f"Corrupted frame received, skipping it. Cause:{e}")
                traceback.print_exc()
            else:
                # AUTHENTIFICATION
                if isinstance(command, Auth):
                    self.handle_authentification(command)

                # DISCONNECT CONFIRMATION
                if isinstance(command, Disconnect):
                    self.disconnect(reason=command.data)

                # CLIENTS INFO UPDATE
                if isinstance(command, UpdateClientsState):
                    self.repository.remote.online_users = command.data
                
                if isinstance(command, UpdateUserMetadata):
                    user = self.repository.remote.online_users.get(command.owner)
                    user['metadata'].update(command.data)

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

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

                # GRAPH OPERATION (DELETE, CHANGE_RIGHT)
                if type(command) in [Delete, Right]:
                    command.execute(self.repository.object_store)

        # DATA IN
        if self.remote.data in sockets:
            try:
                datablock = self.repository.fetch_data(self.remote.data)
            except Exception as e:
                logging.error(
                    f"Corrupted frame received, skipping it. Cause:{e}")
                traceback.print_exc()
            else:
                # Client snapshot
                if self._state == STATE_SYNCING:
                        # If the snapshot is expected in the snapshot catalog we store
                        logging.info("Adding an update for the late one...")
                        self._snapshot_late_updates.put(datablock)

                # Store received updates
                if self._state == STATE_ACTIVE:
                    self.repository.do_commit(datablock)

        # Various timeout checks
        # auth
        if self._state == STATE_AUTH:
            if (current_milli_time()-self._connection_start_time) > self._connection_timeout:
                self.disconnect(
                    reason='Connection failed, server not found')

    @property
    def state(self):
        """Get active session state

        :return: session state
        """
        return self._state

    @property
    def state_progress(self):
        """Get active session state

        :return: session state
        """
        return self._state_progress

    @property
    def online_users(self):
        return self.repository.remote.online_users

    @property
    def id(self):
        return self._id

    def is_admin(self):
        return self.repository.remote.online_users[self.id]['admin']


this.session = Session()
