import agent_util
import logging
import traceback

logger = logging.getLogger(__name__)

### Mapping of JMX URI entries to their agent readable counterparts
JMX_MAPPING = {
    "broker.bips.oneminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesInPerSec", "OneMinuteRate", None),
    "broker.bips.fiveminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesInPerSec", "FiveMinuteRate", None),
    "broker.bips.fifteenminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesInPerSec", "FifteenMinuteRate", None),
    "broker.bips.meanrate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesInPerSec", "MeanRate", None),
    "broker.bops.oneminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesOutPerSec", "OneMinuteRate", None),
    "broker.bops.fiveminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesOutPerSec", "FiveMinuteRate", None),
    "broker.bops.fifteenminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesOutPerSec", "FifteenMinuteRate", None),
    "broker.bops.meanrate" : ("kafka.server:type=BrokerTopicMetrics,name=BytesOutPerSec", "MeanRate", None),
    "broker.mips.oneminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=MessagesInPerSec", "OneMinuteRate", None),
    "broker.mips.fiveminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=MessagesInPerSec", "FiveMinuteRate", None),
    "broker.mips.fifteenminuterate" : ("kafka.server:type=BrokerTopicMetrics,name=MessagesInPerSec", "FifteenMinuteRate", None),
    "broker.mips.meanrate" : ("kafka.server:type=BrokerTopicMetrics,name=MessagesInPerSec", "MeanRate", None),
    "underreplicatedpartitions" : ("kafka.server:type=ReplicaManager,name=UnderReplicatedPartitions", "Value", None),
    "fetch.queue-size" : ("kafka.server:type=Fetch", "queue-size", None),
    "memory.heap.committed" : ("java.lang:type=Memory", "HeapMemoryUsage", "committed"),
    "memory.heap.used" : ("java.lang:type=Memory", "HeapMemoryUsage", "used"),
    "memory.heap.max" : ("java.lang:type=Memory", "HeapMemoryUsage", "max"),
}
####

def discover_beans(connection):
    allowed_beans = ['kafka.server:type=BrokerTopicMetrics', 'kafka.server:type=ReplicaManager', 
        'kafka.log:type=LogFlushStats', 'java.lang:type=Memory', 'kafka.server:type=Fetch']
    ignored_topics = ['__consumer_offsets', 'ReplicaFetcherThread-0-2']

    discovered_beans = []
    discovered_topics = []
    avail_beans = connection.queryMBeans(None, None)
    for bean in avail_beans:
        name = bean.objectName.toString()
        if any(b in name for b in allowed_beans):
            discovered_beans.append(bean.toString())
            if 'topic' in name:
                topic = name.split(',')[2].split('=')[1]
                if topic not in ignored_topics:
                    discovered_topics.append(topic)

    return discovered_beans, list(set(discovered_topics))


class ApacheKafkaPlugin(agent_util.Plugin):
    textkey = "apache_kafka_jmx"
    label = "Apache Kafka (JMX)"
    connection = None

    @classmethod
    def get_metadata(self, config):
        status = agent_util.SUPPORTED
        msg = None

        # Check for jmx configuration block
        if not config:
            self.log.info("No JMX configuration found")
            return {}

        # make sure jpype1 is installed first
        try:
            import jpype
            from jpype import java, javax
        except:
            msg = "Unable to access JMX metrics due to missing jpype library."
            self.log.info(msg)
            status = agent_util.MISCONFIGURED

        # Check for config setting sin jmx configuration block
        for key in ['port', 'host']:
            if key not in config:
                msg = "Missing value for %s in the [jmx] block of the agent config file." % key
                self.log.info(msg)
                status = agent_util.MISCONFIGURED

        # we'll need to get the default JVM path if not specified. If that doesn't work, throw an error
        if 'jvm_path' not in config:
            try:
                jvm_path = jpype.getDefaultJVMPath()
                if not jvm_path:
                    msg = "Unable to find JVM, please specify 'jvm_path' in the [jmx] block of the agent config file."
                    self.log.info(msg)
                    status = agent_util.MISCONFIGURED
            except:                
                msg = "Unable to find JVM, please specify 'jvm_path' in the [jmx] block of the agent config file."
                self.log.info(msg)
                self.log.error(traceback.format_exc())
                status = agent_util.MISCONFIGURED
        elif 'jvm_path' in config:
            jvm_path = config['jvm_path']


        try:
            if status == agent_util.SUPPORTED and not jpype.isJVMStarted():
                jpype.startJVM(jvm_path)
        except:
            msg = "Unable to access JMX metrics because JVM cannot be started."
            self.log.info(msg)
            status = agent_util.MISCONFIGURED

        if status == agent_util.SUPPORTED:
            try:
                if not jpype.isJVMStarted(): jpype.startJVM(config['jvm_path'])
                jhash = java.util.HashMap()
                if config.get('username') and config.get('password'):
                    jarray = jpype.JArray(java.lang.String)([config['username'], config['password']])
                    jhash.put(javax.management.remote.JMXConnector.CREDENTIALS, jarray)
                url = "service:jmx:rmi:///jndi/rmi://%s:%d/jmxrmi" % (config['host'], int(config['port']))
                jmxurl = javax.management.remote.JMXServiceURL(url)
                # Gather the topics we can monitor
                jmxsoc = javax.management.remote.JMXConnectorFactory.connect(jmxurl, jhash)
                self.connection = jmxsoc.getMBeanServerConnection()
                
            except:
                msg = "Unable to access JMX metrics, JMX is not running or not installed."
                self.log.info(msg)
                status = agent_util.MISCONFIGURED
                return {}

        beans, topics = discover_beans(self.connection)

        metadata = {
            "broker.bips.oneminuterate" : {
                "label": "Topic Byes In/sec - 1 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.bips.fiveminuterate" : {
                "label": "Topic Byes In/sec - 5 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.bips.fifteenminuterate" : {
                "label": "Topic Byes In/sec - 15 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.bips.meanrate" : {
                "label": "Topic Byes In/sec - Avg",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.bops.oneminuterate" : {
                "label": "Topic Byes Out/sec - 1 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.bops.fiveminuterate" : {
                "label": "Topic Byes Out/sec - 5 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.bops.fifteenminuterate" : {
                "label": "Topic Byes Out/sec - 15 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.bops.meanrate" : {
                "label": "Topic Byes Out/sec - Avg",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "broker.mips.oneminuterate" : {
                "label": "Topic Messages In/sec - 1 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "messages"
            },

            "broker.mips.fiveminuterate" : {
                "label": "Topic Messages In/sec - 5 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "messages"
            },

            "broker.mips.fifteenminuterate" : {
                "label": "Topic Messages In/sec - 15 min",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "messages"
            },

            "broker.mips.meanrate" : {
                "label": "Topic Messages In/sec - Avg",
                "options": topics,
                "status": status,
                "error_message": msg,
                "unit": "messages"
            },

            "underreplicatedpartitions" : {
                "label": "Replica Manager Unreplicated Partitions",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "partitions"
            },
            
            "fetch.queue-size" : {
                "label": "Queued messages",
                "options": None,
                "status": status,
                "error_message": msg,
            },

            "memory.heap.committed" : {
                "label": "Heap Memory - Committed",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "memory.heap.used" : {
                "label": "Heap Memory - Used",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

            "memory.heap.max" : {
                "label": "Heap Memory - Max",
                "options": None,
                "status": status,
                "error_message": msg,
                "unit": "bytes"
            },

        }
        return metadata

    def check(self, textkey, data, config):
        try:
            import jpype
            from jpype import java, javax
        except:
            self.log.error("Unable to import jpype! Is it installed?")
            return None

        try:
           # we'll need to get the default JVM path if not specified. If that doesn't work, throw an error
            if 'jvm_path' not in config:
                try:
                    jvm_path = jpype.getDefaultJVMPath()
                    if not jvm_path:
                        msg = "Unable to find JVM, please specify 'jvm_path' in the [jmx] block of the agent config file."
                        self.log.info(msg)
                except:                
                    msg = "Unable to find JVM, please specify 'jvm_path' in the [jmx] block of the agent config file."
                    self.log.info(msg)
                    self.log.error(traceback.format_exc())
            elif 'jvm_path' in config:
                jvm_path = config['jvm_path']

            if not jpype.isJVMStarted(): 
                jpype.startJVM(jvm_path)
            jhash = java.util.HashMap()
            if config.get('username') and config.get('password'):
                jarray = jpype.JArray(java.lang.String)([config['username'], config['password']])
                jhash.put(javax.management.remote.JMXConnector.CREDENTIALS, jarray)
            url = "service:jmx:rmi:///jndi/rmi://%s:%d/jmxrmi" % (config['host'], int(config['port']))
            jmxurl = javax.management.remote.JMXServiceURL(url)
            jmxsoc = javax.management.remote.JMXConnectorFactory.connect(jmxurl, jhash)
            connection = jmxsoc.getMBeanServerConnection()

            parts = JMX_MAPPING.get(textkey, None)
            if parts is None or not parts:
                self.log.error("Unable to find Kafka metric %s in known metrics!" % textkey)
                return None

            # start building the JMX object
            obj = parts[0]
            if data:
                obj +=',topic=%s' % data
            # get the actual metric
            attribute = parts[1]
            # if the metric is buried deeper in a dict, grab it
            val = parts[2]

            res = connection.getAttribute(javax.management.ObjectName(obj), attribute)
            log_msg = "Checking Kafka metric %s" % attribute

            if val is not None:
                return res.contents.get(val).floatValue()
            else:
                log_msg += " with key %s" % val
                return res.floatValue()

            self.log.debug(log_msg)
        except:
            self.log.critical("Error checking Kafka metric %s - %s \n%s" % (textkey, data, traceback.format_exc()))
            return None