# 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.
#
# Basic object model of a Impala Services (impalad + statestored). Provides a way to
# programatically interact with the services and perform operations such as querying
# the debug webpage, getting metric values, or creating client connections.

import json
import logging
import re
import urllib
from time import sleep, time

from tests.common.impala_connection import create_connection, create_ldap_connection
from TCLIService import TCLIService
from thrift.transport.TSocket import TSocket
from thrift.transport.TTransport import TBufferedTransport
from thrift.protocol import TBinaryProtocol, TCompactProtocol, TProtocol
from thrift.TSerialization import deserialize
from RuntimeProfile.ttypes import TRuntimeProfileTree
import base64
import zlib

logging.basicConfig(level=logging.ERROR, format='%(threadName)s: %(message)s')
LOG = logging.getLogger('impala_service')
LOG.setLevel(level=logging.DEBUG)

# Base class for all Impala services
# TODO: Refactor the retry/timeout logic into a common place.
class BaseImpalaService(object):
  def __init__(self, hostname, webserver_port):
    self.hostname = hostname
    self.webserver_port = webserver_port

  def open_debug_webpage(self, page_name, timeout=10, interval=1):
    start_time = time()

    while (time() - start_time < timeout):
      try:
        return urllib.urlopen("http://%s:%d/%s" %
            (self.hostname, int(self.webserver_port), page_name))
      except Exception:
        LOG.info("Debug webpage not yet available.")
      sleep(interval)
    assert 0, 'Debug webpage did not become available in expected time.'

  def read_debug_webpage(self, page_name, timeout=10, interval=1):
    return self.open_debug_webpage(page_name, timeout=timeout, interval=interval).read()

  def get_thrift_profile(self, query_id, timeout=10, interval=1):
    """Returns thrift profile of the specified query ID, if available"""
    page_name = "query_profile_encoded?query_id=%s" % (query_id)
    try:
      response = self.open_debug_webpage(page_name, timeout=timeout, interval=interval)
      tbuf = response.read()
    except Exception as e:
      LOG.info("Thrift profile for query %s not yet available: %s", query_id, str(e))
      return None
    else:
      tree = TRuntimeProfileTree()
      try:
        deserialize(tree, zlib.decompress(base64.b64decode(tbuf)),
                  protocol_factory=TCompactProtocol.TCompactProtocolFactory())
        tree.validate()
        return tree
      except Exception as e:
        LOG.info("Exception while deserializing query profile of %s: %s", query_id,
                str(e));
        # We should assert that the response code is not 200 once
        # IMPALA-6332: Impala webserver should return HTTP error code for missing query
        # profiles, is fixed.
        if str(e) == 'Incorrect padding':
          assert "Could not obtain runtime profile" in tbuf, tbuf
        return None

  def get_debug_webpage_json(self, page_name):
    """Returns the json for the given Impala debug webpage, eg. '/queries'"""
    return json.loads(self.read_debug_webpage(page_name + "?json"))

  def get_metric_value(self, metric_name, default_value=None):
    """Returns the value of the the given metric name from the Impala debug webpage"""
    metrics = json.loads(self.read_debug_webpage('jsonmetrics?json'))
    return metrics.get(metric_name, default_value)

  def wait_for_metric_value(self, metric_name, expected_value, timeout=10, interval=1):
    start_time = time()
    while (time() - start_time < timeout):
      LOG.info("Getting metric: %s from %s:%s" %
          (metric_name, self.hostname, self.webserver_port))
      value = None
      try:
        value = self.get_metric_value(metric_name)
      except Exception, e:
        LOG.error(e)

      if value == expected_value:
        LOG.info("Metric '%s' has reached desired value: %s" % (metric_name, value))
        return value
      else:
        LOG.info("Waiting for metric value '%s'=%s. Current value: %s" %
            (metric_name, expected_value, value))
      LOG.info("Sleeping %ds before next retry." % interval)
      sleep(interval)
    assert 0, 'Metric value %s did not reach value %s in %ss\nDumping impalad debug ' \
              'pages:\nmemz: %s\nmetrics: %s\nqueries: %s\nthreadz: %s\nrpcz: %s' % \
              (metric_name, expected_value, timeout,
               json.dumps(self.read_debug_webpage('memz?json')),
               json.dumps(self.read_debug_webpage('metrics?json')),
               json.dumps(self.read_debug_webpage('queries?json')),
               json.dumps(self.read_debug_webpage('threadz?json')),
               json.dumps(self.read_debug_webpage('rpcz?json')))

# Allows for interacting with an Impalad instance to perform operations such as creating
# new connections or accessing the debug webpage.
class ImpaladService(BaseImpalaService):
  def __init__(self, hostname, webserver_port=25000, beeswax_port=21000, be_port=22000,
               hs2_port=21050):
    super(ImpaladService, self).__init__(hostname, webserver_port)
    self.beeswax_port = beeswax_port
    self.be_port = be_port
    self.hs2_port = hs2_port

  def get_num_known_live_backends(self, timeout=30, interval=1):
    LOG.info("Getting num_known_live_backends from %s:%s" %
        (self.hostname, self.webserver_port))
    result = json.loads(self.read_debug_webpage('backends?json', timeout, interval))
    num = len(result['backends'])
    return None if num is None else int(num)

  def get_in_flight_queries(self, timeout=30, interval=1):
    result = json.loads(self.read_debug_webpage('queries?json', timeout, interval))
    return result['in_flight_queries']

  def get_num_in_flight_queries(self, timeout=30, interval=1):
    LOG.info("Getting num_in_flight_queries from %s:%s" %
        (self.hostname, self.webserver_port))
    result = self.read_debug_webpage('inflight_query_ids?raw', timeout, interval)
    return None if result is None else len([l for l in result.split('\n') if l])

  def wait_for_num_in_flight_queries(self, expected_val, timeout=10):
    """Waits for the number of in-flight queries to reach a certain value"""
    start_time = time()
    while (time() - start_time < timeout):
      num_in_flight_queries = self.get_num_in_flight_queries()
      if num_in_flight_queries == expected_val: return True
      sleep(1)
    LOG.info("The number of in flight queries: %s, expected: %s" %
        (num_in_flight_queries, expected_val))
    return False

  def wait_for_num_known_live_backends(self, expected_value, timeout=30, interval=1):
    start_time = time()
    while (time() - start_time < timeout):
      value = None
      try:
        value = self.get_num_known_live_backends(timeout=timeout, interval=interval)
      except Exception, e:
        LOG.error(e)
      if value == expected_value:
        LOG.info("num_known_live_backends has reached value: %s" % value)
        return value
      else:
        LOG.info("Waiting for num_known_live_backends=%s. Current value: %s" %\
            (expected_value, value))
      sleep(1)
    assert 0, 'num_known_live_backends did not reach expected value in time'

  def read_query_profile_page(self, query_id, timeout=10, interval=1):
    """Fetches the raw contents of the query's runtime profile webpage.
    Fails an assertion if Impala's webserver is unavailable or the query's
    profile page doesn't exist."""
    return self.read_debug_webpage("query_profile?query_id=%s&raw" % (query_id))

  def get_query_status(self, query_id):
    """Gets the 'Query Status' section of the query's runtime profile."""
    page = self.read_query_profile_page(query_id)
    status_line =\
        next((x for x in page.split('\n') if re.search('Query Status:', x)), None)
    return status_line.split('Query Status:')[1].strip()

  def wait_for_query_state(self, client, query_handle, target_state,
                              timeout=10, interval=1):
    """Keeps polling for the query's state using client in the given interval until
    the query's state reaches the target state or the given timeout has been reached."""
    start_time = time()
    while (time() - start_time < timeout):
      try:
        query_state = client.get_state(query_handle)
      except Exception as e:
        pass
      if query_state == target_state:
        return
      sleep(interval)
    assert target_state == query_state, 'Did not reach query state in time'
    return

  def wait_for_query_status(self, client, query_id, expected_content,
                               timeout=30, interval=1):
    """Polls for the query's status in the query profile web page to contain the
    specified content. Returns False if the timeout was reached before a successful
    match, True otherwise."""
    start_time = time()
    query_status = ""
    while (time() - start_time < timeout):
      try:
        query_status = self.get_query_status(query_id)
        if query_status is None:
          assert False, "Could not find 'Query Status' section in profile of "\
                    "query with id %s:\n%s" % (query_id)
      except Exception as e:
        pass
      if expected_content in query_status:
        return True
      sleep(interval)
    return False

  def create_beeswax_client(self, use_kerberos=False):
    """Creates a new beeswax client connection to the impalad"""
    client = create_connection('%s:%d' % (self.hostname, self.beeswax_port), use_kerberos)
    client.connect()
    return client

  def create_ldap_beeswax_client(self, user, password, use_ssl=False):
    client = create_ldap_connection('%s:%d' % (self.hostname, self.beeswax_port),
                                    user=user, password=password, use_ssl=use_ssl)
    client.connect()
    return client

  def create_hs2_client(self):
    """Creates a new HS2 client connection to the impalad"""
    host, port = (self.hostname, self.hs2_port)
    socket = TSocket(host, port)
    transport = TBufferedTransport(socket)
    transport.open()
    protocol = TBinaryProtocol.TBinaryProtocol(transport)
    hs2_client = TCLIService.Client(protocol)
    return hs2_client

  def get_catalog_object_dump(self, object_type, object_name):
    return self.read_debug_webpage('catalog_objects?object_type=%s&object_name=%s' %\
        (object_type, object_name))


# Allows for interacting with the StateStore service to perform operations such as
# accessing the debug webpage.
class StateStoredService(BaseImpalaService):
  def __init__(self, hostname, webserver_port):
    super(StateStoredService, self).__init__(hostname, webserver_port)

  def wait_for_live_subscribers(self, num_subscribers, timeout=15, interval=1):
    self.wait_for_metric_value('statestore.live-backends', num_subscribers,
                               timeout=timeout, interval=interval)


# Allows for interacting with the Catalog service to perform operations such as
# accessing the debug webpage.
class CatalogdService(BaseImpalaService):
  def __init__(self, hostname, webserver_port, service_port):
    super(CatalogdService, self).__init__(hostname, webserver_port)
    self.service_port = service_port

  def get_catalog_object_dump(self, object_type, object_name):
    return self.read_debug_webpage('catalog_objects?object_type=%s&object_name=%s' %\
        (object_type, object_name))