#
# 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.
#

# pytype: skip-file

from typing import Optional

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):
    # type: () -> ScopedState

    """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,  # type: counters.CounterName
                    name_context,  # type: common.NameContext
                    output_counter,
                    metrics_container=None):
    # type: (...) -> ScopedState
    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):
    metrics_container = self.current_state().metrics_container
    if metrics_container is not None:
      metrics_container.get_metric_cell(typed_metric_name).update(value)

  def _enter_state(self, state):
    # type: (ScopedState) -> None
    self.state_transition_count += 1
    self._state_stack.append(state)

  def _exit_state(self):
    # type: () -> None
    self.state_transition_count += 1
    self._state_stack.pop()

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

  def stop(self):
    # type: () -> None
    pass

  def reset(self):
    # type: () -> None
    pass


class ScopedState(object):

  def __init__(self,
               sampler,  # type: StateSampler
               name,  # type: counters.CounterName
               step_name_context,  # type: Optional[common.NameContext]
               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):
    # type: () -> float
    return 1e-9 * self.nsecs

  def sampled_msecs_int(self):
    # type: () -> int
    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()