#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) 2018 Procube Co., Ltd.

import docker
import logging
from logging import handlers
import json
import os
import argparse
from datetime import datetime
import re
import uptime
# logging.basicConfig(level=os.environ.get('HIVE_LOG_LEVEL', logging.INFO))
DAEMON = None


# python 3.6 does not support fromisoformat
def fromisoformat(str):
  # strptime does not support nano second
  return datetime.strptime(re.sub(r'^(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d{6}).*$', r'\1', str), '%Y-%m-%dT%H:%M:%S.%f')


def service_uptime(client, logger, service_name):
  try:
    for service in client.services.list():
      logger.debug(f'traverse services in {service.name}')
      if service.name == service_name:
        logger.debug(f'found service: {service_name}')
        min = -1
        for task in service.tasks():
          logger.debug(task.get('DesiredState'))
          if task.get('DesiredState') == 'running':
            status = task.get('Status')
            if status.get('State') == 'running':
              logger.debug(f'found runinig task start_time={status.get("Timestamp")}')
              # s_uptime = int((datetime.now() - datetime.fromisoformat(status.get('Timestamp'))).total_seconds())
              s_uptime = int((datetime.utcnow() - fromisoformat(status.get('Timestamp'))).total_seconds())
              if min == -1 or s_uptime < min:
                min = s_uptime
        return min
    logger.error(f'service {service_name} is not found')
  except Exception as e:
    logger.exception(f'fail to get uptime for "{service_name}": {e}')
  return 0


def replicas(client, logger, service_name):
  try:
    for service in client.services.list():
      logger.debug(f'traverse services in {service.name}')
      if service.name == service_name:
        logger.debug(f'found service: {service_name}')
        desired_count = 0
        running_count = 0
        for task in service.tasks():
          if task.get('DesiredState') == 'running':
            desired_count += 1
            status = task.get('Status')
            if status.get('State') == 'running':
              running_count += 1
        return running_count * 100 / desired_count if desired_count > 0 else 0
    logger.error(f'service {service_name} is not found')
  except Exception as e:
    logger.exception(f'fail to get replicas for "{service_name}": {e}')
  return 0


def discover(client, logger, standalone=False):
  try:
    for service in client.services.list():
      logger.debug(f'found services : {service.name}')
      if (not standalone) or service.attrs.get('Spec', {}).get('Labels', {}).get('HIVE_STANDALONE', "False") == 'True':
        yield {'{#SERVICE_NAME}': service.name}
  except Exception as e:
    logger.exception(f'fail to get list of services: {e}')


def read_blacklist():
  try:
    with open('/etc/zabbix/service_discovery_blacklist') as f:
      return list(map(lambda x: re.compile(x.replace('\n', '')), f.readlines()))
  except FileNotFoundError:
    return []


def discover_innerservice(clients, logger):
  blacklist = read_blacklist()
  try:
    for service in next(iter(clients.values())).services.list():
      labels = service.attrs.get('Spec', {}).get('Labels', {})
      if labels.get('HIVE_STANDALONE', "False") == 'True':
        logger.debug(f'traverse inner services in standalone container : {service.name}')
        innerservices = set()
        for task in service.tasks():
          if task.get('DesiredState') == 'running':
            status = task.get('Status')
            if status.get('State') == 'running':
              logger.debug(f'found runinig task on {task.get("NodeID")}')
              container_id = task.get('Status', {}).get('ContainerStatus', {}).get('ContainerID', '')
              if len(container_id) == 0:
                logger.error(f'fail to get container ID for service "{service.name}" on node "{task.get("NodeID")}"')
                return
              container = clients[task.get('NodeID')].containers.get(container_id)
              (rc, data) = container.exec_run(
                  'systemctl list-unit-files --type service --no-pager --no-legend',
                  user='root')
              decoded_data = ''
              try:
                decoded_data = data.decode('utf-8')
              except Exception:
                pass
              if rc != 0:
                logger.error(f'fail to execute list-units on node "{task.get("NodeID")}": {decoded_data}')
                return
              for l in decoded_data.splitlines():
                sname = l.split()[0]
                status = l.split()[1]
                if status == 'enabled' and ('@' not in sname) and not any(map(lambda pattern: pattern.match(sname), blacklist)):
                  (rc, data) = container.exec_run(
                      'systemctl show -p Type ' + sname, user='root')
                  decoded_data = ''
                  try:
                    decoded_data = data.decode('utf-8')
                  except Exception:
                    pass
                  if rc != 0:
                    logger.error(f'fail to execute show -p Type {sname} in service "{service.name}" on node "{task.get("NodeID")}": {decoded_data}')
                    return
                  key_value = decoded_data.split('=')
                  logger.debug(f'type of innner service {sname} of service "{service.name}" on node "{task.get("NodeID")}" is "{key_value[1]}"')
                  if len(key_value) != 2 or key_value[0] != 'Type':
                    logger.error(f'fail to parse output "{decoded_data}" of command systemctl show -p Type {sname}' +
                                 f' in "{service.name}" on node "{task.get("NodeID")}"')
                    return -1
                  if key_value[1] not in ['oneshot', 'dbus'] and sname not in innerservices:
                    innerservices.add(sname)
        for sname in innerservices:
          yield {'{#SERVICE_NAME}': service.name, '{#INNER}': sname.replace('@', '%')}
  except Exception as e:
    logger.exception(f'fail to get inner services: {e}')


def service_uptime_innerservice(clients, logger, service_name, inner):
  try:
    sl = [s for s in next(iter(clients.values())).services.list(filters={'name': service_name}) if s.name == service_name]
    if len(sl) != 1:
      logger.error(f'fail to get service {service_name} count of service={len(sl)}')
      return -1
    min = -1
    for task in sl[0].tasks():
      if task.get('DesiredState') == 'running':
        status = task.get('Status')
        if status.get('State') == 'running':
          logger.debug(f'found runinig task on {task.get("NodeID")}')
          container_id = task.get('Status', {}).get('ContainerStatus', {}).get('ContainerID', '')
          if len(container_id) == 0:
            logger.error(f'fail to get container ID for service "{service_name}" on node "{task.get("NodeID")}"')
            return
          container = clients[task.get('NodeID')].containers.get(container_id)
          (rc, data) = container.exec_run('systemctl show -p ActiveEnterTimestampMonotonic ' + inner, user='root')
          decoded_data = ''
          try:
            decoded_data = data.decode('utf-8')
          except Exception:
            pass
          if rc != 0:
            logger.error(f'fail to get Timestamp of service "{inner}@{service_name}" on node "{task.get("NodeID")}": {decoded_data}')
            return -1
          key_value = decoded_data.split('=')
          if len(key_value) != 2:
            logger.error(f'fail to parse output "{decoded_data}" of command systemctl show -p ActiveEnterTimestampMonotonic {inner}' +
                         f' in "{service_name}" on node "{task.get("NodeID")}"')
            return -1
          s_uptime = int(uptime.uptime() - (float(key_value[1]) / 1000000))
          if min == -1 or s_uptime < min:
            min = s_uptime
    return min
  except Exception as e:
    logger.exception(f'fail to get uptime for inner service "{inner}" in container "{service_name}": {e}')
  return -1


def failed_innerservice_count(clients, logger, service_name):
  blacklist = read_blacklist()
  try:
    sl = [s for s in next(iter(clients.values())).services.list(filters={'name': service_name}) if s.name == service_name]
    if len(sl) != 1:
      logger.error(f'fail to get service {service_name} count of service={len(sl)}')
      return -1
    failed_count = 0
    for task in sl[0].tasks():
      if task.get('DesiredState') == 'running':
        status = task.get('Status')
        if status.get('State') == 'running':
          logger.debug(f'found runinig task on {task.get("NodeID")}')
          container_id = task.get('Status', {}).get('ContainerStatus', {}).get('ContainerID', '')
          if len(container_id) == 0:
            logger.error(f'fail to get container ID for service "{service_name}" on node "{task.get("NodeID")}"')
            return
          container = clients[task.get('NodeID')].containers.get(container_id)
          (rc, data) = container.exec_run(
              'systemctl list-units --type=service --no-pager --no-legend --state=failed --all',
              user='root')
          decoded_data = ''
          try:
            decoded_data = data.decode('utf-8')
          except Exception:
            pass
          if rc != 0:
            logger.error(f'fail to execute list-units in service "{service_name}" on node "{task.get("NodeID")}: {decoded_data}"')
            return 99
          for l in decoded_data.splitlines():
            sname = l.split()[0]
            logger.debug(f'found failed service "{sname}" in service "{service_name}" on node "{task.get("NodeID")}"')
            if not any(map(lambda pattern: pattern.match(sname), blacklist)):
              failed_count += 1
    return failed_count
  except Exception as e:
    logger.exception(f'fail to get failed service count in container "{service_name}": {e}')
  return 99


def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('servers', metavar='Server', nargs='+',
                      help='list of docker swarm node server')
  group = parser.add_mutually_exclusive_group()
  group.add_argument("--discover", help="discover services and print zabbix discover data format json.", action="store_true")
  group.add_argument("--discover-standalone", help="discover standalone type services and print zabbix discover data format json.", action="store_true")
  group.add_argument("--discover-innerservice",
                     help="discover inner-services in a standalone type container and print zabbix discover data format json.", action="store_true")
  group.add_argument("--uptime", help="print least uptime for the service.", metavar='service')
  group.add_argument("--replicas", help="print available running task persentage for the service.", metavar='service')
  parser.add_argument("--inner", help="name of inner service for standalone container.", metavar='inner')
  group.add_argument("--failed-innerservice-count",
                     help="count number of failed serivces in standalone type container and print", metavar='service')
  args = parser.parse_args()
  logger = logging.getLogger('docker-service')
  logger.setLevel(os.environ.get('HIVE_LOG_LEVEL', logging.INFO))
  syslog_handler = handlers.SysLogHandler(address="/dev/log", facility=handlers.SysLogHandler.LOG_LOCAL2)
  logger.addHandler(syslog_handler)
  clients = {}
  dot_docker = os.environ['HOME'] + '/.docker'
  tls_config = docker.tls.TLSConfig(ca_cert=dot_docker + '/ca.pem', verify=dot_docker + '/ca.pem',
                                    client_cert=(dot_docker + '/cert.pem', dot_docker + '/key.pem'))
  for server in args.servers:
    try:
      client = docker.DockerClient(base_url=f'tcp://{server}:2376', tls=tls_config)
      info = client.info()
    except (docker.errors.APIError, docker.errors.DockerException) as e:
      logger.error(f'fail to initialize docker client for server {server}: {e}')
      continue
    node_id = info.get('Swarm', {}).get('NodeID', '')
    if len(node_id) > 0:
      clients[node_id] = client
    else:
      logger.error(f'fail to get node id for server {server}')

  if args.discover:
    print(json.dumps(dict(data=[v for v in discover(next(iter(clients.values())), logger)])))
  elif args.discover_standalone:
    print(json.dumps(dict(data=[v for v in discover(next(iter(clients.values())), logger, standalone=True)])))
  elif args.discover_innerservice:
    print(json.dumps(dict(data=[v for v in discover_innerservice(clients, logger)])))
  elif args.uptime:
    if args.inner:
      print(json.dumps(service_uptime_innerservice(clients, logger, args.uptime, args.inner.replace('%', '@'))))
    else:
      print(json.dumps(service_uptime(next(iter(clients.values())), logger, args.uptime)))
  elif args.replicas:
    print(json.dumps(replicas(next(iter(clients.values())), logger, args.replicas)))
  elif args.failed_innerservice_count:
    print(json.dumps(failed_innerservice_count(clients, logger, args.failed_innerservice_count)))
  else:
    logger.error('command option is required')


if __name__ == "__main__":
    main()
