#!/usr/local/sbin/charm-env python3

import os
import re
import shlex
import subprocess
import sys

from charms import layer

from etcdctl import EtcdCtl

from charmhelpers.core.hookenv import action_get, action_set, action_fail, action_name


CTL = EtcdCtl()


def action_fail_now(*args, **kw):
    """Call action_fail() and exit immediately."""
    action_fail(*args, **kw)
    sys.exit(0)


def requires_etcd_version(version_regex, human_version=None):
    """Decorator that enforces a specific version of etcdctl be present.

    The decorated function will only be executed if the required version
    of etcdctl is present. Otherwise, action_fail() will be called and
    the process will exit immediately.

    """

    def wrap(f):
        def wrapped_f(*args):
            version = CTL.version()
            if not re.match(version_regex, version):
                required_version = human_version or version_regex
                action_fail_now(
                    "This action requires etcd version {}".format(required_version)
                )
            f(*args)

        return wrapped_f

    return wrap


requires_etcd_v2 = requires_etcd_version(r"2\..*", human_version="2.x")
requires_etcd_v3 = requires_etcd_version(r"3\..*", human_version="3.x")


@requires_etcd_v3
def alarm_disarm():
    """Call `etcdctl alarm disarm`."""
    try:
        output = CTL.run("alarm disarm")
        action_set(dict(output=output))
    except subprocess.CalledProcessError as e:
        action_fail_now(e.output)


@requires_etcd_v3
def alarm_list():
    """Call `etcdctl alarm list`."""
    try:
        output = CTL.run("alarm list")
        action_set(dict(output=output))
    except subprocess.CalledProcessError as e:
        action_fail_now(e.output)


@requires_etcd_v3
def compact():
    """Call `etcdctl compact`."""

    def get_latest_revision():
        try:
            output = CTL.run("endpoint status --write-out json")
        except subprocess.CalledProcessError as e:
            action_fail_now(
                "Failed to determine latest revision for " "compaction: {}".format(e)
            )

        m = re.search(r'"revision":(\d*)', output)
        if not m:
            action_fail_now(
                "Failed to get revision from 'endpoint status' "
                "output: {}".format(output)
            )
        return m.group(1)

    revision = action_get("revision") or get_latest_revision()
    physical = "true" if action_get("physical") else "false"
    command = "compact {} --physical={}".format(revision, physical)
    try:
        output = CTL.run(command)
        action_set(dict(output=output))
    except subprocess.CalledProcessError as e:
        action_fail_now(e.output)


@requires_etcd_v3
def defrag():
    """Call `etcdctl defrag`."""
    try:
        output = CTL.run("defrag")
        action_set(dict(output=output))
    except subprocess.CalledProcessError as e:
        action_fail_now(e.output)


def health():
    """Call etcdctl cluster-health"""
    try:
        output = CTL.cluster_health(True)
        action_set(dict(output=output))
    except subprocess.CalledProcessError as e:
        action_fail_now(e.output)


if __name__ == "__main__":
    ACTIONS = {
        "alarm-disarm": alarm_disarm,
        "alarm-list": alarm_list,
        "compact": compact,
        "defrag": defrag,
        "health": health,
    }

    action = action_name()
    ACTIONS[action]()