import os
import datetime
import logging
import logging.config
import yaml
try:
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader
import redis
import uuid
import traceback
import threading

import consul
import pika
from flask import Flask
from flask import jsonify

from biomaj_download.download.curl import CurlDownload
from biomaj_download.download.direct import DirectFTPDownload
from biomaj_download.download.direct import DirectHTTPDownload
from biomaj_download.download.localcopy import LocalDownload
from biomaj_download.message import downmessage_pb2
from biomaj_download.download.rsync import RSYNCDownload
from biomaj_core.utils import Utils
from biomaj_zipkin.zipkin import Zipkin
from biomaj_download.download.protocolirods import IRODSDownload

app = Flask(__name__)
app_log = logging.getLogger('werkzeug')
app_log.setLevel(logging.ERROR)

# Classify protocols from downmessage.proto
# Note: those lists are based on the protocol numbers, not the protocol names
ALL_PROTOCOLS = [item for key, item in downmessage_pb2.DownloadFile.Protocol.items()]
DIRECT_PROTOCOLS = [
    item for key, item in downmessage_pb2.DownloadFile.Protocol.items()
    if key.startswith("DIRECT")
]


@app.route('/api/download-message')
def ping():
    return jsonify({'msg': 'pong'})


def start_web(config):
    app.run(host='0.0.0.0', port=config['web']['port'])


def consul_declare(config):
    if config['consul']['host']:
        consul_agent = consul.Consul(host=config['consul']['host'])
        consul_agent.agent.service.register(
            'biomaj-download-message',
            service_id=config['consul']['id'],
            address=config['web']['hostname'],
            port=config['web']['port'],
            tags=['biomaj']
        )
        check = consul.Check.http(
            url='http://' + config['web']['hostname'] + ':' + str(config['web']['port']) + '/api/download-message',
            interval=20
        )
        consul_agent.agent.check.register(
            config['consul']['id'] + '_check',
            check=check,
            service_id=config['consul']['id']
        )
        return True
    else:
        return False


class DownloadService(object):

    channel = None
    redis_client = None

    def supervise(self):
        if consul_declare(self.config):
            web_thread = threading.Thread(target=start_web, args=(self.config,))
            web_thread.start()

    def __init__(self, config_file=None, rabbitmq=True):
        self.logger = logging
        self.session = None
        self.bank = None
        self.download_callback = None
        with open(config_file, 'r') as ymlfile:
            self.config = yaml.load(ymlfile, Loader=Loader)
            Utils.service_config_override(self.config)

        Zipkin.set_config(self.config)

        if 'log_config' in self.config:
            for handler in list(self.config['log_config']['handlers'].keys()):
                self.config['log_config']['handlers'][handler] = dict(self.config['log_config']['handlers'][handler])
            logging.config.dictConfig(self.config['log_config'])
            self.logger = logging.getLogger('biomaj')

        if not self.redis_client:
            self.redis_client = redis.StrictRedis(host=self.config['redis']['host'],
                                                  port=self.config['redis']['port'],
                                                  db=self.config['redis']['db'],
                                                  decode_responses=True)

        if rabbitmq and not self.channel:
            connection = None
            rabbitmq_port = self.config['rabbitmq']['port']
            rabbitmq_user = self.config['rabbitmq']['user']
            rabbitmq_password = self.config['rabbitmq']['password']
            rabbitmq_vhost = self.config['rabbitmq']['virtual_host']
            if rabbitmq_user:
                credentials = pika.PlainCredentials(rabbitmq_user, rabbitmq_password)
                connection = pika.BlockingConnection(pika.ConnectionParameters(self.config['rabbitmq']['host'], rabbitmq_port, rabbitmq_vhost, credentials, heartbeat_interval=0))
            else:
                connection = pika.BlockingConnection(pika.ConnectionParameters(self.config['rabbitmq']['host'], heartbeat_interval=0))
            self.channel = connection.channel()
            self.logger.info('Download service started')

    def close(self):
        if self.channel:
            try:
                self.channel.close()
            except Exception as e:
                logging.warn('Download:Service:Exception:' + str(e))

    def on_download_callback(self, func):
        self.download_callback = func

    def get_handler(self, protocol_name, server, remote_dir, remote_files=[],
                    credentials=None, http_parse=None, http_method=None, param=None,
                    proxy=None, proxy_auth='',
                    save_as=None, timeout_download=None, offline_dir=None,
                    options={}):
        protocol = downmessage_pb2.DownloadFile.Protocol.Value(protocol_name.upper())
        downloader = None
        if protocol in [0, 1]:  # FTP, SFTP
            downloader = CurlDownload(protocol_name, server, remote_dir)
        if protocol in [2, 3]:  # HTTP, HTTPS (could be factored with previous case)
            downloader = CurlDownload(protocol_name, server, remote_dir, http_parse)
        if protocol == 4:  # DirectFTP
            downloader = DirectFTPDownload("ftp", server, '/')
        if protocol == 5:  # DirectHTTP
            downloader = DirectHTTPDownload("http", server, '/')
        if protocol == 6:  # DirectHTTPS
            downloader = DirectHTTPDownload("https", server, '/')
        if protocol == 10:  # DirectFTPS
            downloader = DirectFTPDownload('ftps', server, '/')
        if protocol == 7:  # Local
            downloader = LocalDownload(remote_dir)
        if protocol == 8:  # RSYNC
            downloader = RSYNCDownload(server, remote_dir)
        if protocol == 9:  # iRods
            downloader = IRODSDownload(server, remote_dir)
        if downloader is None:
            return None

        for remote_file in remote_files:
            if remote_file['save_as']:
                save_as = remote_file['save_as']
        # For direct protocols, we only keep base name
        if protocol in DIRECT_PROTOCOLS:
            tmp_remote = []
            for remote_file in remote_files:
                tmp_remote.append(remote_file['name'])
            remote_files = tmp_remote

        if http_method is not None:
            downloader.set_method(http_method)

        if offline_dir:
            downloader.set_offline_dir(offline_dir)

        if proxy is not None and proxy:
            downloader.set_proxy(proxy, proxy_auth)

        if timeout_download is not None and timeout_download:
            downloader.set_timeout(timeout_download)

        if credentials:
            downloader.set_credentials(credentials)

        if save_as:
            downloader.set_save_as(save_as)

        if param:
            downloader.set_param(param)

        downloader.set_server(server)

        # Set the name of the BioMAJ protocol to which we respond.
        downloader.set_protocol(protocol_name)

        if options is not None:
            self.logger.debug("Received options: " + str(options))
            downloader.set_options(options)

        downloader.logger = self.logger
        downloader.set_files_to_download(remote_files)
        return downloader

    def _get_handler(self, biomaj_file_info):
        """
        Get a protocol download handler
        """

        protocol = biomaj_file_info.remote_file.protocol
        server = biomaj_file_info.remote_file.server
        remote_dir = biomaj_file_info.remote_file.remote_dir

        protocol_name = downmessage_pb2.DownloadFile.Protocol.Name(protocol).lower()
        self.logger.debug('%s request to download from %s://%s' % (biomaj_file_info.bank, protocol_name, server))

        remote_files = []
        for remote_file in biomaj_file_info.remote_file.files:
            remote_files.append({
                                'name': remote_file.name,
                                'save_as': remote_file.save_as,
                                'year': remote_file.metadata.year,
                                'month': remote_file.metadata.month,
                                'day': remote_file.metadata.day,
                                'root': remote_file.root
                                })

        proxy = None
        proxy_auth = ''
        if biomaj_file_info.proxy is not None:
            proxy = biomaj_file_info.proxy.proxy
            proxy_auth = biomaj_file_info.proxy.proxy_auth

        params = None
        if biomaj_file_info.remote_file.param:
            params = {}
            for param in biomaj_file_info.remote_file.param:
                params[param.name] = param.value
        return self.get_handler(protocol_name, server, remote_dir,
                                remote_files=remote_files,
                                credentials=biomaj_file_info.remote_file.credentials,
                                http_parse=biomaj_file_info.remote_file.http_parse,
                                http_method=downmessage_pb2.DownloadFile.HTTP_METHOD.Name(biomaj_file_info.http_method),
                                param=params,
                                proxy=proxy,
                                proxy_auth=proxy_auth,
                                save_as=biomaj_file_info.remote_file.save_as,
                                timeout_download=biomaj_file_info.timeout_download,
                                offline_dir=biomaj_file_info.local_dir,
                                options=biomaj_file_info.options
                                )

    def clean(self, biomaj_file_info=None):
        '''
        Clean session and download info
        '''
        session = self.session
        bank = self.bank
        if biomaj_file_info:
            session = biomaj_file_info.session
            bank = biomaj_file_info.bank

        self.logger.debug('Clean %s session %s' % (bank, session))
        self.redis_client.delete(self.config['redis']['prefix'] + ':' + bank + ':session:' + session)
        self.redis_client.delete(self.config['redis']['prefix'] + ':' + bank + ':session:' + session + ':error')
        self.redis_client.delete(self.config['redis']['prefix'] + ':' + bank + ':session:' + session + ':progress')
        self.redis_client.delete(self.config['redis']['prefix'] + ':' + bank + ':session:' + session + ':files')
        self.redis_client.delete(self.config['redis']['prefix'] + ':' + bank + ':session:' + session + ':error:info')

    def _create_session(self, bank):
        '''
        Creates a unique session
        '''
        self.session = str(uuid.uuid4())
        self.redis_client.set(self.config['redis']['prefix'] + ':' + bank + ':session:' + self.session, 1)
        self.logger.debug('Create %s new session %s' % (bank, self.session))
        self.bank = bank
        return self.session

    def download_errors(self, biomaj_file_info):
        '''
        Get errors
        '''
        errors = []
        error = self.redis_client.rpop(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':error:info')
        while error:
            errors.append(error)
            error = self.redis_client.rpop(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':error:info')
        return errors

    def download_status(self, biomaj_file_info):
        '''
        Get current status
        '''

        error = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':error')
        progress = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':progress')
        if error is None:
            error = -1
        if progress is None:
            progress = -1
        return (int(progress), int(error))

    def list_status(self, biomaj_file_info):

        list_progress = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':progress')
        if list_progress:
            return True
        else:
            return False

    def list_result(self, biomaj_file_info, protobuf_decode=True):
        '''
        Get file list result
        '''

        file_list = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':files')
        if protobuf_decode:
            file_list_pb2 = downmessage_pb2.FileList()
            file_list_pb2.ParseFromString(file_list_pb2)
            return file_list_pb2

        return file_list

    def _list(self, download_handler, biomaj_file_info):
        '''
        List remote content, no session management
        '''
        file_list = []
        dir_list = []
        file_list_pb2 = downmessage_pb2.FileList()

        try:
            (file_list, dir_list) = download_handler.list()
            download_handler.match(biomaj_file_info.remote_file.matches, file_list, dir_list)
        except Exception as e:
            self.logger.error('List exception for bank %s: %s' % (biomaj_file_info.bank, str(e)))
            self.redis_client.set(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':error', 1)
            self.redis_client.lpush(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':error:info', str(e))
        else:
            self.logger.debug('End of download for %s session %s' % (biomaj_file_info.bank, biomaj_file_info.session))
            for file_elt in download_handler.files_to_download:
                file_pb2 = file_list_pb2.files.add()
                file_pb2.name = file_elt['name']
                file_pb2.root = file_elt['root']
                if 'save_as' in file_elt:
                    file_pb2.save_as = file_elt['save_as']
                if 'url' in file_elt:
                    file_pb2.url = file_elt['url']
                if 'param' in file_elt and file_elt['param']:
                    for key in list(file_elt['param'].keys()):
                        param = file_list_pb2.param.add()
                        param.name = key
                        param.value = file_elt['param'][key]
                metadata = downmessage_pb2.File.MetaData()
                metadata.permissions = file_elt['permissions']
                metadata.group = file_elt['group']
                metadata.size = int(file_elt['size'])
                metadata.hash = file_elt['hash']
                metadata.year = int(file_elt['year'])
                metadata.month = int(file_elt['month'])
                metadata.day = int(file_elt['day'])
                if 'format' in file_elt:
                    metadata.format = file_elt['format']
                file_pb2.metadata.MergeFrom(metadata)
        return file_list_pb2

    def list(self, biomaj_file_info):
        '''
        List remote content
        '''
        self.logger.debug('New list request %s session %s' % (biomaj_file_info.bank, biomaj_file_info.session))
        session = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session)
        if not session:
            self.logger.debug('Session %s for bank %s has expired, skipping download of %s' % (biomaj_file_info.session, biomaj_file_info.bank, biomaj_file_info.remote_file.files))
            return
        download_handler = self._get_handler(biomaj_file_info)
        if download_handler is None:
            self.logger.error('Could not get a handler for %s with session %s' % (biomaj_file_info.bank, biomaj_file_info.session))

        file_list_pb2 = self._list(download_handler, biomaj_file_info)

        self.redis_client.set(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':files', str(file_list_pb2.SerializeToString()))
        self.redis_client.incr(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':progress')

    def local_download(self, biomaj_file_info):
        '''
        Download files, no session
        '''
        download_handler = self._get_handler(biomaj_file_info)
        if download_handler is None:
            self.logger.error('Could not get a handler for %s with session %s' % (biomaj_file_info.bank, biomaj_file_info.session))
        downloaded_files = download_handler.download(biomaj_file_info.local_dir)
        self.logger.debug("Downloaded " + str(len(downloaded_files)) + " file in " + biomaj_file_info.local_dir)
        self.get_file_info(biomaj_file_info.local_dir, downloaded_files)
        return downloaded_files

    def download(self, biomaj_file_info):
        '''
        Download files

        Store in redis the progress and count of errors under:
         - prefix:bank_name:session:session_id:error
         - prefix:bank_name:session:session_id:progress
        '''

        self.logger.debug('New download request %s session %s' % (biomaj_file_info.bank, biomaj_file_info.session))
        session = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session)
        if not session:
            self.logger.debug('Session %s for bank %s has expired, skipping download of %s' % (biomaj_file_info.session, biomaj_file_info.bank, biomaj_file_info.remote_file.files))
            return
        downloaded_files = []
        try:
            downloaded_files = self.local_download(biomaj_file_info)
        except Exception as e:
            self.logger.exception("Download error:%s:%s:%s" % (biomaj_file_info.bank, biomaj_file_info.session, str(e)))
            session = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session)
            if session:
                # If session deleted, do not track
                self.redis_client.incr(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':error')
                self.redis_client.lpush(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':error:info', str(e))
        else:
            self.logger.debug('End of download for %s session %s' % (biomaj_file_info.bank, biomaj_file_info.session))

        session = self.redis_client.get(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session)
        if session:
            # If session deleted, do not track
            self.redis_client.incr(self.config['redis']['prefix'] + ':' + biomaj_file_info.bank + ':session:' + biomaj_file_info.session + ':progress')
        return downloaded_files

    def get_file_info(self, local_dir, downloaded_files):
        if downloaded_files is None:
            return
        for downloaded_file in downloaded_files:
            # file_dir = local_dir + '/' + os.path.dirname(downloaded_file['save_as'])
            file_path = local_dir + '/' + downloaded_file['save_as']
            fstat = os.stat(file_path)
            downloaded_file['permissions'] = str(fstat.st_mode)
            downloaded_file['group'] = str(fstat.st_gid)
            downloaded_file['user'] = str(fstat.st_uid)
            downloaded_file['size'] = str(fstat.st_size)
            fstat_mtime = datetime.datetime.fromtimestamp(fstat.st_mtime)
            downloaded_file['month'] = fstat_mtime.month
            downloaded_file['day'] = fstat_mtime.day
            downloaded_file['year'] = fstat_mtime.year

    def ask_download(self, biomaj_info_file):
        self.channel.basic_publish(
            exchange='',
            routing_key='biomajdownload',
            body=biomaj_info_file.SerializeToString(),
            properties=pika.BasicProperties(
                # make message persistent
                delivery_mode=2
            ))

    def callback_messages(self, ch, method, properties, body):
        '''
        Manage download and send ACK message
        '''
        try:
            operation = downmessage_pb2.Operation()
            operation.ParseFromString(body)
            message = operation.download
            span = None
            if operation.trace and operation.trace.trace_id:
                url = str(message.remote_file.protocol) + ':' + str(message.remote_file.server) + ':' + str(message.remote_file.remote_dir)
                span = Zipkin('biomaj-download-executor', str(message.remote_file.server), trace_id=operation.trace.trace_id, parent_id=operation.trace.span_id)
                span.add_binary_annotation('url', url)
                span.add_binary_annotation('local_dir', str(message.local_dir))

            self.logger.debug('Received message: %s' % (message))
            if operation.type == 0:
                message = operation.download
                self.logger.debug('List operation %s, %s' % (message.bank, message.session))
                if len(message.remote_file.matches) == 0:
                    self.logger.error('No pattern match for a list operation')
                else:
                    self.list(message)
            elif operation.type == 1:
                message = operation.download
                self.logger.debug('Download operation %s, %s' % (message.bank, message.session))
                downloaded_files = self.download(message)
                if self.download_callback is not None:
                    self.download_callback(message.bank, downloaded_files)
            else:
                self.logger.warn('Wrong message type, skipping')
            if span:
                span.trace()
        except Exception as e:
            self.logger.error('Error with message: %s' % (str(e)))
            traceback.print_exc()
        ch.basic_ack(delivery_tag=method.delivery_tag)

    def wait_for_messages(self):
        '''
        Loop queue waiting for messages
        '''
        self.channel.queue_declare(queue='biomajdownload', durable=True)
        self.channel.basic_qos(prefetch_count=1)
        self.channel.basic_consume(
            self.callback_messages,
            queue='biomajdownload')
        self.channel.start_consuming()
