#!/usr/bin/env python3
import os
import sys
import time
import json
import pickle
import codecs
import logging
import argparse
import traceback
import concurrent.futures

from ZabbixDockerAgent import *

parser = argparse.ArgumentParser(description='Zabbix Docker Agent')
parser.add_argument('-t', dest='test', action='store_true', default=False,
                    help='Test metrics collection. ( | jq \'.\' )')

args = parser.parse_args()

if args.test:
    TestRun()
    sys.exit(0)

logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
logLevel = logging.getLevelName(os.getenv('LOGLEVEL', 'INFO'))
ch.setLevel(logLevel)
logger.addHandler(ch)


def parse_key(itemKey):
    docker, section, key = itemKey.split('.')
    return section, key


if os.getenv('DOCKERHOST') is None or \
   os.getenv('ZBX_SERVER_HOST') is None:
    logger.error("Required env variables:"
                 " DOCKERHOST, ZBX_SERVER_HOST")
    sys.exit(1)


def doInstance(dckr):
    stats = {
        'containerCount': len(dckr.instancesData)
    }

    zbx = Zabbix()
    instanceItems = zbx.getItemList(
        host=os.getenv('DOCKERHOST'),
        hostMetadata=os.getenv('HOSTMETADATA', 'DockerInstance')
    )

    logger.debug(instanceItems)
    logger.debug("Containers found: {}".format(
        len(dckr.instancesData)))

    for instanceName, instanceData  in dckr.instancesData.items():
        logger.debug("Instance data: {} -> {}".format(
            instanceName,instanceData))

    ItemsSender = zbx.initSender(dataType='items')
    for instanceItem in instanceItems:
        if instanceItem['key'] == 'docker.instanceDiscovery':
            logger.debug("SENDING - docker.instanceDiscovery")
            logger.debug("{}".format(dckr.instancesData))

            payload = codecs.encode(
                pickle.dumps(dckr.instancesData), "base64").decode()
            ItemsSender.add_item(os.getenv('DOCKERHOST'),
                                 'docker.instanceDiscovery',
                                 payload)
        elif instanceItem['key'] == 'docker.containerCount':
            logger.debug("SENDING - docker.containerCount - {}".format(
                len(dckr.instancesData)
            ))
            ItemsSender.add_item(os.getenv('DOCKERHOST'),
                                 'docker.containerCount',
                                 len(dckr.instancesData))

    resp = ItemsSender.send()
    logger.debug("Sending items to Zabbix response: {}".format(resp))

    del zbx

    return stats


def doContainer(containerId, zbx, dckr):
    logger.debug("Collecting metrics for container: {} {}".format(
        containerId, dckr.containers[containerId]['name']))

    stats = {
        'itemCount': 0
    }
    items = zbx.getItemList(host=containerId)
    stats['itemCount'] = len(items)
    if len(items) == 0:
        return stats
    logger.debug("Item list retrieved: {}".format(items))
    metrics = dckr.metrics(containerId=containerId)
    logger.debug("Metrics collected: {}".format(metrics))

    for item in items:
        if item['key'].startswith('docker.'):
            try:
                section, key = parse_key(item['key'])

                itemValue, itemState = metrics[section][key]
            except Exception:
                itemValue = 'Unknown key'
                itemState = 1
            finally:
                sender.add_item(
                    containerId,
                    item['key'],
                    itemValue,
                    state=itemState
                )
                logger.debug("Sending: key:{} value:{} state:{}".format(
                    item['key'],
                    itemValue,
                    itemState
                ))
    return stats


logger.info("Starting up ...")

while True:
    stats = {'itemCount': 0}
    start_time = time.time()
    try:
        dckr = dockerData()

        dckr.discover_containers()

        zbx = Zabbix()

        sender = zbx.initSender(dataType='items')
        with concurrent.futures.ThreadPoolExecutor(
                max_workers=os.getenv('MAX_WORKERS', 5)) as executor:
            future_to_container = {
                executor.submit(
                    doContainer, containerId, zbx, dckr
                ): containerId for containerId in dckr.containers
            }

            for future in concurrent.futures.as_completed(future_to_container):
                containerId = future_to_container[future]
                try:
                    data = future.result()
                    stats['itemCount'] += data['itemCount']
                except Exception as exc:
                    logger.error('Container {0} generated an exception'.format(
                        containerId))
                    logger.debug(traceback.print_exc())
        sender.send()

        del zbx

        stats['instance'] = doInstance(dckr)

        del dckr
    except AppError as e:
        logger.error("ERROR:", str(e))

    sleep_time = (60 - (time.time() - start_time))

    logger.debug("Sleeping for {0}".format(sleep_time))
    time.sleep(sleep_time)
