#!/usr/bin/env python3

# Copyright (C) 2017  The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information

import datetime
from fnmatch import fnmatch
from operator import itemgetter
import os
import sys

import click


def list_remote_workers(inspect):
    ping_replies = inspect.ping()
    if not ping_replies:
        return {}
    workers = list(sorted(ping_replies))
    ret = {}

    for worker_name in workers:
        if not worker_name.startswith('celery@'):
            print('Unsupported worker: %s' % worker_name, file=sys.stderr)
            continue
        type, host = worker_name[len('celery@'):].split('.', 1)
        worker = {
            'name': worker_name,
            'host': host,
            'type': type,
        }
        ret[worker_name] = worker

    return ret


def make_filters(filter_host, filter_type):
    """Parse the filters and create test functions"""

    def include(field, value):
        def filter(worker, field=field, value=value):
            return fnmatch(worker[field], value)
        return filter

    def exclude(field, value):
        def filter(worker, field=field, value=value):
            return not fnmatch(worker[field], value)
        return filter

    filters = []
    for host in filter_host:
        if host.startswith('-'):
            filters.append(exclude('host', host[1:]))
        else:
            filters.append(include('host', host))

    for type_ in filter_type:
        if type_.startswith('-'):
            filters.append(exclude('type', type_[1:]))
        else:
            filters.append(include('type', type_))

    return filters


def filter_workers(workers, filters):
    """Filter workers according to the set criteria"""
    return {name: worker
            for name, worker in workers.items()
            if all(check(worker) for check in filters)}


def get_clock_offsets(workers, inspect):
    """Add a clock_offset entry for each worker"""
    err_msg = 'Could not get monotonic clock for {worker}'

    t = datetime.datetime.now(tz=datetime.timezone.utc)
    for worker, clock in inspect._request('monotonic').items():
        monotonic = clock.get('monotonic')
        if monotonic is None:
            monotonic = 0
            click.echo(err_msg.format(worker=worker), err=True)
        dt = datetime.timedelta(seconds=monotonic)
        workers[worker]['clock_offset'] = t - dt


def worker_to_wallclock(worker, monotonic):
    """Convert a monotonic timestamp from a worker to a wall clock time"""
    dt = datetime.timedelta(seconds=monotonic)
    return worker['clock_offset'] + dt


@click.group()
@click.option('--instance-config', metavar='CONFIG', default=None,
              help='Use this worker instance configuration')
@click.option('--host', metavar='HOSTNAME_FILTER', multiple=True,
              help='Filter by hostname')
@click.option('--type', metavar='WORKER_TYPE_FILTER', multiple=True,
              help='Filter by worker type')
@click.option('--timeout', metavar='TIMEOUT', type=float, default=1.0,
              help='Timeout for remote control communication')
@click.option('--debug/--no-debug', default=False, help='Turn on debugging')
@click.pass_context
def cli(ctx, debug, timeout, instance_config, host, type):
    """Manage the Software Heritage workers

    Filters support globs; a filter starting with a "-" excludes the
    corresponding values.

    """
    if instance_config:
        os.environ['SWH_WORKER_INSTANCE'] = instance_config

    from swh.scheduler.celery_backend.config import app
    full_inspect = app.control.inspect(timeout=timeout)

    workers = filter_workers(
        list_remote_workers(full_inspect),
        make_filters(host, type)
    )
    ctx.obj['workers'] = workers

    destination = list(workers)
    inspect = app.control.inspect(destination=destination,
                                  timeout=timeout)
    ctx.obj['inspect'] = inspect

    get_clock_offsets(workers, inspect)

    ctx.obj['control'] = app.control
    ctx.obj['destination'] = destination
    ctx.obj['timeout'] = timeout
    ctx.obj['debug'] = debug


@cli.command()
@click.pass_context
def list_workers(ctx):
    """List the currently running workers"""
    workers = ctx.obj['workers']

    for worker_name, worker in sorted(workers.items()):
        click.echo("{type} alive on {host}".format(**worker))

    if not workers:
        sys.exit(2)


@cli.command()
@click.pass_context
def list_tasks(ctx):
    """List the tasks currently running on workers"""
    task_template = ('{worker} {name}'
                     '[{id} '
                     'started={started:%Y-%m-%mT%H:%M:%S} '
                     'pid={worker_pid}] {args} {kwargs}')
    inspect = ctx.obj['inspect']
    workers = ctx.obj['workers']
    active = inspect.active()

    if not active:
        click.echo('No reply from workers', err=True)
        sys.exit(2)

    has_tasks = False
    for worker_name, tasks in sorted(active.items()):
        worker = workers[worker_name]
        if not tasks:
            click.echo("No active tasks on {name}".format(**worker), err=True)
        for task in sorted(tasks, key=itemgetter('time_start')):
            task['started'] = worker_to_wallclock(worker, task['time_start'])
            click.echo(task_template.format(worker=worker_name, **task))
            has_tasks = True

    if not has_tasks:
        sys.exit(2)


@cli.command()
@click.pass_context
def list_queues(ctx):
    """List all the queues currently enabled on the workers"""
    inspect = ctx.obj['inspect']
    active = inspect.active_queues()

    if not active:
        click.echo('No reply from workers', err=True)
        sys.exit(2)

    has_queues = False
    for worker_name, queues in sorted(active.items()):
        queues = sorted(queue['name'] for queue in queues)
        if queues:
            click.echo('{worker} {queues}'.format(worker=worker_name,
                                                  queues=' '.join(queues)))
            has_queues = True
        else:
            click.echo('No queues for {worker}'.format(worker=worker_name),
                       err=True)

    if not has_queues:
        sys.exit(2)


@cli.command()
@click.option('--noop', is_flag=True, default=False, help='Do not proceed')
@click.argument('queues', nargs=-1)
@click.pass_context
def remove_queues(ctx, noop, queues):
    """Cancel the queue for the given workers"""
    msg_template = 'Canceling queue {queue} on worker {worker}{noop}'

    inspect = ctx.obj['inspect']
    control = ctx.obj['control']
    timeout = ctx.obj['timeout']
    active = inspect.active_queues()

    if not queues:
        queues = ['*']

    if not active:
        click.echo('No reply from workers', err=True)
        sys.exit(2)

    for worker, active_queues in sorted(active.items()):
        for queue in sorted(active_queues, key=itemgetter('name')):
            if any(fnmatch(queue['name'], name) for name in queues):
                msg = msg_template.format(queue=queue['name'], worker=worker,
                                          noop=' (noop)' if noop else '')
                click.echo(msg, err=True)
                if not noop:
                    control.cancel_consumer(queue['name'],
                                            destination=[worker],
                                            timeout=timeout)


@cli.command()
@click.option('--noop', is_flag=True, default=False, help='Do not proceed')
@click.argument('queues', nargs=-1)
@click.pass_context
def add_queues(ctx, noop, queues):
    """Start the queue for the given workers"""
    msg_template = 'Starting queue {queue} on worker {worker}{noop}'

    control = ctx.obj['control']
    timeout = ctx.obj['timeout']
    workers = ctx.obj['workers']

    if not workers:
        click.echo('No reply from workers', err=True)
        sys.exit(2)

    for worker in sorted(workers):
        for queue in queues:
            msg = msg_template.format(queue=queue, worker=worker,
                                      noop=' (noop)' if noop else '')
            click.echo(msg, err=True)
            if not noop:
                ret = control.add_consumer(queue,
                                           destination=[worker],
                                           timeout=timeout)
                print(ret)


if __name__ == '__main__':
    cli(obj={})