#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#

# This module is experimental. No backwards-compatibility guarantees.

from __future__ import absolute_import

from builtins import object

from apache_beam.runners import common
from apache_beam.utils import counters


class StateSampler(object):

  def __init__(self, sampling_period_ms):
    self._state_stack = [ScopedState(self,
                                     counters.CounterName('unknown'),
                                     None)]
    self.state_transition_count = 0
    self.time_since_transition = 0

  def current_state(self):
    """Returns the current execution state.

    This operation is not thread safe, and should only be called from the
    execution thread."""
    return self._state_stack[-1]

  def _scoped_state(self,
                    counter_name,
                    name_context,
                    output_counter,
                    metrics_container=None):
    assert isinstance(name_context, common.NameContext)
    return ScopedState(
        self, counter_name, name_context, output_counter, metrics_container)

  def update_metric(self, typed_metric_name, value):
    self.current_state().metrics_container.get_metric_cell(
        typed_metric_name).update(value)

  def _enter_state(self, state):
    self.state_transition_count += 1
    self._state_stack.append(state)

  def _exit_state(self):
    self.state_transition_count += 1
    self._state_stack.pop()

  def start(self):
    # Sampling not yet supported. Only state tracking at the moment.
    pass

  def stop(self):
    pass

  def reset(self):
    for state in self._states_by_name.values():
      state.nsecs = 0


class ScopedState(object):

  def __init__(self, sampler, name, step_name_context,
               counter=None, metrics_container=None):
    self.state_sampler = sampler
    self.name = name
    self.name_context = step_name_context
    self.counter = counter
    self.nsecs = 0
    self.metrics_container = metrics_container

  def sampled_seconds(self):
    return 1e-9 * self.nsecs

  def sampled_msecs_int(self):
    return int(1e-6 * self.nsecs)

  def __repr__(self):
    return "ScopedState[%s, %s]" % (self.name, self.nsecs)

  def __enter__(self):
    self.state_sampler._enter_state(self)

  def __exit__(self, exc_type, exc_value, traceback):
    self.state_sampler._exit_state()