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

# -*- encoding: utf-8 -*-

"""contextimpl.py: ContextImpl class that implements the Context interface
"""

import re
import time
import os
import json
import log

import pulsar
import util

from prometheus_client import Summary
from function_stats import Stats
from functools import partial

Log = log.Log

class ContextImpl(pulsar.Context):

  # add label to indicate user metric
  user_metrics_label_names = Stats.metrics_label_names + ["metric"]

  def __init__(self, instance_config, logger, pulsar_client, user_code, consumers,
               secrets_provider, metrics_labels, state_context, stats):
    self.instance_config = instance_config
    self.log = logger
    self.pulsar_client = pulsar_client
    self.user_code_dir = os.path.dirname(user_code)
    self.consumers = consumers
    self.secrets_provider = secrets_provider
    self.state_context = state_context
    self.publish_producers = {}
    self.publish_serializers = {}
    self.message = None
    self.current_start_time = None
    self.user_config = json.loads(instance_config.function_details.userConfig) \
      if instance_config.function_details.userConfig \
      else []
    self.secrets_map = json.loads(instance_config.function_details.secretsMap) \
      if instance_config.function_details.secretsMap \
      else {}

    self.metrics_labels = metrics_labels
    self.user_metrics_map = dict()
    self.user_metrics_summary = Summary("pulsar_function_user_metric",
                                    'Pulsar Function user defined metric',
                                        ContextImpl.user_metrics_label_names)
    self.stats = stats

  # Called on a per message basis to set the context for the current message
  def set_current_message_context(self, message, topic):
    self.message = message
    self.current_start_time = time.time()

  def get_message_id(self):
    return self.message.message_id()

  def get_message_key(self):
    return self.message.partition_key()

  def get_message_eventtime(self):
    return self.message.event_timestamp()

  def get_message_properties(self):
    return self.message.properties()

  def get_current_message_topic_name(self):
    return self.message.topic_name()

  def get_message_sequence_id(self):
    if not self.get_message_id():
      return None
    ledger_id = self.get_message_id().ledger_id()
    entry_id = self.get_message_id().entry_id()
    offset = (ledger_id << 28) | entry_id
    return offset

  def get_message_partition_index(self):
    if not self.get_message_id():
      return None
    return self.get_message_id().partition()

  def get_partition_key(self):
    return self.message.partition_key()

  def get_function_name(self):
    return self.instance_config.function_details.name

  def get_function_tenant(self):
    return self.instance_config.function_details.tenant

  def get_function_namespace(self):
    return self.instance_config.function_details.namespace

  def get_function_id(self):
    return self.instance_config.function_id

  def get_instance_id(self):
    return self.instance_config.instance_id

  def get_function_version(self):
    return self.instance_config.function_version

  def get_logger(self):
    return self.log

  def get_user_config_value(self, key):
    if key in self.user_config:
      return self.user_config[key]
    else:
      return None

  def get_user_config_map(self):
    return self.user_config

  def get_secret(self, secret_key):
    if not secret_key in self.secrets_map:
      return None
    return self.secrets_provider.provide_secret(secret_key, self.secrets_map[secret_key])

  def record_metric(self, metric_name, metric_value):
    if metric_name not in self.user_metrics_map:
      user_metrics_labels = self.metrics_labels + [metric_name]
      self.user_metrics_map[metric_name] = self.user_metrics_summary.labels(*user_metrics_labels)

    self.user_metrics_map[metric_name].observe(metric_value)

  def get_input_topics(self):
    return list(self.instance_config.function_details.source.inputSpecs.keys())

  def get_output_topic(self):
    return self.instance_config.function_details.sink.topic

  def get_output_serde_class_name(self):
    return self.instance_config.function_details.sink.serDeClassName

  def callback_wrapper(self, callback, topic, message_id, result, msg):
    if result != pulsar.Result.Ok:
      error_msg = "Failed to publish to topic [%s] with error [%s] with src message id [%s]" % (topic, result, message_id)
      Log.error(error_msg)
      self.stats.incr_total_sys_exceptions(Exception(error_msg))
    if callback:
      callback(result, msg)

  def publish(self, topic_name, message, serde_class_name="serde.IdentitySerDe", properties=None, compression_type=None, callback=None, message_conf=None):
    # Just make sure that user supplied values are properly typed
    topic_name = str(topic_name)
    serde_class_name = str(serde_class_name)
    pulsar_compression_type = pulsar._pulsar.CompressionType.NONE
    if compression_type is not None:
      pulsar_compression_type = compression_type
    if topic_name not in self.publish_producers:
      self.publish_producers[topic_name] = self.pulsar_client.create_producer(
        topic_name,
        block_if_queue_full=True,
        batching_enabled=True,
        batching_max_publish_delay_ms=10,
        compression_type=pulsar_compression_type,
        properties=util.get_properties(util.getFullyQualifiedFunctionName(
          self.instance_config.function_details.tenant,
          self.instance_config.function_details.namespace,
          self.instance_config.function_details.name),
          self.instance_config.instance_id)
      )

    if serde_class_name not in self.publish_serializers:
      serde_klass = util.import_class(self.user_code_dir, serde_class_name)
      self.publish_serializers[serde_class_name] = serde_klass()

    output_bytes = bytes(self.publish_serializers[serde_class_name].serialize(message))

    if properties:
      # The deprecated properties args was passed. Need to merge into message_conf
      if not message_conf:
        message_conf = {}
      message_conf['properties'] = properties

    if message_conf:
      self.publish_producers[topic_name].send_async(
        output_bytes, partial(self.callback_wrapper, callback, topic_name, self.get_message_id()), **message_conf)
    else:
      self.publish_producers[topic_name].send_async(
        output_bytes, partial(self.callback_wrapper, callback, topic_name, self.get_message_id()))

  def ack(self, msgid, topic):
    topic_consumer = None
    if topic in self.consumers:
      topic_consumer = self.consumers[topic]
    else:
      # if this topic is a partitioned topic
      m = re.search(r'(.+)-partition-(\d+)', topic)
      if not m:
        raise ValueError('Invalid topicname %s' % topic)
      elif m.group(1) in self.consumers:
        topic_consumer = self.consumers[m.group(1)]
      else:
        raise ValueError('Invalid topicname %s' % topic)
    topic_consumer.acknowledge(msgid)

  def get_and_reset_metrics(self):
    metrics = self.get_metrics()
    # TODO(sanjeev):- Make this thread safe
    self.reset_metrics()
    return metrics

  def reset_metrics(self):
    # TODO: Make it thread safe
    for user_metric in self.user_metrics_map.values():
      user_metric._sum.set(0.0)
      user_metric._count.set(0.0)

  def get_metrics(self):
    metrics_map = {}
    for metric_name, user_metric in self.user_metrics_map.items():
      metrics_map["%s%s_sum" % (Stats.USER_METRIC_PREFIX, metric_name)] = user_metric._sum.get()
      metrics_map["%s%s_count" % (Stats.USER_METRIC_PREFIX, metric_name)] = user_metric._count.get()

    return metrics_map

  def incr_counter(self, key, amount):
    return self.state_context.incr(key, amount)

  def get_counter(self, key):
    return self.state_context.get_amount(key)

  def del_counter(self, key):
    return self.state_context.delete_key(key)

  def put_state(self, key, value):
    return self.state_context.put(key, value)

  def get_state(self, key):
    return self.state_context.get_value(key)

  def get_pulsar_client(self):
    return self.pulsar_client
