#!/usr/bin/env python3
import argparse
import configparser
import os
import re
import subprocess
import sys

from lib.zulip_tools import assert_running_as_root, parse_os_release

assert_running_as_root()
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

parser = argparse.ArgumentParser(description="Run puppet")
parser.add_argument('--force', '-f', action="store_true",
                    help="Do not prompt with proposed changes")
parser.add_argument('--noop', action="store_true",
                    help="Do not apply the changes")
parser.add_argument('--config', action="store",
                    default="/etc/zulip/zulip.conf",
                    help="Alternate zulip.conf path")
args, extra_args = parser.parse_known_args()

config = configparser.RawConfigParser()
config.read(args.config)

distro_info = parse_os_release()
puppet_config = """
Exec { path => "/usr/sbin:/usr/bin:/sbin:/bin" }
"""

for pclass in re.split(r'\s*,\s*', config.get('machine', 'puppet_classes')):
    puppet_config += f"include {pclass}\n"

# We use the puppet configuration from the same Zulip checkout as this script
scripts_path = os.path.join(BASE_DIR, "scripts")
puppet_module_path = os.path.join(BASE_DIR, "puppet")
puppet_cmd = ["puppet", "apply", "--modulepath", puppet_module_path, "-e", puppet_config]
if args.noop:
    puppet_cmd += ["--noop"]
puppet_cmd += extra_args

# Set the scripts path to be a factor so it can be used by puppet code
puppet_env = os.environ.copy()
puppet_env["FACTER_zulip_conf_path"] = args.config
puppet_env["FACTER_zulip_scripts_path"] = scripts_path

# This is to suppress puppet warnings with ruby 2.7.
if (distro_info['ID'], distro_info['VERSION_ID']) in [('ubuntu', '20.04')]:
    puppet_env["RUBYOPT"] = "-W0"

if not args.noop and not args.force:
    subprocess.check_call([*puppet_cmd, '--noop', '--show_diff'], env=puppet_env)

    do_apply = None
    while do_apply != 'y':
        sys.stdout.write("Apply changes? [y/N] ")
        sys.stdout.flush()
        do_apply = sys.stdin.readline().strip().lower()
        if do_apply == '' or do_apply == 'n':
            sys.exit(0)

ret = subprocess.call([*puppet_cmd, '--detailed-exitcodes'], env=puppet_env)
# ret = 0 => no changes, no errors
# ret = 2 => changes, no errors
# ret = 4 => no changes, yes errors
# ret = 6 => changes, yes errors
if ret != 0 and ret != 2:
    sys.exit(1)