import gevent
import gevent.pool
import os
import signal
import datetime
import time
import socket
import traceback
import psutil
import sys
import pymongo
from bson import ObjectId
from gevent.pywsgi import WSGIServer

from .job import Job
from .exceptions import JobTimeoutException, StopRequested, JobInterrupt
from .context import set_current_worker, set_current_job, get_current_job, connections, enable_greenlet_tracing
from .queue import Queue

# https://groups.google.com/forum/#!topic/gevent/EmZw9CVBC2g
# if "__pypy__" in sys.builtin_module_names:
#   def _reuse(self):
#       self._sock._reuse()
#   def _drop(self):
#       self._sock._drop()
#   gevent.socket.socket._reuse = _reuse
#   gevent.socket.socket._drop = _drop


class Worker(object):

  # Allow easy overloading
  job_class = Job

  def __init__(self, config):

    self.config = config

    set_current_worker(self)

    if self.config.get("trace_greenlets"):
      enable_greenlet_tracing()

    self.datestarted = datetime.datetime.utcnow()
    self.status = "init"
    self.queues = [x for x in self.config["queues"] if x]
    self.redis_queues = [Queue(x).redis_key for x in self.queues]
    self.done_jobs = 0
    self.max_jobs = self.config["max_jobs"]

    self.connected = False  # MongoDB + Redis

    self.exitcode = 0
    self.process = psutil.Process(os.getpid())
    self.greenlet = gevent.getcurrent()

    self.id = ObjectId()
    if config["name"]:
      self.name = self.config["name"]
    else:
      self.name = self.make_name()

    self.pool_size = self.config["gevent"]

    from .logger import LogHandler
    self.log_handler = LogHandler(quiet=self.config["quiet"])
    self.log = self.log_handler.get_logger(worker=self.id)

    self.log.info("Starting Gevent pool with %s worker greenlets (+ 1 monitoring)" % self.pool_size)

    self.gevent_pool = gevent.pool.Pool(self.pool_size)

    # Keep references to main greenlets
    self.greenlets = {}

    self.profiler = None
    if self.config["profile"]:
      print "Starting profiler..."
      import cProfile
      self.profiler = cProfile.Profile()
      self.profiler.enable()

  def connect(self, force=False):

    if self.connected and not force:
      return

    # Accessing connections attributes will automatically connect
    self.redis = connections.redis
    self.mongodb_jobs = connections.mongodb_jobs
    self.mongodb_logs = connections.mongodb_logs

    self.log_handler.set_collection(self.mongodb_logs.mrq_logs)

    self.connected = True

    # Be mindful that this is done each time a worker starts
    self.ensure_indexes()

  def ensure_indexes(self):

    if self.config["mongodb_logs_size"] > 0:

      try:
        self.mongodb_logs.create_collection("mrq_logs", capped=True, size=self.config["mongodb_logs_size"])
      except:
        pass

      try:
        self.mongodb_logs.command("convertToCapped", "mrq_logs", size=self.config["mongodb_logs_size"])
      except:
        pass

    self.mongodb_logs.mrq_logs.ensure_index([("job", 1)], background=True)
    self.mongodb_logs.mrq_logs.ensure_index([("worker", 1)], background=True, sparse=True)

    self.mongodb_logs.mrq_workers.ensure_index([("status", 1)], background=True)
    self.mongodb_logs.mrq_workers.ensure_index([("datereported", 1)], background=True, expireAfterSeconds=3600)

    self.mongodb_jobs.mrq_jobs.ensure_index([("status", 1)], background=True)
    self.mongodb_jobs.mrq_jobs.ensure_index([("path", 1)], background=True)
    self.mongodb_jobs.mrq_jobs.ensure_index([("worker", 1)], background=True)
    self.mongodb_jobs.mrq_jobs.ensure_index([("queue", 1)], background=True)
    self.mongodb_jobs.mrq_jobs.ensure_index([("dateexpires", 1)], sparse=True, background=True, expireAfterSeconds=0)

  def make_name(self):
    """ Generate a human-readable name for this worker. """
    return "%s.%s" % (socket.gethostname().split(".")[0], os.getpid())

  def greenlet_scheduler(self):

    from .scheduler import Scheduler
    scheduler = Scheduler(self.mongodb_jobs.mrq_scheduled_jobs)

    scheduler.sync_tasks(self.config.get("scheduler_tasks") or [])

    while True:
      scheduler.check()
      time.sleep(int(self.config["scheduler_interval"]))

  def greenlet_monitoring(self):
    """ This greenlet always runs in background to update current status in MongoDB every 10 seconds.

    Caution: it might get delayed when doing long blocking operations. Should we do this in a thread instead?
     """

    while True:

      # print "Monitoring..."

      self.report_worker()
      self.flush_logs(w=0)
      time.sleep(int(self.config["report_interval"]))

  def get_worker_report(self):

    greenlets = []

    for greenlet in self.gevent_pool:
      g = {}
      short_stack = []
      stack = traceback.format_stack(greenlet.gr_frame)
      for s in stack[1:]:
        if "/gevent/hub.py" in s:
          break
        short_stack.append(s)
      g["stack"] = short_stack

      job = get_current_job(id(greenlet))
      if job:
        if job.data:
          g["path"] = job.data["path"]
        g["datestarted"] = job.datestarted
        g["id"] = job.id
        g["time"] = getattr(greenlet, "_trace_time", 0)
        g["switches"] = getattr(greenlet, "_trace_switches", None)
      greenlets.append(g)

    cpu = self.process.get_cpu_times()

    # Avoid sharing passwords or sensitive config!
    whitelisted_config = [
      "max_jobs",
      "gevent",
      "processes",
      "queues",
      "scheduler",
      "name"
    ]

    return {
      "status": self.status,
      "config": {k: v for k, v in self.config.iteritems() if k in whitelisted_config},
      "done_jobs": self.done_jobs,
      "datestarted": self.datestarted,
      "datereported": datetime.datetime.utcnow(),
      "name": self.name,
      "process": {
        "pid": self.process.pid,
        "cpu": {
          "user": cpu.user,
          "system": cpu.system,
          "percent": self.process.get_cpu_percent(0)
        },
        "mem": {
          "rss": self.process.get_memory_info().rss
        }
        # https://code.google.com/p/psutil/wiki/Documentation
        # get_open_files
        # get_connections
        # get_num_ctx_switches
        # get_num_fds
        # get_io_counters
        # get_nice
      },
      "jobs": greenlets
    }

  def report_worker(self, w=0):

    try:
      self.mongodb_logs.mrq_workers.update({
        "_id": ObjectId(self.id)
      }, {"$set": self.get_worker_report()}, upsert=True, w=w)
    except pymongo.errors.AutoReconnect:
      self.log.debug("Worker report failed.")

  def greenlet_admin(self):
    """ This greenlet is used to get status information about the worker when --admin_port was given
    """

    if self.config["processes"] > 1:
      self.log.debug("Admin server disabled because of multiple processes.")
      return

    from flask import Flask
    from mrq.dashboard.utils import jsonify
    app = Flask("admin")

    @app.route('/')
    def route_index():
      report = self.get_worker_report()
      report.update({
        "_id": self.id
      })
      return jsonify(report)

    self.log.debug("Starting admin server on port %s" % self.config["admin_port"])
    try:
      server = WSGIServer(("0.0.0.0", self.config["admin_port"]), app)
      server.serve_forever()
    except Exception, e:
      self.log.debug("Error in admin server : %s" % e)

  def flush_logs(self, w=0):
    self.log_handler.flush(w=w)

  def dequeue_jobs(self, max_jobs=1):
    """ Fetch a maximum of max_jobs from this worker's queues. """

    self.log.debug("Fetching %s jobs from Redis" % max_jobs)

    jobs = []
    self.status = "idle"
    queue, job_id = self.redis.blpop(self.redis_queues, 0)
    self.status = "started"

    # From this point until job.fetch_and_start(), job is only local to this worker.
    # If we die here, job will be lost in redis without having been marked as "started".

    jobs.append(self.job_class(job_id, worker=self, queue=queue, start=True))

    # Bulk-fetch other jobs from that queue to fill the pool.
    # We take the chance that if there was one job on that queue, there should be more.
    if max_jobs > 1:

      with self.redis.pipeline(transaction=False) as pipe:
        for _ in range(max_jobs - 1):
          pipe.lpop(queue)
        job_ids = pipe.execute()

      jobs += [self.job_class(_job_id, worker=self, queue=queue, start=True)
               for _job_id in job_ids if _job_id]

    return jobs

  def work_loop(self):
    """Starts the work loop.

    """

    self.connect()

    self.status = "started"

    self.greenlets["monitoring"] = gevent.spawn(self.greenlet_monitoring)

    if self.config["scheduler"]:
      self.greenlets["scheduler"] = gevent.spawn(self.greenlet_scheduler)

    if self.config["admin_port"]:
      self.greenlets["admin"] = gevent.spawn(self.greenlet_admin)

    self.install_signal_handlers()

    try:

      while True:

        while True:
          free_pool_slots = self.gevent_pool.free_count()
          if free_pool_slots > 0:
            self.status = "started"
            break
          self.status = "idle"
          gevent.sleep(0.01)

        self.log.info('Listening on %s' % self.queues)

        jobs = self.dequeue_jobs(max_jobs=free_pool_slots)

        for job in jobs:

          # TODO investigate spawn_raw?
          self.gevent_pool.spawn(self.perform_job, job)

          self.done_jobs += 1

        if self.max_jobs and self.max_jobs >= self.done_jobs:
          self.log.info("Reached max_jobs=%s" % self.done_jobs)
          break

    except StopRequested:
      pass

    finally:

      try:

        self.log.debug("Joining the greenlet pool...")
        self.status = "stopping"

        self.gevent_pool.join(timeout=None, raise_error=False)
        self.log.debug("Joined.")

      except StopRequested:
        pass

      self.gevent_pool.kill(exception=JobInterrupt, block=True)

      for g in self.greenlets:
        g_time = getattr(self.greenlets[g], "_trace_time", 0)
        g_switches = getattr(self.greenlets[g], "_trace_switches", None)
        self.greenlets[g].kill(block=True)
        self.log.debug("Greenlet for %s killed (%0.5fs, %s switches)." % (g, g_time, g_switches))

      self.report_worker(w=1)
      self.flush_logs(w=1)

      if self.profiler:
        self.profiler.print_stats(sort="cumulative")

      g_time = getattr(self.greenlet, "_trace_time", 0)
      g_switches = getattr(self.greenlet, "_trace_switches", None)
      self.log.debug("Exiting main worker greenlet (%0.5fs, %s switches)." % (g_time, g_switches))

    return self.exitcode

  def perform_job(self, job):
    """ Wraps a job.perform() call with timeout logic and exception handlers.

        This is the first call happening inside the greenlet.
    """

    set_current_job(job)

    gevent_timeout = gevent.Timeout(job.timeout, JobTimeoutException(
      'Gevent Job exceeded maximum timeout value (%d seconds).' % job.timeout
    ))

    gevent_timeout.start()

    try:
      job.perform()

    except job.retry_on_exceptions:
      trace = traceback.format_exc()
      self.log.error("Caught exception => retry")
      self.log.error(trace)
      job.save_retry(sys.exc_info()[1], traceback=trace)

    except job.cancel_on_exceptions:
      trace = traceback.format_exc()
      self.log.error("Job cancelled")
      self.log.error(trace)
      job.save_status("cancel", traceback=trace)

    except JobTimeoutException:
      trace = traceback.format_exc()
      self.log.error("Job timeouted after %s seconds" % job.timeout)
      self.log.error(trace)
      job.save_status("timeout", traceback=trace)

    except JobInterrupt:
      trace = traceback.format_exc()
      self.log.error(trace)
      job.save_status("interrupt", traceback=trace)

    except Exception:
      trace = traceback.format_exc()
      self.log.error(trace)
      job.save_status("failed", traceback=trace)

    finally:
      gevent_timeout.cancel()
      set_current_job(None)

  def shutdown_graceful(self):
    """ Graceful shutdown: waits for all the jobs to finish. """

    # This is in the 'exitcodes' list in supervisord so processes
    # exiting gracefully won't be restarted.
    self.exitcode = 2

    self.log.info("Graceful shutdown...")
    raise StopRequested()

  def shutdown_now(self):
    """ Forced shutdown: interrupts all the jobs. """

    # This is in the 'exitcodes' list in supervisord so processes
    # exiting gracefully won't be restarted.
    self.exitcode = 3

    self.log.info("Forced shutdown...")
    self.status = "killing"

    self.gevent_pool.kill(exception=JobInterrupt, block=False)

    raise StopRequested()

  def install_signal_handlers(self):
    """ Handle events like Ctrl-C from the command line. """

    self.graceful_stop = False

    def request_shutdown_now():
      self.shutdown_now()

    def request_shutdown_graceful():

      # Second time CTRL-C, shutdown now
      if self.graceful_stop:
        request_shutdown_now()
      else:
        self.graceful_stop = True
        self.shutdown_graceful()

    # First time CTRL-C, try to shutdown gracefully
    gevent.signal(signal.SIGINT, request_shutdown_graceful)

    # User (or Heroku) requests a stop now, just mark tasks as interrupted.
    gevent.signal(signal.SIGTERM, request_shutdown_now)

