# This try/except block will make this compatible with any Agent version
try:
    # first, try to import the base class from new versions of the Agent...
    from datadog_checks.base import AgentCheck
except ImportError:
    # ...if the above failed, the check is running in Agent version < 6.6.0
    from checks import AgentCheck

from datadog_checks.base.utils.subprocess_output import get_subprocess_output

# content of the special variable __version__ is shown in the Agent status page
__version__ = "1.0.0"

TAGS = <%= @tags %>


class Neo4jMemoryCheck(AgentCheck):
    def check(self, instance):
        neo4j_major_version = int("<%= node['datadog']['neo4j']['version'] %>"[0])
        protocol = "bolt+ssc" if neo4j_major_version >= 4 else "bolt"
        out, _err, _retcode = get_subprocess_output(
            [
                "cypher-shell",
                "--format",
                "plain",
                "-a",
                "{}://<%= node['ipaddress'] %>:7687".format(protocol),
                "-d",
                "graph.db",
                "-u",
                "<%= node['datadog']['neo4j']['user'] %>",
                "-p",
                "<%= node['datadog']['neo4j']['password'] %>",
                "CALL dbms.queryJmx('java.lang:type=Memory') yield attributes\
                WITH attributes.HeapMemoryUsage as heap, attributes.NonHeapMemoryUsage as nonHeap\
                WITH heap.value.properties as heapProps,nonHeap.value.properties as nonHeapProps\
                CALL dbms.queryJmx('neo4j.metrics:name=neo4j.database.system.transaction.last_committed_tx_id') yield attributes\
                RETURN heapProps.init,heapProps.committed,heapProps.used,heapProps.max,nonHeapProps.init,nonHeapProps.committed,\
                    nonHeapProps.used,nonHeapProps.max,heapProps.used + nonHeapProps.used, attributes['Count']['value'] as last_committed_tx_id\
                "
            ],
            self.log,
            raise_on_empty_output=True
        )
        formatted_string = out.replace(' ', '').replace('\n', ',')
        stats = formatted_string.split(',')
        # Just cherry-pick the numerical elements. 0-9 elements are the column names.
        self.gauge('neo4j.heap.init', int(stats[10]), tags=TAGS)
        self.gauge('neo4j.heap.committed', int(stats[11]), tags=TAGS)
        self.gauge('neo4j.heap.used', int(stats[12]), tags=TAGS)
        self.gauge('neo4j.heap.max', int(stats[13]), tags=TAGS)
        self.gauge('neo4j.non_heap.init', int(stats[14]), tags=TAGS)
        self.gauge('neo4j.non_heap.committed', int(stats[15]), tags=TAGS)
        self.gauge('neo4j.non_heap.used', int(stats[16]), tags=TAGS)
        self.gauge('neo4j.non_heap.max', int(stats[17]), tags=TAGS)
        self.gauge('neo4j.totalmem', int(stats[18]), tags=TAGS)
        self.gauge('neo4j.last_committed_tx_id', int(stats[19]), tags=TAGS)