from aardvark import create_app, db
from aardvark.updater import AccountToUpdate
from flask import current_app
from swag_client.swag import get_all_accounts
import json
from bunch import Bunch
from flask_script import Manager, Command, Option
import better_exceptions # noqa
import Queue
import threading

manager = Manager(create_app)

ACCOUNT_QUEUE = Queue.Queue()
DB_LOCK = threading.Lock()
QUEUE_LOCK = threading.Lock()
UPDATE_DONE = False


class UpdateAccountThread(threading.Thread):
    global ACCOUNT_QUEUE, DB_LOCK, QUEUE_LOCK, UPDATE_DONE

    def __init__(self, thread_ID):
        self.thread_ID = thread_ID
        threading.Thread.__init__(self)
        self.app = current_app._get_current_object()

    def run(self):
        while not UPDATE_DONE:

            QUEUE_LOCK.acquire()

            if not ACCOUNT_QUEUE.empty():
                (account_num, role_name, arns) = ACCOUNT_QUEUE.get()

                self.app.logger.info("Thread #{} updating account {} with {} arns".format(
                                     self.thread_ID, account_num, 'all' if arns[0] == 'all' else len(arns)))

                QUEUE_LOCK.release()

                account = AccountToUpdate(self.app, account_num, role_name, arns)
                ret_code, aa_data = account.update_account()

                if ret_code != 0:  # retrieve wasn't successful, put back on queue
                    QUEUE_LOCK.acquire()
                    ACCOUNT_QUEUE.put((account_num, role_name, arns))
                    QUEUE_LOCK.release()

                self.app.logger.info("Thread #{} persisting data for account {}".format(self.thread_ID, account_num))

                DB_LOCK.acquire()
                persist_aa_data(self.app, aa_data)
                DB_LOCK.release()

            else:
                QUEUE_LOCK.release()


def persist_aa_data(app, aa_data):
    """
    Reads access advisor JSON file & persists to our database
    """
    from aardvark.model import AWSIAMObject, AdvisorData

    aa = json.loads(aa_data)

    with app.app_context():
        arn_cache = {}
        for arn, data in aa.items():
            if arn in arn_cache:
                item = arn_cache[arn]
            else:
                item = AWSIAMObject.get_or_create(arn)
                arn_cache[arn] = item
            for service in data:
                AdvisorData.create_or_update(item.id,
                                             service['lastAuthenticated'],
                                             service['serviceName'],
                                             service['serviceNamespace'],
                                             service['lastAuthenticatedEntity'],
                                             service['totalAuthenticatedEntities'])
        db.session.commit()


@manager.command
def drop_db():
    """ Drops the database. """
    db.drop_all()


@manager.command
def create_db():
    """ Creates the database. """
    db.create_all()


@manager.command
def config():
    """
    Creates a config file.

    SWAG_BUCKET = '...'
    SWAG_FILTER = '...'
    ROLENAME = '<ASSUME_ROLE_HERE>'
    REGION = 'us-east-1'
    NUM_THREADS = 5
    SQLALCHEMY_DATABASE_URI = 'postgresql://user:pass@localhost:5432/db' or
    SQLALCHEMY_DATABASE_URI = 'sqlite:///tmp/aardvark.db'
    PHANTOMJS = '/usr/local/bin/phantomjs'
    """
    bucket = raw_input('SWAG_BUCKET: ')
    swag_filter = raw_input('SWAG_FILTER:  ')
    role_name = raw_input('ROLENAME: ')
    region = raw_input('AWS REGION: ')
    db_uri = raw_input('DATABASE: ')
    num_threads = raw_input('Number of threads (5)') or 5  # testing shows problems with more than 6 threads
    phantom = raw_input('Path to phantomjs: ')

    log = """LOG_CFG = {
    'version': 1,
    'disable_existing_loggers': False,
    'formatters': {
        'standard': {
            'format': '%(asctime)s %(levelname)s: %(message)s '
                '[in %(pathname)s:%(lineno)d]'
        }
    },
    'handlers': {
        'file': {
            'class': 'logging.handlers.RotatingFileHandler',
            'level': 'DEBUG',
            'formatter': 'standard',
            'filename': 'aardvark.log',
            'maxBytes': 10485760,
            'backupCount': 100,
            'encoding': 'utf8'
        },
        'console': {
            'class': 'logging.StreamHandler',
            'level': 'DEBUG',
            'formatter': 'standard',
            'stream': 'ext://sys.stdout'
        }
    },
    'loggers': {
        'aardvark': {
            'handlers': ['file', 'console'],
            'level': 'DEBUG'
        }
    }
}"""

    with open('config.py', 'w') as filedata:
        filedata.write('# Autogenerated config file\n')
        filedata.write('SWAG_BUCKET = "{bucket}"\n'.format(bucket=bucket))
        filedata.write('SWAG_FILTER = "{filter}"\n'.format(filter=swag_filter))
        filedata.write('ROLENAME = "{role}"\n'.format(role=role_name))
        filedata.write('REGION = "{region}"\n'.format(region=region))
        filedata.write('SQLALCHEMY_DATABASE_URI = "{uri}"\n'.format(uri=db_uri))
        filedata.write('PHANTOMJS = "{phantom}"\n'.format(phantom=phantom))
        filedata.write('SQLALCHEMY_TRACK_MODIFICATIONS = False\n')
        filedata.write('NUM_THREADS = "{num_threads}"\n'.format(num_threads=num_threads))
        filedata.write(log)


@manager.option('-a', '--accounts', dest='accounts', type=unicode, default='all')
@manager.option('-r', '--arns', dest='arns', type=unicode, default='all')
def update(accounts, arns):
    """
    Asks AWS for new Access Advisor information.
    """
    accounts = _prep_accounts(accounts)
    arns = arns.split(',')
    app = create_app()

    global ACCOUNT_QUEUE, QUEUE_LOCK, UPDATE_DONE

    role_name = app.config.get('ROLENAME')
    num_threads = app.config.get('NUM_THREADS') or 5

    if num_threads > 6:
        current_app.logger.warn('Greater than 6 threads seems to cause problems')

    QUEUE_LOCK.acquire()
    for account in accounts:
        ACCOUNT_QUEUE.put((account.metadata['account_number'], role_name, arns))
    QUEUE_LOCK.release()

    threads = []
    for thread_num in range(num_threads):
        thread = UpdateAccountThread(thread_num + 1)
        thread.start()
        threads.append(thread)

    while not ACCOUNT_QUEUE.empty():
        pass
    UPDATE_DONE = True


def _prep_accounts(account_names):
    """
    Convert CLI provided account names into list of accounts from SWAG.
    Considers account aliases as well as account names.
    Returns a set of Bunch'd accounts.
    """
    current_app.logger.info('getting bucket {}'.format(
                            current_app.config.get('SWAG_BUCKET')))

    service = {'services': {current_app.config.get('SWAG_FILTER') or
               'security_monkey': {'enabled': True}}, 'ours': True}

    accounts = get_all_accounts(bucket=current_app.config.get('SWAG_BUCKET'),
                                **service).get('accounts')

    if account_names == 'all':
        return [Bunch(account) for account in accounts]

    lookup = {account['name']: Bunch(account) for account in accounts}
    for account in accounts:
        for alias in account['alias']:
            lookup[alias] = Bunch(account)

    matching_accounts = list()
    account_names = account_names.split(',')
    account_names = {name.lower() for name in account_names}

    for name in account_names:
        if name not in lookup:
            current_app.logger.warn('Could not find an account named %s'
                                    % name)
            continue

        matching_accounts.append(lookup[name])

    return matching_accounts


class GunicornServer(Command):
    """
    This is the main GunicornServer server, it runs the flask app with gunicorn and
    uses any configuration options passed to it.
    You can pass all standard gunicorn flags to this command as if you were
    running gunicorn itself.
    For example:
    aardvark start_api -w 4 -b 127.0.0.0:8002
    Will start gunicorn with 4 workers bound to 127.0.0.0:8002
    """
    description = 'Run the app within Gunicorn'

    def get_options(self):
        options = []
        try:
            from gunicorn.config import make_settings
        except ImportError:
            # Gunicorn does not yet support Windows.
            # See issue #524. https://github.com/benoitc/gunicorn/issues/524
            # For dev on Windows, make this an optional import.
            print('Could not import gunicorn, skipping.')
            return options

        settings = make_settings()
        for setting, klass in settings.items():
            if klass.cli:
                if klass.action:
                    if klass.action == 'store_const':
                        options.append(Option(*klass.cli, const=klass.const, action=klass.action))
                    else:
                        options.append(Option(*klass.cli, action=klass.action))
                else:
                    options.append(Option(*klass.cli))
        return options

    def run(self, *args, **kwargs):
        from gunicorn.app.wsgiapp import WSGIApplication

        app = WSGIApplication()

        app.app_uri = 'aardvark:create_app()'
        return app.run()


def main():
    manager.add_command("start_api", GunicornServer())
    manager.run()


if __name__ == '__main__':
    manager.add_command("start_api", GunicornServer())
    manager.run()
