# (C) Datadog, Inc. 2019
# All rights reserved
# Licensed under a 3-clause BSD style license (see LICENSE)
from __future__ import division

import os
import ssl
from collections import defaultdict
from datetime import datetime
from itertools import chain
from os.path import expanduser, isdir

import vertica_python as vertica
from six import iteritems
from vertica_python.vertica.column import timestamp_tz_parse

from datadog_checks.base import AgentCheck, is_affirmative
from datadog_checks.base.utils.containers import iter_unique

from . import views
from .utils import kilobytes_to_bytes, node_state_to_service_check

# Python 3 only
PROTOCOL_TLS_CLIENT = getattr(ssl, 'PROTOCOL_TLS_CLIENT', ssl.PROTOCOL_TLS)


class VerticaCheck(AgentCheck):
    __NAMESPACE__ = 'vertica'
    SERVICE_CHECK_CONNECT = 'can_connect'
    SERVICE_CHECK_NODE_STATE = 'node_state'

    def __init__(self, name, init_config, instances):
        super(VerticaCheck, self).__init__(name, init_config, instances)

        self._db = self.instance.get('db', '')
        self._server = self.instance.get('server', '')
        self._port = int(self.instance.get('port', 5433))
        self._username = self.instance.get('username', '')
        self._password = self.instance.get('password', '')
        self._backup_servers = [
            (bs.get('server', self._server), int(bs.get('port', self._port)))
            for bs in self.instance.get('backup_servers', [])
        ]
        self._connection_load_balance = is_affirmative(self.instance.get('connection_load_balance', False))
        self._timeout = float(self.instance.get('timeout', 10))
        self._tags = self.instance.get('tags', [])

        self._client_lib_log_level = self.instance.get('client_lib_log_level')

        self._tls_verify = is_affirmative(self.instance.get('tls_verify', False))
        self._validate_hostname = is_affirmative(self.instance.get('validate_hostname', True))

        self._cert = self.instance.get('cert', '')
        if self._cert:  # no cov
            self._cert = expanduser(self._cert)
            self._tls_verify = True

        self._private_key = self.instance.get('private_key', '')
        if self._private_key:  # no cov
            self._private_key = expanduser(self._private_key)

        self._cafile = None
        self._capath = None
        ca_cert = self.instance.get('ca_cert', '')
        if ca_cert:  # no cov
            ca_cert = expanduser(ca_cert)
            if isdir(ca_cert):
                self._capath = ca_cert
            else:
                self._cafile = ca_cert

            self._tls_verify = True

        custom_queries = self.instance.get('custom_queries', [])
        use_global_custom_queries = self.instance.get('use_global_custom_queries', True)

        # Handle overrides
        if use_global_custom_queries == 'extend':
            custom_queries.extend(self.init_config.get('global_custom_queries', []))
        elif 'global_custom_queries' in self.init_config and is_affirmative(use_global_custom_queries):
            custom_queries = self.init_config.get('global_custom_queries', [])

        # Deduplicate
        self._custom_queries = list(iter_unique(custom_queries))

        # Add global database tag
        self._tags.append('db:{}'.format(self._db))

        # We'll connect on the first check run
        self._connection = None

        # Cache database results for re-use among disparate functions
        self._view = defaultdict(list)

    def check(self, instance):
        if self._connection is None:
            connection = self.get_connection()
            if connection is None:
                return

            self._connection = connection
        elif self._connection_load_balance or self._connection.closed():
            self._connection.reset_connection()

        # The order of queries is important as some results are cached for later re-use
        try:
            self.query_licenses()
            self.query_license_audits()
            self.query_system()
            self.query_nodes()
            self.query_projections()
            self.query_projection_storage()
            self.query_storage_containers()
            self.query_host_resources()
            self.query_query_metrics()
            self.query_resource_pool_status()
            self.query_disk_storage()
            self.query_resource_usage()
            self.query_custom()
        finally:
            self._view.clear()

    def query_licenses(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/CATALOG/LICENSES.htm
        for db_license in self.iter_rows(views.Licenses):
            self._view[views.Licenses].append(db_license)

            tags = ['license_type:{}'.format(db_license['licensetype'])]
            tags.extend(self._tags)

            expiration = db_license['end_date']
            if expiration and expiration != 'Perpetual':
                expiration = timestamp_tz_parse(expiration)
                seconds_until_expiration = (expiration - datetime.now(tz=expiration.tzinfo)).total_seconds()
            else:
                seconds_until_expiration = -1

            self.gauge('license.expiration', seconds_until_expiration, tags=tags)

    def query_license_audits(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/CATALOG/LICENSE_AUDITS.htm
        for license_audit in self.iter_rows(views.LicenseAudits):
            last_audit = license_audit['audit_start_timestamp']
            if last_audit:
                seconds_since_last_audit = (datetime.now(tz=last_audit.tzinfo) - last_audit).total_seconds()
            else:
                seconds_since_last_audit = -1
            self.gauge('license.latest_audit', seconds_since_last_audit, tags=self._tags)

            size = int(license_audit['license_size_bytes'])
            used = int(license_audit['database_size_bytes'])
            self.gauge('license.size', size, tags=self._tags)
            self.gauge('license.used', used, tags=self._tags)
            self.gauge('license.usable', size - used, tags=self._tags)
            self.gauge('license.utilized', used / size * 100, tags=self._tags)

    def query_system(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/SYSTEM.htm

        # Will only be 1 system
        for system in self.iter_rows(views.System):
            total_nodes = system['node_count']
            self.gauge('node.total', total_nodes, tags=self._tags)
            self.gauge('node.down', system['node_down_count'], tags=self._tags)

            # Is is possible for there to be no restriction
            allowed_nodes = self._view[views.Licenses][0]['node_restriction']
            if allowed_nodes is not None:
                self.gauge('node.allowed', allowed_nodes, tags=self._tags)
                self.gauge('node.available', allowed_nodes - total_nodes, tags=self._tags)

            self.gauge('ksafety.current', system['current_fault_tolerance'], tags=self._tags)
            self.gauge('ksafety.intended', system['designed_fault_tolerance'], tags=self._tags)

            self.gauge('epoch.ahm', system['ahm_epoch'], tags=self._tags)
            self.gauge('epoch.current', system['current_epoch'], tags=self._tags)
            self.gauge('epoch.last_good', system['last_good_epoch'], tags=self._tags)

    def query_nodes(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/CATALOG/NODES.htm
        for node in self.iter_rows(views.Nodes):
            tags = ['node_name:{}'.format(node['node_name'])]
            tags.extend(self._tags)

            node_state = node['node_state']
            self.service_check(
                self.SERVICE_CHECK_NODE_STATE, node_state_to_service_check(node_state), message=node_state, tags=tags
            )

    def query_projections(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/CATALOG/PROJECTIONS.htm
        total_projections = 0
        unsegmented_projections = 0
        unsafe_projections = 0

        for projection in self.iter_rows(views.Projections):
            total_projections += 1

            if not projection['is_segmented']:
                unsegmented_projections += 1

            if not projection['is_up_to_date']:
                unsafe_projections += 1

        self.gauge('projection.total', total_projections, tags=self._tags)

        self.gauge('projection.unsegmented', unsegmented_projections, tags=self._tags)
        if total_projections:
            unsegmented_percent = unsegmented_projections / total_projections * 100
        else:
            unsegmented_percent = 0
        self.gauge('projection.unsegmented_percent', unsegmented_percent, tags=self._tags)

        self.gauge('projection.unsafe', unsafe_projections, tags=self._tags)
        if total_projections:
            unsafe_percent = unsafe_projections / total_projections * 100
        else:
            unsafe_percent = 0
        self.gauge('projection.unsafe_percent', unsafe_percent, tags=self._tags)

    def query_projection_storage(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/PROJECTION_STORAGE.htm
        projection_data = defaultdict(
            lambda: defaultdict(
                lambda: defaultdict(
                    lambda: {'ros_count': 0, 'rows_ros': 0, 'rows_wos': 0, 'used_ros': 0, 'used_wos': 0}
                )
            )
        )

        for ps in self.iter_rows(views.ProjectionStorage):
            projection = projection_data[ps['node_name']][ps['anchor_table_name']][ps['projection_name']]

            projection['ros_count'] += ps['ros_count']
            projection['rows_ros'] += ps['ros_row_count']
            projection['rows_wos'] += ps['wos_row_count']
            projection['used_ros'] += ps['ros_used_bytes']
            projection['used_wos'] += ps['wos_used_bytes']

        total_rows_ros = 0
        total_rows_wos = 0
        total_used_ros = 0
        total_used_wos = 0

        # My understanding is that nodes have multiple tables, which in turn can have multiple projections
        for node, tables in iteritems(projection_data):
            node_tags = ['node_name:{}'.format(node)]
            node_tags.extend(self._tags)

            node_rows_ros = 0
            node_rows_wos = 0
            node_used_ros = 0
            node_used_wos = 0

            for table, projections in iteritems(tables):
                table_tags = ['table_name:{}'.format(table)]
                table_tags.extend(node_tags)

                table_rows_ros = 0
                table_rows_wos = 0
                table_used_ros = 0
                table_used_wos = 0

                for projection, data in iteritems(projections):
                    projection_tags = ['projection_name:{}'.format(projection)]
                    projection_tags.extend(table_tags)

                    projection_rows_ros = data['rows_ros']
                    projection_rows_wos = data['rows_wos']
                    projection_used_ros = data['used_ros']
                    projection_used_wos = data['used_wos']

                    self.gauge('projection.ros.containers', data['ros_count'], tags=projection_tags)
                    self.gauge('projection.row.ros', projection_rows_ros, tags=projection_tags)
                    self.gauge('projection.row.wos', projection_rows_wos, tags=projection_tags)
                    self.gauge('projection.row.total', projection_rows_ros + projection_rows_wos, tags=projection_tags)
                    self.gauge('projection.disk.used.ros', projection_used_ros, tags=projection_tags)
                    self.gauge('projection.disk.used.wos', projection_used_wos, tags=projection_tags)
                    self.gauge('projection.disk.used', projection_used_ros + projection_used_wos, tags=projection_tags)

                    table_rows_ros += projection_rows_ros
                    table_rows_wos += projection_rows_wos
                    table_used_ros += projection_used_ros
                    table_used_wos += projection_used_wos

                self.gauge('table.row.ros', table_rows_ros, tags=table_tags)
                self.gauge('table.row.wos', table_rows_wos, tags=table_tags)
                self.gauge('table.row.total', table_rows_ros + table_rows_wos, tags=table_tags)
                self.gauge('table.disk.used.ros', table_used_ros, tags=table_tags)
                self.gauge('table.disk.used.wos', table_used_wos, tags=table_tags)
                self.gauge('table.disk.used', table_used_ros + table_used_wos, tags=table_tags)

                node_rows_ros += table_rows_ros
                node_rows_wos += table_rows_wos
                node_used_ros += table_used_ros
                node_used_wos += table_used_wos

            self.gauge('node.row.ros', node_rows_ros, tags=node_tags)
            self.gauge('node.row.wos', node_rows_wos, tags=node_tags)
            self.gauge('node.row.total', node_rows_ros + node_rows_wos, tags=node_tags)
            self.gauge('node.disk.used.ros', node_used_ros, tags=node_tags)
            self.gauge('node.disk.used.wos', node_used_wos, tags=node_tags)
            self.gauge('node.disk.used', node_used_ros + node_used_wos, tags=node_tags)

            total_rows_ros += node_rows_ros
            total_rows_wos += node_rows_wos
            total_used_ros += node_used_ros
            total_used_wos += node_used_wos

        self.gauge('row.ros', total_rows_ros, tags=self._tags)
        self.gauge('row.wos', total_rows_wos, tags=self._tags)
        self.gauge('row.total', total_rows_ros + total_rows_wos, tags=self._tags)
        self.gauge('disk.used.ros', total_used_ros, tags=self._tags)
        self.gauge('disk.used.wos', total_used_wos, tags=self._tags)
        self.gauge('disk.used', total_used_ros + total_used_wos, tags=self._tags)

    def query_storage_containers(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/STORAGE_CONTAINERS.htm
        container_data = defaultdict(lambda: defaultdict(lambda: defaultdict(lambda: {'delete_vectors': 0})))

        for sc in self.iter_rows(views.StorageContainers):
            container = container_data[sc['node_name']][sc['projection_name']][sc['storage_type'].lower()]

            container['delete_vectors'] += sc['delete_vector_count']

        total_delete_vectors = 0

        for node, projections in iteritems(container_data):
            node_tags = ['node_name:{}'.format(node)]
            node_tags.extend(self._tags)

            node_delete_vectors = 0

            for projection, containers in iteritems(projections):
                projection_tags = ['projection_name:{}'.format(projection)]
                projection_tags.extend(node_tags)

                projection_delete_vectors = 0

                for container_type, data in iteritems(containers):
                    container_tags = ['container_type:{}'.format(container_type)]
                    container_tags.extend(projection_tags)

                    container_type_delete_vectors = data['delete_vectors']

                    self.gauge('projection.delete_vectors', container_type_delete_vectors, tags=container_tags)

                    projection_delete_vectors += container_type_delete_vectors

                node_delete_vectors += projection_delete_vectors

            self.gauge('node.delete_vectors', node_delete_vectors, tags=node_tags)

            total_delete_vectors += node_delete_vectors

        self.gauge('delete_vectors', total_delete_vectors, tags=self._tags)

    def query_host_resources(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/HOST_RESOURCES.htm
        for host in self.iter_rows(views.HostResources):
            tags = ['host_name:{}'.format(host['host_name'])]
            tags.extend(self._tags)

            self.gauge('processor.cpu', host['processor_count'], tags=tags)
            self.gauge('processor.core', host['processor_core_count'], tags=tags)
            self.gauge('file.max', host['open_files_limit'], tags=tags)
            self.gauge('file.open', host['opened_file_count'], tags=tags)
            self.gauge('socket.open', host['opened_socket_count'], tags=tags)
            self.gauge('thread.max', host['threads_limit'], tags=tags)

            # Memory
            total = host['total_memory_bytes']
            usable = host['total_memory_free_bytes']
            used = total - usable

            self.gauge('memory.total', total, tags=tags)
            self.gauge('memory.usable', usable, tags=tags)
            self.gauge('memory.used', used, tags=tags)

            if total:
                utilized = used / total * 100
            else:
                utilized = 0
            self.gauge('memory.utilized', utilized, tags=tags)

            # Swap
            total = host['total_swap_memory_bytes']
            usable = host['total_swap_memory_free_bytes']
            used = total - usable

            self.gauge('memory.swap.total', total, tags=tags)
            self.gauge('memory.swap.usable', usable, tags=tags)
            self.gauge('memory.swap.used', used, tags=tags)

            if total:
                utilized = used / total * 100
            else:
                utilized = 0
            self.gauge('memory.swap.utilized', utilized, tags=tags)

    def query_query_metrics(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/QUERY_METRICS.htm
        for node in self.iter_rows(views.QueryMetrics):
            tags = ['node_name:{}'.format(node['node_name'])]
            tags.extend(self._tags)

            self.gauge('connection.active', node['active_user_session_count'], tags=tags)
            self.monotonic_count('connection.total', node['total_user_session_count'], tags=tags)
            self.gauge('query.active', node['running_query_count'], tags=tags)
            self.monotonic_count('query.total', node['executed_query_count'], tags=tags)

    def query_resource_pool_status(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/RESOURCE_POOL_STATUS.htm
        for pool in self.iter_rows(views.ResourcePoolStatus):
            tags = ['node_name:{}'.format(pool['node_name']), 'pool_name:{}'.format(pool['pool_name'])]
            tags.extend(self._tags)

            self.gauge(
                'resource_pool.memory.borrowed', kilobytes_to_bytes(pool['general_memory_borrowed_kb']), tags=tags
            )
            self.gauge('resource_pool.memory.max', kilobytes_to_bytes(pool['max_memory_size_kb']), tags=tags)
            self.gauge('resource_pool.memory.used', kilobytes_to_bytes(pool['memory_inuse_kb']), tags=tags)
            self.gauge('resource_pool.query.running', pool['running_query_count'], tags=tags)

    def query_disk_storage(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/DISK_STORAGE.htm
        for storage in self.iter_rows(views.DiskStorage):
            tags = [
                'node_name:{}'.format(storage['node_name']),
                'storage_status:{}'.format(storage['storage_status']),
                'storage_usage:{}'.format(storage['storage_usage']),
            ]
            tags.extend(self._tags)

            # Space
            block_size = storage['disk_block_size_bytes']
            usable = block_size * storage['disk_space_free_blocks']
            used = block_size * storage['disk_space_used_blocks']
            total = usable + used

            self.gauge('storage.size', total, tags=tags)
            self.gauge('storage.usable', usable, tags=tags)
            self.gauge('storage.used', used, tags=tags)

            if total:
                utilized = used / total * 100
            else:
                utilized = 0
            self.gauge('storage.utilized', utilized, tags=tags)

            # Latency
            latency = storage['latency']
            self.gauge('storage.latency', latency, tags=tags)

            if latency:
                latency_reciprocal = 1 / latency
            else:
                latency_reciprocal = 0

            # Throughput
            throughput = storage['throughput']
            self.gauge('storage.throughput', throughput, tags=tags)

            if throughput:
                throughput_reciprocal = 1 / throughput
            else:
                throughput_reciprocal = 0

            # Time to read 1 MiB
            self.gauge('storage.speed', latency_reciprocal + throughput_reciprocal, tags=tags)

    def query_resource_usage(self):
        # https://www.vertica.com/docs/9.2.x/HTML/Content/Authoring/SQLReferenceManual/SystemTables/MONITOR/RESOURCE_USAGE.htm
        for node in self.iter_rows(views.ResourceUsage):
            tags = ['node_name:{}'.format(node['node_name'])]
            tags.extend(self._tags)

            self.gauge('node.resource_requests', node['request_count'], tags=tags)
            self.gauge('thread.active', node['active_thread_count'], tags=tags)

    def query_custom(self):
        for custom_query in self._custom_queries:
            query = custom_query.get('query')
            if not query:  # no cov
                self.log.error('Custom query field `query` is required')
                continue

            columns = custom_query.get('columns')
            if not columns:  # no cov
                self.log.error('Custom query field `columns` is required')
                continue

            self.log.debug('Running custom query for Vertica')
            cursor = self._connection.cursor()
            cursor.execute(query)

            rows = cursor.iterate()

            # Trigger query execution
            try:
                first_row = next(rows)
            except Exception as e:  # no cov
                self.log.error('Error executing custom query: {}'.format(e))
                continue

            for row in chain((first_row,), rows):
                if not row:  # no cov
                    self.log.debug('Custom query returned an empty result')
                    continue

                if len(columns) != len(row):  # no cov
                    self.log.error('Custom query result expected {} columns, got {}'.format(len(columns), len(row)))
                    continue

                metric_info = []
                query_tags = list(self._tags)
                query_tags.extend(custom_query.get('tags', []))

                for column, value in zip(columns, row):
                    # Columns can be ignored via configuration.
                    if not column:  # no cov
                        continue

                    name = column.get('name')
                    if not name:  # no cov
                        self.log.error('Column field `name` is required')
                        break

                    column_type = column.get('type')
                    if not column_type:  # no cov
                        self.log.error('Column field `type` is required for column `{}`'.format(name))
                        break

                    if column_type == 'tag':
                        query_tags.append('{}:{}'.format(name, value))
                    else:
                        if not hasattr(self, column_type):
                            self.log.error(
                                'Invalid submission method `{}` for metric column `{}`'.format(column_type, name)
                            )
                            break
                        try:
                            metric_info.append((name, float(value), column_type))
                        except (ValueError, TypeError):  # no cov
                            self.log.error('Non-numeric value `{}` for metric column `{}`'.format(value, name))
                            break

                # Only submit metrics if there were absolutely no errors - all or nothing.
                else:
                    for info in metric_info:
                        metric, value, method = info
                        getattr(self, method)(metric, value, tags=query_tags)

    def get_connection(self):
        connection_options = {
            'database': self._db,
            'host': self._server,
            'port': self._port,
            'user': self._username,
            'password': self._password,
            'backup_server_node': self._backup_servers,
            'connection_load_balance': self._connection_load_balance,
            'connection_timeout': self._timeout,
        }
        if self._client_lib_log_level:
            connection_options['log_level'] = self._client_lib_log_level
            # log_path is required by vertica client for using logging
            # we still get vertica client logs in the agent log even if `log_path` is set to os.devnull
            connection_options['log_path'] = os.devnull

        if self._tls_verify:  # no cov
            # https://docs.python.org/3/library/ssl.html#ssl.SSLContext
            # https://docs.python.org/3/library/ssl.html#ssl.PROTOCOL_TLS
            tls_context = ssl.SSLContext(protocol=PROTOCOL_TLS_CLIENT)

            # https://docs.python.org/3/library/ssl.html#ssl.SSLContext.verify_mode
            tls_context.verify_mode = ssl.CERT_REQUIRED

            # https://docs.python.org/3/library/ssl.html#ssl.SSLContext.check_hostname
            tls_context.check_hostname = self._validate_hostname

            # https://docs.python.org/3/library/ssl.html#ssl.SSLContext.load_verify_locations
            if self._cafile or self._capath:
                tls_context.load_verify_locations(self._cafile, self._capath, None)

            # https://docs.python.org/3/library/ssl.html#ssl.SSLContext.load_default_certs
            else:
                tls_context.load_default_certs(ssl.Purpose.SERVER_AUTH)

            # https://docs.python.org/3/library/ssl.html#ssl.SSLContext.load_cert_chain
            if self._cert:
                tls_context.load_cert_chain(self._cert, keyfile=self._private_key)

            connection_options['ssl'] = tls_context

        try:
            connection = vertica.connect(**connection_options)
        except Exception as e:
            self.log.error('Unable to connect to database `{}` as user `{}`: {}'.format(self._db, self._username, e))
            self.service_check(self.SERVICE_CHECK_CONNECT, self.CRITICAL, tags=self._tags)
        else:
            self.service_check(self.SERVICE_CHECK_CONNECT, self.OK, tags=self._tags)
            return connection

    def iter_rows(self, view):
        cursor = self._connection.cursor('dict')
        cursor.execute(view.query)

        for row in cursor.iterate():
            yield row
