#! /usr/bin/env python

import argparse
import multiprocessing as mp
import os
import sys

import qless
from qless import logger
from qless.workers.forking import ForkingWorker


# First off, read the arguments
parser = argparse.ArgumentParser(
    description="Run qless workers.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)

# Options specific to the client we're making'
parser.add_argument(
    "--host",
    dest="host",
    default="redis://localhost:6379",
    help="The redis:// url to connect to",
)
parser.add_argument(
    "-n",
    "--name",
    default=None,
    type=str,
    help="The hostname to identify your worker as",
)

# Options that we consume in this binary
parser.add_argument(
    "-p",
    "--path",
    action="append",
    default=[],
    help="Path(s) to include when loading jobs",
)
parser.add_argument(
    "-v", "--verbose", default=False, action="store_true", help="Be extra talkative"
)
parser.add_argument(
    "-l", "--logger", default=False, type=str, help="Log out to this file"
)
parser.add_argument(
    "-m",
    "--import",
    action="append",
    default=[],
    help="The modules to preemptively import",
)
parser.add_argument("-d", "--workdir", default=".", help="The base work directory path")

# Options specific to the worker we're instantiating
parser.add_argument(
    "-w",
    "--workers",
    default=mp.cpu_count(),
    type=int,
    help="How many processes to run. Set to 0 to use all available cores",
)
parser.add_argument(
    "-q", "--queue", action="append", default=[], help="The queues to pull work from"
)

# Options specific to forking greenlet workers
parser.add_argument(
    "-g",
    "--greenlets",
    default=0,
    type=int,
    help="How many greenlets to run in each process (if used, uses gevent)",
)
parser.add_argument(
    "-i", "--interval", default=60, type=int, help="The polling interval"
)
parser.add_argument(
    "-r",
    "--resume",
    default=False,
    action="store_true",
    help="Try to resume jobs that this worker had previously been working on",
)
args = parser.parse_args()

# Build up the kwargs that we'll pass to the worker
kwargs = {"workers": args.workers, "interval": args.interval, "resume": args.resume}

# If we're supposed to use greenlets...
if args.greenlets:
    from gevent import monkey

    monkey.patch_all()
    kwargs.update(
        {"klass": "qless.workers.greenlet.GeventWorker", "greenlets": args.greenlets}
    )


# Add each of the paths to the python search path
sys.path = [os.path.abspath(p) for p in args.path] + sys.path

# Be verbose if requested.
if args.verbose:
    import logging

    logger.setLevel(logging.DEBUG)

# Log to the provided file, if requested.
if args.logger:
    import logging

    from qless import formatter

    handler = logging.FileHandler(args.logger)
    handler.setFormatter(formatter)
    handler.setLevel(logging.DEBUG)
    logger.addHandler(handler)

# Import all the modules and packages we've been asked to import
for module in getattr(args, "import"):
    try:
        logger.info("Loaded %s" % repr(__import__(module)))
    except Exception:
        logger.exception("Failed to import %s" % module)

# Change path to our working directory
os.chdir(args.workdir)

# And now run the worker
ForkingWorker(args.queue, qless.Client(args.host, hostname=args.name), **kwargs).run()
