#!/usr/bin/env python
# -*- encoding: utf-8 -*-

# Copyright 2016 Twitter. All rights reserved.
#
# 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.
'''metrics_helper: helper classes for managing common metrics'''

from heron.common.src.python.utils.log import Log
import heron.instance.src.python.utils.system_constants as constants
from heron.instance.src.python.utils import system_config

from heron.proto import metrics_pb2

from heronpy.api.metrics import (CountMetric, MultiCountMetric, MeanReducedMetric,
                                 ReducedMetric, MultiMeanReducedMetric, MultiReducedMetric)

class BaseMetricsHelper(object):
  """Helper class for metrics management

  It registers metrics to the metrics collector and provides methods for
  updating metrics. This supports updating of: CountMetric, MultiCountMetric,
  ReducedMetric and MultiReducedMetric class.
  """
  def __init__(self, metrics):
    self.metrics = metrics

  def register_metrics(self, metrics_collector, interval):
    """Registers its metrics to a given metrics collector with a given interval"""
    for field, metrics in self.metrics.items():
      metrics_collector.register_metric(field, metrics, interval)

  def update_count(self, name, incr_by=1, key=None):
    """Update the value of CountMetric or MultiCountMetric

    :type name: str
    :param name: name of the registered metric to be updated.
    :type incr_by: int
    :param incr_by: specifies how much to increment. Default is 1.
    :type key: str or None
    :param key: specifies a key for MultiCountMetric. Needs to be `None` for updating CountMetric.
    """
    if name not in self.metrics:
      Log.error("In update_count(): %s is not registered in the metric", name)

    if key is None and isinstance(self.metrics[name], CountMetric):
      self.metrics[name].incr(incr_by)
    elif key is not None and isinstance(self.metrics[name], MultiCountMetric):
      self.metrics[name].incr(key, incr_by)
    else:
      Log.error("In update_count(): %s is registered but not supported with this method", name)

  def update_reduced_metric(self, name, value, key=None):
    """Update the value of ReducedMetric or MultiReducedMetric

    :type name: str
    :param name: name of the registered metric to be updated.
    :param value: specifies a value to be reduced.
    :type key: str or None
    :param key: specifies a key for MultiReducedMetric. Needs to be `None` for updating
                ReducedMetric.
    """
    if name not in self.metrics:
      Log.error("In update_reduced_metric(): %s is not registered in the metric", name)

    if key is None and isinstance(self.metrics[name], ReducedMetric):
      self.metrics[name].update(value)
    elif key is not None and isinstance(self.metrics[name], MultiReducedMetric):
      self.metrics[name].update(key, value)
    else:
      Log.error("In update_count(): %s is registered but not supported with this method", name)


class GatewayMetrics(BaseMetricsHelper):
  """Metrics helper class for Gateway metric"""
  RECEIVED_PKT_SIZE = '__gateway-received-packets-size'
  SENT_PKT_SIZE = '__gateway-sent-packets-size'
  RECEIVED_PKT_COUNT = '__gateway-received-packets-count'
  SENT_PKT_COUNT = '__gateway-sent-packets-count'

  SENT_METRICS_SIZE = '__gateway-sent-metrics-size'
  SENT_METRICS_PKT_COUNT = '__gateway-sent-metrics-packets-count'
  SENT_METRICS_COUNT = '__gateway-sent-metrics-count'
  SENT_EXCEPTION_COUNT = '__gateway-sent-exceptions-count'

  IN_STREAM_QUEUE_SIZE = '__gateway-in-stream-queue-size'
  OUT_STREAM_QUEUE_SIZE = '__gateway-out-stream-queue-size'

  IN_STREAM_QUEUE_EXPECTED_CAPACITY = '__gateway-in-stream-queue-expected-capacity'
  OUT_STREAM_QUEUE_EXPECTED_CAPACITY = '__gateway-out-stream-queue-expected-capacity'

  IN_QUEUE_FULL_COUNT = '__gateway-in-queue-full-count'

  metrics = {RECEIVED_PKT_SIZE: CountMetric(),
             SENT_PKT_SIZE: CountMetric(),
             RECEIVED_PKT_COUNT: CountMetric(),
             SENT_PKT_COUNT: CountMetric(),

             SENT_METRICS_SIZE: CountMetric(),
             SENT_METRICS_PKT_COUNT: CountMetric(),
             SENT_METRICS_COUNT: CountMetric(),
             SENT_EXCEPTION_COUNT: CountMetric(),

             IN_STREAM_QUEUE_SIZE: MeanReducedMetric(),
             OUT_STREAM_QUEUE_SIZE: MeanReducedMetric(),
             IN_STREAM_QUEUE_EXPECTED_CAPACITY: MeanReducedMetric(),
             OUT_STREAM_QUEUE_EXPECTED_CAPACITY: MeanReducedMetric()}

  def __init__(self, metrics_collector):
    sys_config = system_config.get_sys_config()
    super(GatewayMetrics, self).__init__(self.metrics)
    interval = float(sys_config[constants.HERON_METRICS_EXPORT_INTERVAL_SEC])
    self.register_metrics(metrics_collector, interval)

  def update_received_packet(self, received_pkt_size_bytes):
    """Update received packet metrics"""
    self.update_count(self.RECEIVED_PKT_COUNT)
    self.update_count(self.RECEIVED_PKT_SIZE, incr_by=received_pkt_size_bytes)

  def update_sent_packet(self, sent_pkt_size_bytes):
    """Update sent packet metrics"""
    self.update_count(self.SENT_PKT_COUNT)
    self.update_count(self.SENT_PKT_SIZE, incr_by=sent_pkt_size_bytes)

  def update_sent_metrics_size(self, size):
    self.update_count(self.SENT_METRICS_SIZE, size)

  def update_sent_metrics(self, metrics_count, exceptions_count):
    self.update_count(self.SENT_METRICS_PKT_COUNT)
    self.update_count(self.SENT_METRICS_COUNT, metrics_count)
    self.update_count(self.SENT_EXCEPTION_COUNT, exceptions_count)

  def update_in_out_stream_metrics(self, in_size, out_size, in_expect_size, out_expect_size):
    self.update_reduced_metric(self.IN_STREAM_QUEUE_SIZE, in_size)
    self.update_reduced_metric(self.OUT_STREAM_QUEUE_SIZE, out_size)
    self.update_reduced_metric(self.IN_STREAM_QUEUE_EXPECTED_CAPACITY, in_expect_size)
    self.update_reduced_metric(self.OUT_STREAM_QUEUE_EXPECTED_CAPACITY, out_expect_size)


class ComponentMetrics(BaseMetricsHelper):
  """Metrics to be collected for both Bolt and Spout"""
  FAIL_LATENCY = "__fail-latency"
  FAIL_COUNT = "__fail-count"
  EMIT_COUNT = "__emit-count"
  TUPLE_SERIALIZATION_TIME_NS = "__tuple-serialization-time-ns"
  OUT_QUEUE_FULL_COUNT = "__out-queue-full-count"

  component_metrics = {FAIL_LATENCY: MultiMeanReducedMetric(),
                       FAIL_COUNT: MultiCountMetric(),
                       EMIT_COUNT: MultiCountMetric(),
                       TUPLE_SERIALIZATION_TIME_NS: MultiCountMetric(),
                       OUT_QUEUE_FULL_COUNT: CountMetric()}

  def __init__(self, additional_metrics):
    metrics = self.component_metrics
    metrics.update(additional_metrics)
    super(ComponentMetrics, self).__init__(metrics)

  # pyli
nt: disable=arguments-differ
  def register_metrics(self, context):
    """Registers metrics to context

    :param context: Topology Context
    """
    sys_config = system_config.get_sys_config()
    interval = float(sys_config[constants.HERON_METRICS_EXPORT_INTERVAL_SEC])
    collector = context.get_metrics_collector()
    super(ComponentMetrics, self).register_metrics(collector, interval)

  def update_out_queue_full_count(self):
    """Apply update to the out-queu
e full count"""
    self.update_count(self.OUT_QUEUE_FULL_COUNT)

  def update_emit_count(self, stream_id):
    """Apply update to emit count"""
    self.update_count(self.EMIT_COUNT, key=stream_id)

  def serialize_data_tuple(self, stream_id, latency_in_ns):
    """Apply update to serialization metrics"""
    self.update_count(self.TUPLE_SERIALIZATION_TIME_NS, incr_by=latency_in_ns, key=stream_id)

class SpoutMetrics(ComponentMetrics):
  """Metrics helper class for Spout"""
  ACK_COUNT = "__ack-