#!/usr/bin/python3
# -*- coding: utf-8 -*-

from __future__ import annotations
import subprocess
import json
import sys
from ipaddress import ip_address
from datetime import datetime

LOG = "/tmp/mmxHANodeSetup.log"
PsmClient = "/usr/local/bluecat/PsmClient"


def log(input):
    inputStr = str(input).strip()
    if not inputStr:
        return
    msg = f"{datetime.now().isoformat(sep=' ', timespec='milliseconds')}: {inputStr}"
    with open(LOG, "a") as log:
        log.write(f"{msg}\n")


def get_valid_ip(input: str | dict):
    """
    Extracts valid IP from the input. If it's a dict, tries to use the "address" key
    Throws exception if IP or default are not valid in any way
    """
    if isinstance(input, dict):
        if not input.get("address"):
            raise Exception(f"Malformed address input: {input}")
        return ip_address(input["address"])
    return ip_address(
        input
    )  # will throw if the string (or stringified object) is not a valid IP


def run_command(command_args: str):
    """Runs a command and returns its output"""
    log(f"Executing `{command_args}`")
    result = subprocess.run(command_args, shell=True, capture_output=True, text=True)
    if result.returncode != 0:
        log(f"Command failed: {result.stderr}")
        log(result.stdout)
        raise Exception(result.stderr.strip() or result.stdout.strip())
    return result


def psm_node(args: str):
    result = run_command(f"{PsmClient} node {args}")
    output = result.stdout.strip()
    if "retcode=ok" not in output:
        log(f"PsmClient node command failed: {output}")
        raise Exception(output)
    log(f"PsmClient node command returned: {output}")
    return result


def psm_node_set(args: str):
    psm_node(f"set {args}")


def psm_node_set_address(args: str, address: str | dict | None):
    """
    `psm_node_set_address("xha-ping-", "1.2.3.4")` calls `PsmClient node set xha-ping-address=1.2.3.4`
    `psm_node_set_address("xha-ping-", "2001::21")` calls `PsmClient node set xha-ping-v6address=2001::21`
    """
    if not address:
        return

    assert args.endswith("-")
    ip = get_valid_ip(address)
    psm_node_set(f"{args}address={ip}" if ip.version == 4 else f"{args}v6address={ip}")


def psm_json(args: dict):
    result = run_command(f"{PsmClient} '{json.dumps(args)}'")
    output = result.stdout.strip()
    if '"result" : "ok"' not in output:
        log(f"PsmClient json command failed: {output}")
        raise Exception(output)
    log(f"PsmClient json command returned: {output}")
    return result


def psm_json_get(destination: str) -> dict:
    output = psm_json({"destination": destination, "operation": "get"}).stdout
    return json.loads(output)["data"]


def psm_json_set(destination: str, data: dict):
    psm_json({"destination": destination, "operation": "set", "data": data})


def run_clish_command(commands: list[str]):
    """Runs a clish command and returns its output (clish always exits with exit code 255, and needs additional environment)."""
    clish_env = {
        "CLISH_PATH": "/usr/local/cli/debug;/usr/local/cli/adonis;/usr/local/cli/adonis_configure;/usr/local/cli/adonis_show;/usr/local/cli/common;/usr/local/cli/configure;/usr/local/cli/show",
    }

    args = []
    for command in commands:
        args.append("-c")
        args.append(command)

    log(f"Executing `clish {args}`")
    result = subprocess.run(
        ["/usr/local/bin/clish"] + args, capture_output=True, text=True, env=clish_env
    )
    output = f"{result.stderr.strip()}\n{result.stdout.strip()}"
    log(f"clish returned: {output}")
    return output


def rollback(interface_backup: dict | None):
    """
    To be run in case of failure to cleanup.
    Sets the xha-mode to 0 and restores the interface configuration to the original, backed-up, state
    """
    log("Initiating rollback")
    psm_node_set("xha-mode=0")
    if interface_backup:
        # restore it to the old state
        log("Restoring interface configuration")
        # for reasons unknown to man, PsmClient hates receiving what it dishes out
        # unless you feed it in bite-sized pieces >_<
        for each in interface_backup["interfaces"]:
            dest = each["name"]
            psm_json_set(f"network/interfaces/{dest}", each)

    try:
        log("Attempting restart of PsmClient node to apply rolled back configuration")
        restart_psm_node_server()
    except:
        pass  # not much we can do here - we failed applying config, and again when applying the rolled back config


def restart_psm_node_server():
    # this temporary unmasking of command server is a workaround because
    # a simple PsmClient node restart fails due to the masked command server
    run_command("systemctl unmask cs.service")
    try:
        psm_node(
            "restart"
        )  # this command will drop the connection Central has to the update agent
    except Exception as e:
        log(f"Error restarting node: {e}")
        raise
    psm_node_set("command-server-enable=0")
    run_command("systemctl mask cs.service")


def apply_csync_key(key: str | None):
    if key:
        log("Applying supplied csync key to /etc/csync2.hagroup.key")
        with open("/etc/csync2.hagroup.key", "w") as key_file:
            key_file.write(key)
        run_command("chmod 600 /etc/csync2.hagroup.key")


def get_interface_backup() -> dict | None:
    """Returns a backup of the interface config or None if that it can't be found"""
    try:
        log("Backing up network interface configuration")
        return psm_json_get("network/interfaces")
    except Exception as e:
        log(f"Got error while backing up interface configuration: {e}")
        return None


def configure_backbone(address_with_mask: dict | None):
    try:
        log(
            f"Configuring backbone address: {address_with_mask['address']}/{address_with_mask['cidr']}"
        )
    except:
        log(
            f"Unexpected backbone config: {address_with_mask}"
            if address_with_mask
            else "Backbone address not configured"
        )
        return

    address_with_mask["flags"] = 2  # const unsigned int IF_FLAG_XHA = 0x02;
    psm_json_set(
        "network/interfaces/eth1",
        {
            "active": 1,
            "name": "eth1",
            "type": "physical",
            "v4addresses": [address_with_mask],
        },
    )


def configure_hostname(input: dict):
    hostname = str(input["local-hostname"])
    assert (
        hostname.lower() != str(input["peer-hostname"]).lower()
    )  # central sends unique names

    log(f"Setting hostname of this MDDS to '{hostname}'")

    # only configure hostname if different from desired hostname
    # in case we're repairing an xHA configured node, where clish rejects changes to hostname
    current_hostname = str(psm_json_get("network/hostname")["hostname"])
    if current_hostname.lower() == hostname.lower():
        log(f"Hostname is already set to '{hostname}'")
        return

    # need to use clish to set the hostname because a hostname set with PsmClient is not properly
    # applied to /etc/ha.d/haresources
    output = run_clish_command(["configure system", f"set hostname {hostname}", "save"])
    if "Saved changes successfully" not in output:
        log(f"clish failed: {output}")
        raise Exception(output)


def fix_mmremote_service_definition():
    # if the mmremote.service definition still has WorkingDirectory, we need to remove it
    # since that'd be pointing to the replicated dir, which is unmounted on standby nodes
    log("Fixing mmremote service definition")
    run_command("sed -i '/^WorkingDirectory=/d' /lib/systemd/system/mmremote.service")
    run_command("systemctl daemon-reload")


def set_xha_node_configuration(input: dict):
    log("Setting xHA configuration parameters")
    psm_node_set("xha-mode=1")

    # these two settings currently set to these hard coded sensible defaults as suggested by Dan - can be made configurable
    psm_node_set("xha-dead-time=5")
    psm_node_set("xha-max-replication-rate=50")

    psm_node_set(f"xha-peer-hostname={input['peer-hostname']}")
    psm_node_set_address("xha-ping-", input.get("mutual-ping-address"))
    psm_node_set_address("xha-pip-", input["local-address"])

    psm_node_set_address(
        "xha-pip-management-", input.get("local-mgmt-address") or input["local-address"]
    )
    psm_node_set(f"xha-preferred-hostname={input['mutual-hostname']}")
    psm_node_set_address("xha-primary-peer-", input["peer-address"])

    if input.get("peer-backbone"):
        psm_node_set("xha-secondary-interface=eth1")
        psm_node_set_address("xha-secondary-peer-", input["peer-backbone"])
    else:
        # no backbone: use eth0 and same address as for xha-primary-peer-address
        psm_node_set("xha-secondary-interface=eth0")
        psm_node_set_address("xha-secondary-peer-", input["peer-address"])

    psm_node_set_address("xha-vip-", input["mutual-address"])
    psm_node_set_address(
        "xha-vip-management-", input.get("mutual-mgmt-address") or input["mutual-address"]
    )


def restart_micetro_services():
    log("Restarting micetro services")
    run_command("service mmremote restart")
    run_command("service mmdhcpremote restart")
    run_command("service mmupdater restart")


def loadJSONinput() -> dict:
    """
    IP addresses can be either v4 or v6
    v4 is used in this example.
    Json should be formatted like this:
    {
    "local-address":{"address":"1.1.1.1","cidr":"..."},
    "mutual-address":"1.1.1.1",
    "peer-address":"1.1.1.1",
    "mutual-hostname":"name",
    "peer-hostname":"another name",
    "local-hostname":"name",
    ---optional fields---
    "local-mgmt-address":{"address":"1.1.1.1","cidr":"..."},
    "local-backbone-address":{"address":"1.1.1.1","cidr":"..."},
    "mutual-mgmt-address":"1.1.1.1",
    "mutual-ping-address":"1.1.1.1",
    "peer-backbone":"1.1.1.1",
    "key":"secret"
    }
    """

    # input can either be an argument or through stdin
    if len(sys.argv) == 2:
        log("Loading JSON input from command line argument")
        return json.loads(sys.argv[1])
    log("Loading JSON input from stdin")
    return json.loads(sys.stdin.read())


def main():
    try:
        # start with a fresh log file
        open(LOG, "w").close()

        log("Starting xHA configuration")
        input = loadJSONinput()
        log(f"input: {input}")

        interface_backup = get_interface_backup()

        # key is only included in configuration sent to standby node. If it's in the json input, we apply it
        apply_csync_key(input.get("key"))

        # backbone must be configured manually - it's not configured automatically from the xHA settings like other addresses
        configure_backbone(input.get("local-backbone-address"))

        # set hostname since the pair cannot have the same hostname
        configure_hostname(input)

        # make the mmremote service definition compatible with xHA, such that it can be started on a standby node
        fix_mmremote_service_definition()

        set_xha_node_configuration(input)
        try:
            # the settings are not actually applied until PsmClient node restart is executed here
            log("Restarting PsmClient node to apply xHA configuration")
            restart_psm_node_server()
        except Exception as e:
            rollback(interface_backup)

        restart_micetro_services()
        log("xHA configuration complete")
    except Exception as e:
        log(f"Unexpected error occurred: {e}")
        sys.stderr.write(f"{e}\n")
        sys.exit(1)


if __name__ == "__main__":
    main()
