# coding=utf-8

"""
Collects all number values from the db.serverStatus() and db.engineStatus()
command, other values are ignored.

#### Dependencies

 * pymongo

"""

import re

import diamond.collector
from diamond.collector import str_to_bool

try:
    import pymongo
except ImportError:
    pymongo = None

try:
    from pymongo import ReadPreference
except ImportError:
    ReadPreference = None


class TokuMXCollector(diamond.collector.Collector):
    def __init__(self, *args, **kwargs):
        self.__totals = {}
        super(TokuMXCollector, self).__init__(*args, **kwargs)

    def get_default_config_help(self):
        config_help = super(TokuMXCollector, self).get_default_config_help()
        config_help.update(
            {
                "hosts": "Array of hostname(:port) elements to get metrics from"
                "Set an alias by prefixing host:port with alias@",
                "host": "A single hostname(:port) to get metrics from"
                " (can be used instead of hosts and overrides it)",
                "user": "Username for authenticated login (optional)",
                "passwd": "Password for authenticated login (optional)",
                "databases": "A regex of which databases to gather metrics for."
                " Defaults to all databases.",
                "ignore_collections": "A regex of which collections to ignore."
                " MapReduce temporary collections (tmp.mr.*)"
                " are ignored by default.",
                "network_timeout": "Timeout for mongodb connection (in seconds)."
                " There is no timeout by default.",
                "simple": "Only collect the same metrics as mongostat.",
                "translate_collections": "Translate dot (.) to underscores (_)"
                " in collection names.",
            }
        )
        return config_help

    def get_default_config(self):
        """
        Returns the default collector settings
        """
        config = super(TokuMXCollector, self).get_default_config()
        config.update(
            {
                "path": "mongo",
                "hosts": ["localhost"],
                "user": None,
                "passwd": None,
                "databases": ".*",
                "ignore_collections": "^tmp\.mr\.",
                "network_timeout": None,
                "simple": "False",
                "translate_collections": "False",
            }
        )
        return config

    def collect(self):
        """Collect number values from db.server_status() and db.engine_status()"""

        if pymongo is None:
            self.log.error("Unable to import pymongo")
            return

        # we need this for backwards compatibility
        if "host" in self.config:
            self.config["hosts"] = [self.config["host"]]

        # convert network_timeout to integer
        if self.config["network_timeout"]:
            self.config["network_timeout"] = int(self.config["network_timeout"])

        # use auth if given
        if "user" in self.config:
            user = self.config["user"]
        else:
            user = None

        if "passwd" in self.config:
            passwd = self.config["passwd"]
        else:
            passwd = None

        for host in self.config["hosts"]:
            if len(self.config["hosts"]) == 1:
                # one host only, no need to have a prefix
                base_prefix = []
            else:
                matches = re.search("((.+)\@)?(.+)?", host)
                alias = matches.group(2)
                host = matches.group(3)

                if alias is None:
                    base_prefix = [re.sub("[:\.]", "_", host)]
                else:
                    base_prefix = [alias]

            try:
                if ReadPreference is None:
                    conn = pymongo.Connection(
                        host,
                        network_timeout=self.config["network_timeout"],
                        slave_okay=True,
                    )
                else:
                    conn = pymongo.Connection(
                        host,
                        network_timeout=self.config["network_timeout"],
                        read_preference=ReadPreference.SECONDARY,
                    )
            except Exception as e:
                self.log.error("Could not connect to mongodb: %s", e)
                continue

            # try auth
            if user:
                try:
                    conn.admin.authenticate(user, passwd)
                except Exception as e:
                    self.log.error(
                        "User auth given, but could not autheticate"
                        + " with host: %s, err: %s" % (host, e)
                    )

                    return {}

            server_status = conn.db.command("server_status")
            engine_status = conn.db.command("engine_status")
            data = dict(server_status.items() + engine_status.items())

            self._publish_transformed(data, base_prefix)

            if str_to_bool(self.config["simple"]):
                data = self._extract_simple_data(data)

            self._publish_dict_with_prefix(data, base_prefix)
            db_name_filter = re.compile(self.config["databases"])
            ignored_collections = re.compile(self.config["ignore_collections"])

            for db_name in conn.database_names():
                if not db_name_filter.search(db_name):
                    continue

                db_stats = conn[db_name].command("dbStats")
                db_prefix = base_prefix + ["databases", db_name]
                self._publish_dict_with_prefix(db_stats, db_prefix)

                for collection_name in conn[db_name].collection_names():
                    if ignored_collections.search(collection_name):
                        continue

                    collection_stats = conn[db_name].command(
                        "collstats", collection_name
                    )

                    if str_to_bool(self.config["translate_collections"]):
                        collection_name = collection_name.replace(".", "_")

                    collection_prefix = db_prefix + [collection_name]
                    self._publish_dict_with_prefix(collection_stats, collection_prefix)

    def _publish_transformed(self, data, base_prefix):
        """Publish values of type: counter or percent"""
        self._publish_dict_with_prefix(
            data.get("opcounters", {}),
            base_prefix + ["opcounters_per_sec"],
            self.publish_counter,
        )
        self._publish_dict_with_prefix(
            data.get("opcountersRepl", {}),
            base_prefix + ["opcountersRepl_per_sec"],
            self.publish_counter,
        )
        self._publish_dict_with_prefix(
            data.get("network", {}),
            base_prefix + ["network_per_sec"],
            self.publish_counter,
        )
        self._publish_metrics(
            base_prefix + ["extra_info_per_sec"],
            "page_faults",
            data.get("extra_info", {}),
            self.publish_counter,
        )

        def get_dotted_value(data, key_name):
            key_name = key_name.split(".")

            for i in key_name:
                data = data.get(i, {})

                if not data:
                    return 0

            return data

        def compute_interval(data, total_name):
            current_total = get_dotted_value(data, total_name)
            total_key = ".".join(base_prefix + [total_name])
            last_total = self.__totals.get(total_key, current_total)
            interval = current_total - last_total
            self.__totals[total_key] = current_total

            return interval

        def publish_percent(value_name, total_name, data):
            value = float(get_dotted_value(data, value_name) * 100)
            interval = compute_interval(data, total_name)
            key = ".".join(base_prefix + ["percent", value_name])
            self.publish_counter(
                key, value, time_delta=bool(interval), interval=interval
            )

        publish_percent("globalLock.lockTime", "globalLock.totalTime", data)

        locks = data.get("locks")

        if locks:
            if "." in locks:
                locks["_global_"] = locks["."]
                del locks["."]

            key_prefix = ".".join(base_prefix + ["percent"])
            db_name_filter = re.compile(self.config["databases"])
            interval = compute_interval(data, "uptimeMillis")

            for db_name in locks:
                if not db_name_filter.search(db_name):
                    continue

                r = get_dotted_value(locks, "%s.timeLockedMicros.r" % db_name)
                R = get_dotted_value(locks, ".%s.timeLockedMicros.R" % db_name)
                value = float(r + R) / 10

                if value:
                    self.publish_counter(
                        key_prefix + ".locks.%s.read" % db_name,
                        value,
                        time_delta=bool(interval),
                        interval=interval,
                    )

                w = get_dotted_value(locks, "%s.timeLockedMicros.w" % db_name)
                W = get_dotted_value(locks, "%s.timeLockedMicros.W" % db_name)
                value = float(w + W) / 10

                if value:
                    self.publish_counter(
                        key_prefix + ".locks.%s.write" % db_name,
                        value,
                        time_delta=bool(interval),
                        interval=interval,
                    )

    def _publish_dict_with_prefix(self, dict, prefix, publishfn=None):
        for key in dict:
            self._publish_metrics(prefix, key, dict, publishfn)

    def _publish_metrics(self, prev_keys, key, data, publishfn=None):
        """Recursively publish keys"""
        if key not in data:
            return

        value = data[key]
        keys = prev_keys + [key]

        if not publishfn:
            publishfn = self.publish

        if isinstance(value, dict):
            for new_key in value:
                self._publish_metrics(keys, new_key, value)
        elif isinstance(value, int) or isinstance(value, float):
            publishfn(".".join(keys), value)

    def _extract_simple_data(self, data):
        return {
            "connections": data.get("connections"),
            "globalLock": data.get("globalLock"),
            "indexCounters": data.get("indexCounters"),
        }
