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

"""A simple example demonstrating usage of the EnvoyRateLimiter in a Beam
pipeline.
"""

import argparse
import logging
import time

import apache_beam as beam
from apache_beam.io.components.rate_limiter import EnvoyRateLimiter
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.utils import shared


# [START RateLimiterSimplePython]
class SampleApiDoFn(beam.DoFn):
  """A DoFn that simulates calling an external API with rate limiting."""
  def __init__(self, rls_address, domain, descriptors):
    self.rls_address = rls_address
    self.domain = domain
    self.descriptors = descriptors
    self._shared = shared.Shared()
    self.rate_limiter = None

  def setup(self):
    # Initialize the rate limiter in setup()
    # We use shared.Shared() to ensure only one RateLimiter instance is created
    # per worker and shared across threads.
    def init_limiter():
      logging.info("Connecting to Envoy RLS at %s", self.rls_address)
      return EnvoyRateLimiter(
          service_address=self.rls_address,
          domain=self.domain,
          descriptors=self.descriptors,
          namespace='example_pipeline')

    self.rate_limiter = self._shared.acquire(init_limiter)

  def process(self, element):
    self.rate_limiter.allow()

    # Process the element mock API call
    logging.info("Processing element: %s", element)
    time.sleep(0.1)
    yield element


# [END RateLimiterSimplePython]


def parse_known_args(argv):
  """Parses args for the workflow."""
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--rls_address',
      default='localhost:8081',
      help='Address of the Envoy Rate Limit Service')
  return parser.parse_known_args(argv)


def run(argv=None):
  known_args, pipeline_args = parse_known_args(argv)
  pipeline_options = PipelineOptions(pipeline_args)

  with beam.Pipeline(options=pipeline_options) as p:
    _ = (
        p
        | 'Create' >> beam.Create(range(100))
        | 'RateLimit' >> beam.ParDo(
            SampleApiDoFn(
                rls_address=known_args.rls_address,
                domain="mongo_cps",
                descriptors=[{
                    "database": "users"
                }])))


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  run()