#!/usr/bin/env python3
# Author: Maurice Makaay, XS4ALL
# {{ ansible_managed }}

import re
import json
import sys
import os
import configparser
from mysql.connector import MySQLConnection


lockfile = "{{ marked_down_lockfile }}"


def _get_mode_from_argv():
    mode = "json"
    if len(sys.argv) > 1:
        if sys.argv[1] == "--haproxy":
            mode = "haproxy"
        elif sys.argv[1] == "--json":
            mode = "json"
        else:
            raise "Invalid argument(s) used (you can only use --haproxy or --json)."
    return mode


def _connect_to_db():
    try:
      config = configparser.ConfigParser()
      config.read("/etc/mysql/debian.cnf")
      user = config["client"]["user"]
      password = config["client"]["password"]
      socket = config["client"]["socket"]
      return MySQLConnection(
          host="localhost",
          database="mysql",
          user=user,
          password=password,
          unix_socket=socket)
    except:
      return None

def _init_response():
    return {
       'cluster_size': 0,
       'cluster_status': None,
       'connected': 'OFF',
       'last_committed': 0,
       'local_state_comment': None,
       'read_only': 'OFF',
       'ready': 'OFF',
       'safe_to_bootstrap': 0,
       'seqno': None,
       'sst_method': None,
       'uuid': None,
       'server_version': None,
       'innodb_version': None,
       'protocol_version': None,
       'wsrep_patch_version': None
    }


def _add_global_status(response, db):
    for key, value in _query("SHOW GLOBAL STATUS LIKE 'wsrep_%'", db):
        key = re.sub('^wsrep_', '', key)
        if key in response:
            response[key] = value


def _add_global_variables(response, db):
    query = """SHOW GLOBAL VARIABLES WHERE Variable_name IN (
            'read_only', 'wsrep_sst_method',
            'innodb_version', 'protocol_version', 'version',
            'wsrep_patch_version'
        )"""
    for key, value in _query(query, db):
        if key == "version":
            key = "server_version"
        if key == "wsrep_sst_method":
            key = "sst_method"
        response[key] = value


def _query(query, db):
    try:
      cursor = db.cursor()
      cursor.execute(query)
      return cursor.fetchall()
    except:
      return []


def _add_grastate(response):
    try:
        f = open("/var/lib/mysql/grastate.dat", "r")
        for line in f:
            if line.startswith('#') or re.match('^\s*$', line):
                continue
            line = re.sub('\s+$', '', line)
            key, value = re.split(':\s+', line, maxsplit=1)
            if key in response:
                response[key] = value
        response['cluster_size'] = int(response['cluster_size'])
        response['seqno'] = int(response['seqno'])
        response['safe_to_bootstrap'] = int(response['safe_to_bootstrap'])
    except:
       pass


def _add_manually_disabled(response):
    response["manually_disabled"] = os.path.isfile(lockfile);


def _evaluate_safe_to_use(response):
    '''
    Evaluate if it is safe to use this node for requests. Inspiration:
    https://severalnines.com/resources/tutorials/mysql-load-balancing-haproxy-tutorial
    '''
    status = response['local_state_comment']
    is_read_only = response['read_only'] != 'OFF'
    is_ready = response['ready'] == 'ON'
    is_connected = response['connected'] == 'ON'
    method = response['sst_method']
    is_using_xtrabackup = method is not None and method.startswith("xtrabackup")

    safe_to_use = False
    comment = None

    if response['manually_disabled']:
        comment = "The node has been manually disabled (file %s exists)" % lockfile
    elif status is None:
        comment = "The MySQL server seems not to be running at all"
    elif status == 'Synced':
        if is_read_only:
            comment = "Status is 'Synced', but database is reported to be read-only"
        elif not is_ready:
            comment = "Status is 'Synced', but database reports WSS not ready"
        elif not is_connected:
            comment = "Status is 'Synced', but database reports WSS not being connected"
        else:
            safe_to_use = True
            comment = "Status is 'Synced' and database is writable"
    elif status == 'Donor':
        if is_using_xtrabackup:
            safe_to_use = True
            comment = "Status is 'Donor', but using safe '%s' as the SST method" % method
        else:
            comment = "Status is 'Donor', and xtrabackup(-v2) is not used for SST"
    else:
        comment = "Galera status is not 'Synced', but '%s'" % status
    response['safe_to_use'] = safe_to_use
    response['safe_to_use_comment'] = comment


def _output_response(response, mode):
    json_data = json.dumps(response, indent=4, sort_keys=True) + "\r\n"
    if mode == "json":
        print(json_data)
    else:
        if response["safe_to_use"]:
            print("HTTP/1.1 200 OK", end="\r\n")
        else:
            print("HTTP/1.1 503 Service Unavailable", end="\r\n")
        print("Content-Length: ", len(json_data), end="\r\n")
        print("Keep-Alive: no", end="\r\n")
        print("Content-Type: Content-Type: application/json", end="\r\n\r\n")
        print(json_data, end="")


response = _init_response()
db = _connect_to_db()
if db is None:
    response['safe_to_use'] = False
    response['safe_to_use_comment'] = "Connection to MySQL server failed"
else:
    _add_global_status(response, db)
    _add_global_variables(response, db)
    db.close()
    _add_grastate(response)
    _add_manually_disabled(response)
    _evaluate_safe_to_use(response)
mode = _get_mode_from_argv()
_output_response(response, mode)