# Copyright 2011 The scales Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Useful utility functions and objects."""

from six.moves.queue import Queue
from math import exp

import logging
import random
import socket
import threading
import time


def lookup(source, keys, fallback = None):
  """Traverses the source, looking up each key.  Returns None if can't find anything instead of raising an exception."""
  try:
    for key in keys:
      source = source[key]
    return source
  except (KeyError, AttributeError, TypeError):
    return fallback



class GraphiteReporter(threading.Thread):
  """A graphite reporter thread."""

  def __init__(self, host, port, maxQueueSize=10000):
    """Connect to a Graphite server on host:port."""
    threading.Thread.__init__(self)

    self.host, self.port = host, port
    self.sock = None
    self.queue = Queue()
    self.maxQueueSize = maxQueueSize
    self.daemon = True


  def run(self):
    """Run the thread."""
    while True:
      try:
        try:
          name, value, valueType, stamp = self.queue.get()
        except TypeError:
          break
        self.log(name, value, valueType, stamp)
      finally:
        self.queue.task_done()


  def connect(self):
    """Connects to the Graphite server if not already connected."""
    if self.sock is not None:
      return
    backoff = 0.01
    while True:
      try:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(5)
        sock.connect((self.host, self.port))
        self.sock = sock
        return
      except socket.error:
        time.sleep(random.uniform(0, 2.0*backoff))
        backoff = min(backoff*2.0, 5.0)


  def disconnect(self):
    """Disconnect from the Graphite server if connected."""
    if self.sock is not None:
      try:
        self.sock.close()
      except socket.error:
        pass
      finally:
        self.sock = None


  def _sendMsg(self, msg):
    """Send a line to graphite. Retry with exponential backoff."""
    if not self.sock:
      self.connect()
    backoff = 0.001
    while True:
      try:
        self.sock.sendall(msg)
        break
      except socket.error:
        logging.warning('Graphite connection error', exc_info = True)
        self.disconnect()
        time.sleep(random.uniform(0, 2.0*backoff))
        backoff = min(backoff*2.0, 5.0)
        self.connect()


  def _sanitizeName(self, name):
    """Sanitize a metric name."""
    return name.replace(' ', '-')


  def log(self, name, value, valueType=None, stamp=None):
    """Log a named numeric value. The value type may be 'value',
    'count', or None."""
    if type(value) == float:
      form = "%s%s %2.2f %d\n"
    else:
      form = "%s%s %s %d\n"

    if valueType is not None and len(valueType) > 0 and valueType[0] != '.':
      valueType = '.' + valueType

    if not stamp:
      stamp = time.time()

    self._sendMsg(form % (self._sanitizeName(name), valueType or '', value, stamp))


  def enqueue(self, name, value, valueType=None, stamp=None):
    """Enqueue a call to log."""
    # If queue is too large, refuse to log.
    if self.maxQueueSize and self.queue.qsize() > self.maxQueueSize:
      return
    # Stick arguments into the queue
    self.queue.put((name, value, valueType, stamp))


  def flush(self):
    """Block until all stats have been sent to Graphite."""
    self.queue.join()


  def shutdown(self):
    """Shut down the background thread."""
    self.queue.put(None)
    self.flush()



class AtomicValue(object):
  """Stores a value, atomically."""

  def __init__(self, val):
    self.lock = threading.RLock()
    self.value = val


  def update(self, function):
    """Atomically apply function to the value, and return the old and new values."""
    with self.lock:
      oldValue = self.value
      self.value = function(oldValue)
      return oldValue, self.value


  def getAndSet(self, newVal):
    """Sets a new value while returning the old value"""
    return self.update(lambda _: newVal)[0]


  def addAndGet(self, val):
    """Adds val to the value and returns the result"""
    return self.update(lambda x: x + val)[1]



class EWMA(object):
  """
  An exponentially-weighted moving average.

  Ported from Yammer metrics.
  """

  M1_ALPHA = 1 - exp(-5 / 60.0)
  M5_ALPHA = 1 - exp(-5 / 60.0 / 5)
  M15_ALPHA = 1 - exp(-5 / 60.0 / 15)

  TICK_RATE = 5 # Every 5 seconds


  @classmethod
  def oneMinute(cls):
    """Creates an EWMA configured for a 1 min decay with a 5s tick"""
    return EWMA(cls.M1_ALPHA, 5)


  @classmethod
  def fiveMinute(cls):
    """Creates an EWMA configured for a 5 min decay with a 5s tick"""
    return EWMA(cls.M5_ALPHA, 5)


  @classmethod
  def fifteenMinute(cls):
    """Creates an EWMA configured for a 15 min decay with a 5s tick"""
    return EWMA(cls.M15_ALPHA, 5)


  def __init__(self, alpha, interval):
    self.alpha = alpha
    self.interval = interval
    self.rate = 0
    self._uncounted = AtomicValue(0)
    self._initialized = False


  def update(self, val):
    """Adds this value to the count to be averaged"""
    self._uncounted.addAndGet(val)


  def tick(self):
    """Updates rates and decays"""
    count = self._uncounted.getAndSet(0)
    instantRate = float(count) / self.interval

    if self._initialized:
      self.rate += (self.alpha * (instantRate - self.rate))
    else:
      self.rate = instantRate
      self._initialized = True
