# Module importation
import os
import logging
import re
import socket
import sys
import time

from oslo_config import cfg


# This generic function will be called by individual actions below
def send_command_to_haproxy(CONF, command):
    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as s:
        s.connect(CONF.socket)
        s.send(command.encode())
        return s.recv(4096).decode()


def usage():
    print("Usage: haproxy-cmd <action> [--socket <admin-socket-file>] [--config <haproxy-config-file>] [--backend <backend>] [--server <server>] [--frontend <frontend>] [--verbose] [--details]")
    print("Action can be one of:")
    print("list-frontends, list-backends, list-config-backends, list-servers, list-connections, enable-server, drain-server, stop-server, check-safe-to-remove, reload-haproxy")
    print("For these actions, the --backend <backend> parameter is mandatory: list-servers")
    print("For these actions, the --backend <backend> and --server <server> parameters are mandatory: list-connections, enable-server, drain-server, stop-server, check-safe-to-remove")
    print("The list-servers action, the --details option can be specified")
    print("For the drain-server action, --wait can be used to wait for server to be fully drained")
    sys.exit(1)


def get_haproxy_cfg(CONF):
    with open(CONF.config, 'r') as f:
        yield from f


def get_socket(CONF):
    """
    Returns first admin socket found
    """
    socket_re = re.compile(r'[^#]\s+socket\s+([^ ]+)\s+.*\badmin\b')
    for line in get_haproxy_cfg(CONF):
        search = socket_re.search(line)
        if search:
            return search.group(1)
    return '/var/lib/haproxy/admin.sock'
            

def list_backends(CONF):
    if CONF.verbose:
        print("Listing backends...")
    backends = send_command_to_haproxy(CONF, "show backend\n")
    lines = backends.split("\n")
    for line in lines:
        if line != "# name" and line != "MASTER" and line != "":
            print(line)

def list_frontends(CONF):
    if CONF.verbose:
        print("Listing frontends...")
    for line in get_haproxy_cfg(CONF):
        line_array = line.strip().split(' ')
        if (line_array[0] == 'frontend'):
            print(line_array[1])

def list_config_backends(CONF):
    # Show backend via socket cannot retrieve backends depending on frontend.
    # Take care with output, as we have to parse haproxy configfile, not sure that service has been restarted after changes or that research is working.
    if CONF.verbose:
        print("Listing config backends...")
    if CONF.frontend:
        text = "".join([line for line in get_haproxy_cfg(CONF)])
        start_index = text.find(f"frontend {CONF.frontend}")
        end_index = text.find("\n\n", start_index)
        result = "\n".join(
            set(re.findall(r"use_backend (\w+)", text[start_index:end_index]))
        )
        result = "\n".join(
            set(re.findall(r"default_backend (\w+)", text[start_index:end_index]))
        )
        print(f"{result}")
    else:
        print("Missing parameter: --frontend")
        usage()


def list_servers(CONF):
    if CONF.backend:
        # List each servers for a specific backend and his state
        op_state_mapping = {
            0: "SRV_ST_STOPPED, The server op state is down",
            1: "SRV_ST_STARTING, The server op state is warming up",
            2: "SRV_ST_RUNNING, The server op state is fully up",
            3: "SRV_ST_STOPPING, The server op state is up but soft-stopping",
        }

        admin_state_mapping = {
            0: "ACTIVE, The server admin state is ACTIVE",
            1: "SRV_ADMF_FMAINT, The server admin state was explicitly forced into maintenance",
            2: "SRV_ADMF_IMAINT, The server admin state has inherited the maintenance status from a tracked server",
            4: "SRV_ADMF_CMAINT, The server admin state is in maintenance because of the configuration",
            8: "SRV_ADMF_FDRAIN, The server admin state was explicitly forced into drain state",
            10: "SRV_ADMF_IDRAIN, The server admin state has inherited the drain status from a tracked server",
            20: "SRV_ADMF_RMAINT, The server admin state is in maintenance because of an IP address resolution failure",
        }

        lines = send_command_to_haproxy(CONF, f"show servers state {CONF.backend}\n").split("\n")
        headers = lines[1].split(" ")
        srv_admin_state_index = headers.index("srv_admin_state")
        srv_op_state_index = headers.index("srv_op_state")
        srv_name_index = headers.index("srv_name")

        for line in lines[2:]:
            fields = line.split(" ")
            if (
                len(fields) > srv_admin_state_index
                and fields[1] == CONF.backend
            ):
                current_srv_op_state = int(fields[srv_op_state_index - 1])
                current_srv_admin_state = int(fields[srv_admin_state_index - 1])
                current_srv_name = str(fields[srv_name_index - 1])

                op_state_description = op_state_mapping.get(
                    current_srv_op_state, "Unknown op state"
                )
                admin_state_description = admin_state_mapping.get(
                    current_srv_admin_state, "Unknown admin state"
                )

                if CONF.details:
                    print(f"{current_srv_name}, {op_state_description}, {admin_state_description}")
                else:
                    print(f"{current_srv_name}")
    else:
        usage()

# Like below, but for internal use, without print...
def _list_conn(CONF):
    lines = send_command_to_haproxy(CONF, f"show servers conn {CONF.backend} \n").split("\n")
    connections_index = (
        lines[0].split(" ").index("used_cur")
    )  # Current connections column index
    for line in lines[1:]:
        fields = line.split(" ")
        if len(fields) > connections_index:
            if fields[0].startswith(CONF.backend + "/" + CONF.server):
                current_connections = int(fields[connections_index - 1])
    return current_connections

def list_connections(CONF):
    if CONF.verbose:
        print(f"Listing servers for backend: {CONF.backend}...")
    if CONF.backend and CONF.server:
        lines = send_command_to_haproxy(CONF, f"show servers conn {CONF.backend} \n").split("\n")
        connections_index = (
            lines[0].split(" ").index("used_cur")
        )  # Current connections column index
        for line in lines[1:]:
            fields = line.split(" ")
            if len(fields) > connections_index:
                if fields[0].startswith(CONF.backend + "/" + CONF.server):
                    current_connections = int(fields[connections_index - 1])
                    if CONF.verbose:
                        print(f"{current_srv_name}, {op_state_description}, {admin_state_description}")
                    else:
                        print(f"{str(current_connections)}")
    else:
        usage()


def check_safe_to_remove(CONF):
    if CONF.verbose:
        print(f"Checking if safe to remove server: {CONF.backend}/{CONF.server}...")
    nb_conn = _list_conn(CONF)
    if nb_conn == 0:
        if CONF.verbose:
            print(f"Server {CONF.server} can be removed from backend {CONF.backend}.")
        sys.exit(0)
    else:
        if CONF.verbose:
            print(f"Server {CONF.server} has ongoing {nb_conn} connection(s) for backend {CONF.backend}.")
        sys.exit(1)

def enable_server(CONF):
    if CONF.verbose:
        print(f"Enabling server: {CONF.backend}/{CONF.server}...")
    if CONF.backend and CONF.server:
        send_command_to_haproxy(CONF, f"enable server {CONF.backend}/{CONF.server}\n")
    else:
        usage()

def drain_server(CONF):
    if CONF.verbose:
        print(f"Draining server: {CONF.backend}/{CONF.server}...")
    if CONF.backend and CONF.server:
        send_command_to_haproxy(CONF, f"set server {CONF.backend}/{CONF.server} state drain\n")
        if CONF.wait:
            nb_conn = _list_conn(CONF)
            cnt = 0
            while nb_conn != 0:
                time.sleep(1)
                cnt = cnt + 1
                # If timeout is zero (the default), then cnt is already 1
                # on first run, so it wont ever timeout (what we want).
                if cnt == CONF.timeout:
                    if CONF.verbose:
                        print("Timed out")
                    sys.exit(1)
                if CONF.verbose:
                    print(".")
                nb_conn = _list_conn(CONF)
            if CONF.verbose:
                print("done.")
    else:
        usage()

def stop_server(CONF):
    if CONF.verbose:
        print(f"Stopping server: {CONF.backend}/{CONF.server}...")
    if CONF.backend and CONF.server:
        send_command_to_haproxy(CONF, f"disable server {CONF.backend}/{CONF.server}\n")
    else:
        usage()


def reload_haproxy(CONF):
    if CONF.verbose:
        print("Reloading HAProxy")
    os.system("systemctl restart haproxy")
    time.sleep(5)
    if os.popen("systemctl is-active haproxy").read().replace("\n", "") != "active":
        raise Exception("Haproxy service not active.")
    uptime = (
        os.popen(r'systemctl status haproxy | grep -Po ".*; \K(.*)(?= ago)"')
        .read()
        .replace("\n", "")
    )
    if CONF.verbose:
        print(f"Haproxy service has been successfully reloaded with {uptime} uptime.")

#############################################
### Command line parsing with oslo.config ###
#############################################

CLI_OPTS = [
    cfg.StrOpt('backend', short='b',
               help='Backend to use'),

    cfg.StrOpt('socket', short='k',
               help='Socket file to use'),

    cfg.StrOpt('config', default='/etc/haproxy/haproxy.cfg', short='c',
               help='Config file to use'),

    cfg.BoolOpt('details', default=False, short='d',
               help='Wait for server to be drained'),

    cfg.BoolOpt('wait', default=False, short='w',
               help='Wait until drain is finished'),

    cfg.IntOpt('timeout', default=0, short='t',
               help='Timeout for drain'),

    cfg.BoolOpt('verbose', default=False, short='v',
               help='Show details'),

    cfg.StrOpt('frontend', short='f',
               help='Frontend to use'),

    cfg.StrOpt('server', short='s',
               help='Server to operate on'),

]

# Map actions to functions
ACTIONS = {
    'list-backends': list_backends,
    'list-frontends': list_frontends,
    'list-config-backends': list_config_backends,
    'list-servers': list_servers,
    'list-connections': list_connections,
    'enable-server': enable_server,
    'drain-server': drain_server,
    'stop-server': stop_server,
    'check-safe-to-remove': check_safe_to_remove,
    'reload-haproxy': reload_haproxy,
}

def main():
    action = sys.argv[1] if len(sys.argv) > 1 else None

    CONF = cfg.ConfigOpts()
    CONF.register_cli_opts(CLI_OPTS)
    CONF(sys.argv[2:])

    if CONF.socket is None:
        CONF.socket = get_socket(CONF)

    # Perform the corresponding action
    if action in ACTIONS:
        ACTIONS[action](CONF)
        sys.exit(0)
    else:
        usage()