"""Module for recording metrics and uploading to Google Cloud Logging.

There are 3 methods of logging:
1) Stackdriver / Cloud Logging API
2) Local Logging
3) Fluentd Agent

Setup:
  1) Stackdriver / Cloud Logging API, during application setup you should call:

    project_name = 'google.com:cast-cq-staging'
    logger_name = 'cast_cq.master'
    cloud_logging.ConfigureLogger(project_name, logger_name)

  2) Local Logging, this is used by default if you do not configure the cloud
  logger or enable the fluentd logger

  3) For Fluentd Agent, during application setup you should call:

    namespace = "buildbot"
    host = "127.0.0.0"
    port = 24224
    cloud_logging.EnableFluentdLogger(namespace, host, port)

  There are few caveats when using the fluentd agent:
  - The Google Cloud project will be inferred by the agent (either by explicitly
    setting the project in the configuration or implicitly by the IAM Role)
    that's how it determines which project to upload the logs to.
  - Since we're using a version of the Google Cloud agent, it assumes that the
    host it's running on is a GCE instance (which it may not). The consequence
    of this is that the logs will be grouped into a GCE resource log like
    "GCE VM Instance, <id>" and NOT "Global"
  - Severity is ignored when using the agent
  - Additional labels are also not passed to the agent

Usage:
  Then, throughout your application, you can use:

  import cloud_logging

  t = cloud_logging.Timer('test_timer').Start()
  # do something interesting...
  t.Stop()

  with cloud_logging.Timer('test_timer') as t:
    # do something interesting...
    t.AddAttribute('a', 'b')

  ev = cloud_logging.LogEvent('test_msg', value='Something to log').Log()
  cloud_logging.Info('info_message', value='Something to log').Log()
  cloud_logging.Error('error_message', value='Something to log').Log()
"""
import atexit
import logging
import os
import socket
import threading
import time
import traceback

from fluent import event as fluentd_event
from fluent import sender as fluentd  # TODO(b/70284159): switch to asyncsender

from google.cloud import logging as gcl  # pylint: disable=ungrouped-imports
from google.gax import errors
from google.gax import grpc
from google.protobuf import json_format
import six

# Subset of supported severity levels:
# https://cloud.google.com/logging/docs/api/reference/rest/v2/LogEntry#LogSeverity
DEBUG = 'DEBUG'
INFO = 'INFO'
WARNING = 'WARNING'
ERROR = 'ERROR'

SEVERITY_LEVELS = [DEBUG, INFO, WARNING, ERROR]
# Global instances of configured cloud client/loggers
_GLOBAL_CLIENT = None
_GLOBAL_LOGGER = None
_GLOBAL_THREAD = None
_GLOBAL_LABELS = {}


class BaseError(Exception):
  """Base error."""


class ConfigurationError(BaseError):
  """Base error for invalid log configurations."""


class EventError(BaseError):
  """Base error for Metrics."""


class TimerError(BaseError):
  """Base error for Timers."""


class StrictAttributesError(BaseError):
  """Base Error for StrictAttributes decorator."""


class LoggerError(BaseError):
  """Base error for errors raised in logger."""

  def __init__(self, original_error):
    super(LoggerError, self).__init__(self)
    self.original_error = original_error
    self.original_trace = traceback.format_exc()

  def __str__(self):
    return self.original_trace


class LoggerGaxError(LoggerError):
  """Class for gax errors raised in logger."""


class LogEvent(object):
  """A high level LogEvent class used for logging events and data."""

  def __init__(self, name, attributes=None, value=None, severity=None):
    """Initialize a LogEvent.

    Args:
      name: The name of the event being timed.
      attributes: A dict representing attributes to be recorded with this timer.
      value: A value to include in the Metric.
      severity: A string representing the log severity. One of:
          'DEBUG', 'INFO', 'WARNING', 'ERROR' or None.
    """
    self._name = self._Sanitize(name)
    self._attributes = {}
    if not attributes:
      attributes = {}
    self.AddAttributes(attributes)
    self._host = socket.getfqdn()
    self._recorded_at = self._Now()
    self._logged = False
    self._value = self._Sanitize(value)
    self._type = 'general'
    self._severity = severity

  @property
  def type(self):
    """Returns the type of event."""
    return self._type

  @property
  def value(self):
    """Returns the value for this event."""
    return self._value

  @value.setter
  def value(self, value):
    """Sets the value for this event."""
    self._value = self._Sanitize(value)

  @property
  def severity(self):
    """Returns the severity of event."""
    return self._severity

  @severity.setter
  def severity(self, severity):
    """Sets the severity of this event.

    Only valid severity values will be accepted, all others will fail silently.

    Args:
      severity: A string representing the log severity. One of:
          'DEBUG', 'INFO', 'WARNING', 'ERROR' or None.
    """
    self._severity = None
    if severity in SEVERITY_LEVELS:
      self._severity = severity

  @property
  def recorded_at(self):
    """Get the timestamp for when this event was recorded."""
    return self._recorded_at

  @property
  def name(self):
    return self._name

  @property
  def host(self):
    return self._host

  @property
  def attributes(self):
    return self._attributes

  def _Sanitize(self, value):
    """Sanitizes arguments so they can be uploaded to GCL.

    If it is a supported primitive that will be returned, otherwise its str
    representation will be used.

    This will also iterate through lists & dicts and sanitize their values.

    Args:
      value: The value to sanitize.

    Returns:
      The sanitized version of value.
    """
    if isinstance(value, (six.integer_types, float, bool)):
      return value
    if isinstance(value, list):
      return [self._Sanitize(v) for v in value]
    if isinstance(value, dict):
      return {self._Sanitize(key): self._Sanitize(v)
              for key, v in six.iteritems(value)}
    try:
      return str(value)
    except Exception as e:  # pylint: disable=broad-except
      return 'Could not sanitize object of type %s: %s' % (type(value), e)

  def AddAttributes(self, attributes):
    """Add more attributes to the event.

    Args:
      attributes: A dict of attributes to add to the current attribute list.
                  They will override other attributes with the same keys.

    Returns:
      self.
    """
    if not attributes:
      return self
    for key, value in six.iteritems(attributes):
      self.AddAttribute(key, value)
    return self

  def AddAttribute(self, name, value):
    """Add a single attribute to the event.

    Args:
      name: The key to be used in the attribute dict.
      value: The value to put for the given key. Note, if the key already
             existed its value will be overridden.

    Returns:
      self.
    """
    self._attributes[self._Sanitize(name)] = self._Sanitize(value)
    return self

  def AsDict(self):
    """Returns a dictionary representation of this event."""
    result = {
        'name': self.name,
        'recorded_at_utc_usec': self.recorded_at,
        'value': self.value,
        'type': self.type,
        'host': self.host,
        'attributes': self.attributes
    }
    return result

  def Valid(self):
    """Returns if this event is valid."""
    return True

  def Log(self):
    """Log this event.

    Raises:
      EventError: If the event is invalid.
    """
    if not self.Valid():
      raise EventError('Event not valid.')
    if not self._logged:
      _SendEvent(self, severity=self._severity)
      self._logged = True

  def _Now(self):
    """Returns the UTC usec timestamp for now."""
    return time.time() * 1000 * 1000


class Timer(LogEvent):
  """A simpler timer object to record & log the duration of something."""

  def __init__(self, name, attributes=None, severity=None):
    """Initialize a timer.

    Remember to call Start() to begin timing.

    Args:
      name: The name of the event being timed.
      attributes: A dict representing attributes to be recorded with this timer.
      severity: A string representing the log severity. One of:
          'DEBUG', 'INFO', 'WARNING', 'ERROR' or None.
    """
    super(Timer, self).__init__(name, attributes=attributes, severity=severity)
    self._started_at = None
    self._finished_at = None
    self._type = 'timer'

  def Start(self):
    """Starts the timer.

    Returns:
      self.
    """
    self._started_at = self._Now()
    return self

  def Stop(self):
    """Stops the timer and records the result.

    Note that the timer must have been started.

    Returns:
      self.

    Raises:
      TimerError: If the timer has not yet been started.
    """
    if not self.started:
      raise TimerError('Timer not yet started.')
    self._finished_at = self._Now()
    self.Log()
    return self

  @property
  def started(self):
    """Returns boolean if this timer has been started."""
    return self._started_at is not None

  @property
  def finished(self):
    """Returns boolean if this timer has been completed."""
    return self._finished_at is not None

  def Time(self):
    """Returns the value recorded by this timer.

    Returns:
      The timed value.

    Raises:
      TimerError: If the timer has not been started or completed.
    """
    if not self.started:
      raise TimerError('Timer not yet started.')
    if not self.finished:
      raise TimerError('Timer not yet finished.')
    return self._finished_at - self._started_at

  @property
  def value(self):
    return self.Time()

  @property
  def recorded_at(self):
    return self._finished_at

  def Valid(self):
    try:
      self.Time()
      return True
    except TimerError:
      return False

  def __enter__(self):
    self.Start()
    return self

  def __exit__(self, *args):
    self.Stop()


class StrictAttributes(object):
  """Decorator for LogEvents that require a fixed set of attributes.

  This decorator should be used when collecting metrics to be uploaded to
  BigQuery. All attributes to be collected must be included in valid_attributes,
  mapped to their default values. Any valid_attributes not added before Log()
  will be included with the specified default value.

  Sample Usage:

  VALID_ATTRIBUTES = {
      'status': StrictAttributes.DEFAULT_STRING,
      'count': StrictAttributes.DEFAULT_INT }

  @StrictAttributes(VALID_ATTRIBUTES)
  class MyLogger(Timer):
    ...
  """
  # TODO(adrexler): Migrate valid_attributes to protobuf.

  # Default values for all accepted attribute datatypes.
  DEFAULT_STRING = ''
  DEFAULT_BOOL = False
  DEFAULT_LIST = []
  DEFAULT_DICT = {}
  DEFAULT_INT = -1
  DEFAULT_FLOAT = float(DEFAULT_INT)
  try:
    DEFAULT_LONG = long(DEFAULT_INT)
  except NameError:
    DEFAULT_LONG = int(DEFAULT_INT)  # pylint: disable=redefined-variable-type

  def __init__(self, valid_attributes):
    """Constructor.

    Args:
      valid_attributes: All acceptable attributes, mapped to default values.
    """
    self._valid_attributes = valid_attributes

  def __call__(self, cls, *args, **kwargs):
    """Returns MetricsCollector class, extending cls."""
    if not issubclass(cls, LogEvent):
      raise StrictAttributesError('%s does not extend LogEvent!', cls)
    valid_attributes = self._valid_attributes

    class MetricsCollector(cls):
      """Extension of LogEvent child cls, with strict attributes."""

      def __init__(self, *args, **kwargs):
        self._valid_attributes = valid_attributes
        super(MetricsCollector, self).__init__(*args, **kwargs)

      def AddAttribute(self, name, value):
        """Adds value for attribute after checking that attribute is valid."""
        if name not in self._valid_attributes:
          logging.error('%s is not a valid attribute!', name)
          return
        super(MetricsCollector, self).AddAttribute(name, value)

      def Log(self):
        """Adds specified default values for missing attributes and logs."""
        for attribute in self._valid_attributes:
          if attribute not in self._attributes:
            self.AddAttribute(attribute, self._valid_attributes[attribute])
        super(MetricsCollector, self).Log()
    return MetricsCollector


class BatchThread(threading.Thread):
  """A simple thread worker to call commit on the Batch Logger."""

  def __init__(self, duration, err_limit=3, sleep=2.0):
    super(BatchThread, self).__init__()
    self._duration = duration
    self._sleep = sleep
    self._stop = threading.Event()
    self._err_limit = err_limit

  def run(self):
    while not self._stop.isSet():
      time.sleep(self._duration)
      if not _CommitWithRetry(self._sleep, self._err_limit):
        logging.error('Terminated cloud logging due to non-transient error')
        self._stop.set()

  def join(self, timeout=None):
    self._stop.set()
    super(BatchThread, self).join(timeout)


def Debug(*args, **kwargs):
  return LogEvent(*args, severity=DEBUG, **kwargs)


def Info(*args, **kwargs):
  return LogEvent(*args, severity=INFO, **kwargs)


def Warn(*args, **kwargs):
  return LogEvent(*args, severity=WARNING, **kwargs)


def Error(*args, **kwargs):
  return LogEvent(*args, severity=ERROR, **kwargs)


def _SendEvent(event, severity=None):
  """Log a metric.

  Args:
    event: A LogEvent to log. Must have an AsDict method implemented.
    severity: A string representing the log severity. This will be ignored if
              fluentd is enabled.

  Raises:
    LoggerGaxError: If GaxError other then DEADLINE_EXCEEDED was
                    raised during logger call.
    LoggerError: If any error other then ParseError or GaxError was
                 raised during logger call.
  """
  if fluentd.get_global_sender():
    fluentd_event.Event(None, event.AsDict())
  elif _GLOBAL_LOGGER:
    try:
      _GLOBAL_LOGGER.log_struct(
          event.AsDict(), labels=_GLOBAL_LABELS, severity=severity)
    except errors.GaxError as e:
      if grpc.exc_to_code(e.cause) != grpc.StatusCode.DEADLINE_EXCEEDED:
        raise LoggerGaxError(e)
      logging.exception('Gax RPC Deadline exceeded while push to log event %s'
                        ' with labels %s, not expecting event loss in cloud',
                        event.AsDict(), _GLOBAL_LABELS)
    except json_format.ParseError as e:
      logging.exception('Unable to log event %s with labels %s'
                        ' event will not be reported to cloud',
                        event.AsDict(), _GLOBAL_LABELS)
    except Exception as e:  # pylint: disable=broad-except
      raise LoggerError(e)
  else:
    logging.info(event.AsDict())


def _CommitWithRetry(sleep=2.0, retries=3):
  """Commiting event queue. Dispatching errors and recover if possible.

  The only transient error we are aware of is a Gax RetryError.
  In 99% cases we succeed from 3rd commit and will not succeed after that
  without clearing commit queue.
  For other then Gax Retry Error we will not recover, so returning False

  Args:
    sleep: Sleep period between retries in case of error.
    retries: Maximum number of commits before clearing queue.
  Return:
    Boolean if the batch process should continue
  """
  while retries > 0:
    try:
      _GLOBAL_LOGGER.commit()
      return True
    except errors.RetryError as e:
      retries -= 1
      queue_size = len(_GLOBAL_LOGGER.entries[:])
      err = str(e)
      logging.exception('Gax Retry Error while commit log buffer,'
                        ' %d retries left,%s %d events in queue', retries,
                        '' if retries else ' dropping', queue_size)
      if retries > 0:
        Warn('RetryError', value=queue_size, attributes={
             'retries_left': retries, 'error': err,
             'dropping_queue': False}).Log()
        time.sleep(sleep)
      else:
        del _GLOBAL_LOGGER.entries[:]
        Warn('RetryError', value=queue_size, attributes={
             'retries_left': 0, 'error': err,
             'dropping_queue': True}).Log()
    except errors.GaxError as e:
      logging.exception('Gax Error while commit log buffer, stopping batch')
      return False
    except Exception as e:  # pylint: disable=broad-except
      logging.exception('Unexpected Error while commit log buffer,'
                        ' stopping batch')
      return False
  return True

def ConfigureLogger(project_name, logger_name, labels=None, batch_time=None,
                    batch_err_limit=3):
  """Configure a Cloud Logger.

  Does NOT enable logging to a fluentd agent. See EnableFluentdLogger()

  Args:
    project_name: The name of the project to log to.
    logger_name: The name of the logger.
    labels: An optional dict of labels to apply to each message.
    batch_time: use a batch logger to send logs every n seconds.
    batch_err_limit: max consecutive errors before batch logging is terminated.

  Raises:
    ConfigurationError: If no credentials were found.
  """
  if 'GOOGLE_APPLICATION_CREDENTIALS' not in os.environ:
    raise ConfigurationError(
        ('No credentials found. Make sure a json based key is available in the '
         'GOOGLE_APPLICATION_CREDENTIALS env var.')
    )

  global _GLOBAL_CLIENT
  global _GLOBAL_LOGGER
  global _GLOBAL_THREAD
  _GLOBAL_CLIENT = gcl.Client(project=project_name)
  _GLOBAL_LOGGER = _GLOBAL_CLIENT.logger(logger_name)
  if batch_time:
    _GLOBAL_LOGGER = _GLOBAL_LOGGER.batch()
    _GLOBAL_THREAD = BatchThread(batch_time, batch_err_limit)
    _GLOBAL_THREAD.start()
  _GLOBAL_LABELS['host'] = socket.getfqdn()
  if labels:
    _GLOBAL_LABELS.update(labels)

def SafeConfigureLogger(*args, **kwargs):
  """Configures a Cloud Logger, if possible. Otherwise fails silently.

     Does NOT enable logging to a fluentd agent. See EnableFluentdLogger()
  """
  try:
    ConfigureLogger(*args, **kwargs)
  except ConfigurationError:
    logging.error('Unable to configure cloud logging: No GAE credentials '
                  'found. Logging to console only.')

def EnableFluentdLogger(namespace, host="127.0.0.1", port=24224):
  """Enable logging to a fluentd agent for logging to Google Cloud.

  This disables any other form of logging, including both the Cloud Logging API
  and local logging.

  Args:
    namespace: The namespace for logs in this application. For example
               "buildbot", which becomes referenced in Google Cloud as
               logName: "projects/<project>/logs/<namespace>"
    host: The hostname or IP the fluentd agent lives on.
    port: The port the fluentd agent is listening on.
  """
  if fluentd.get_global_sender() is None:
    fluentd.setup(namespace, host=host, port=port)

@atexit.register
def StopLogger():
  """
    Stop the batch logger thread and/or fluentd logger if started. No more
    logs will be sent.

    This will be executed on application termination
  """
  if _GLOBAL_THREAD:
    _GLOBAL_THREAD.join()

  if fluentd.get_global_sender():
    fluentd.close()