#!/usr/bin/env python3
import atexit
import logging
import os
import signal
import subprocess
import sys
import time
from datetime import datetime, timezone

import configargparse
from locust_swarm.functions import *

logging.basicConfig(
    format="%(asctime)s,%(msecs)d %(levelname)-4s [%(filename)s:%(lineno)d] %(message)s",
    datefmt="%Y-%m-%d:%H:%M:%S",
    level=logging.INFO,
)

parser = configargparse.ArgumentParser(
    default_config_files=["~/.locust.conf", "locust.conf"],
    auto_env_var_prefix="LOCUST_",
    description="A tool for running locust in a distributed fashion.",
    epilog="Any parameters not listed here are forwarded to locust unmodified, so go ahead and use things like -u, -r, --host, ... Any env vars starting with LOCUST_ will also be forwarded to the workers.",
)

parser.add_argument("-f", "--locustfile", type=str, dest="testplan")
parser.add_argument("--headless", type=bool, dest="ignore_this")
parser.add_argument(
    "--loadgen-list",
    type=str,
    required=True,
    help="A comma-separated list of ssh servers to act as load generators/workers",
)
parser.add_argument(
    "--processes-per-loadgen", type=int, default=4, help="Number of locust worker processes to spawn for each load gen"
)
parser.add_argument("--loadgens", type=int, default=1, help="Number of load gen servers to use")
parser.add_argument("-L", type=str, dest="loglevel")
parser.add_argument("--port", type=str, default="5557")
parser.add_argument(
    "--remote-master",
    type=str,
    help="An ssh server to use as locust master (default is to run the master on the same machine as swarm). This is useful when rurnning swarm on your workstation if it might become disconnected",
)


args, unrecognized_args = parser.parse_known_args()

if args.loglevel:
    logging.getLogger().setLevel(args.loglevel.upper())

testplan = args.testplan or "locustfile.py"

if "/" in testplan:
    parser.error(
        "Testplan (-f) must be a file in the current directory (I'm lazy and havent fixed support for this yet)"
    )

testplan_filename = os.path.split(testplan)[1]
port = int(args.port)
processes_per_loadgen = args.processes_per_loadgen
loadgens = args.loadgens
worker_process_count = processes_per_loadgen * loadgens
server_list = args.loadgen_list.split(",")
workers = []

try:
    subprocess.check_output(f"ssh -o LogLevel=error -o BatchMode=yes {server_list[0]} true", shell=True)
except Exception:
    logging.error(
        f"Error ssh:ing to loadgen ({server_list[0]}). Maybe you dont have permission to log on to them? Or your ssh key requires a password? (in that case, use ssh-agent)"
    )
    raise

signal.signal(signal.SIGTERM, sig_handler)

while is_port_in_use(port):
    port += 2


workers = get_available_servers_and_lock_them(loadgens, server_list)
worker_procs = []
start_time = datetime.now(timezone.utc)
atexit.register(cleanup, workers, args)

os.environ["LOCUST_RUN_ID"] = start_time.isoformat()
locust_env_vars = []

for varname in os.environ:
    if varname.startswith("LOCUST_") or varname.startswith("PG"):
        if varname == "LOCUST_RPS":
            # distribute the rps over the locust worker processes
            # when client count < worker_process count, not all locust processes will get a client,
            # so use the minium of the two when distributing rps.
            locust_env_vars.append(
                f'LOCUST_RPS="{float(os.environ[varname])/min(worker_process_count, int(args.users))}"'
            )
        else:
            locust_env_vars.append(f'{varname}="{os.environ[varname]}"')

locust_env_vars.append("PYTHONUNBUFFERED=1")  # dont buffer locust worker's stdout, show it immediately

if args.remote_master:
    ssh_command = ["ssh", "-q", args.remote_master, "'", *locust_env_vars, "nohup"]
    bind_only_localhost = []
    ssh_command_end = ["'"]
    check_output(f"ssh -q {args.remote_master} 'pkill -9 -u $USER locust' || true")
    upload(args.remote_master)
else:
    # avoid firewall popups by only binding localhost if running local master (ssh port forwarding):
    bind_only_localhost = ["--master-bind-host=127.0.0.1"]
    ssh_command = []
    ssh_command_end = []

master_command = [
    *ssh_command,
    "locust",
    "--master",
    "--master-bind-port",
    str(port),
    *bind_only_localhost,
    "--expect-workers",
    str(worker_process_count),
    "--headless",
    "-f",
    testplan,
    # "-t",
    # run_time,
    # "-u",
    # users,
    "--exit-code-on-error",
    "0",  # return zero even if there were failed samples (locust default is to return 1)
    *unrecognized_args,
    *ssh_command_end,
]

logging.info(f"launching master: {' '.join(master_command)}")
master_proc = subprocess.Popen(" ".join(master_command), shell=True)

for worker in workers:
    # fail early if master has already terminated
    check_proc_running(master_proc)
    worker_procs.extend(
        start_locust_processes(
            worker, port, args.processes_per_loadgen, locust_env_vars, testplan_filename, args.remote_master
        )
    )

# check that worker procs didnt immediately terminate for some reason (like invalid parameters)
time.sleep(2)
for proc in worker_procs:
    check_proc_running(proc)

# wait for test to complete
while True:
    try:
        code = master_proc.wait(timeout=10)
        break
    except subprocess.TimeoutExpired:
        pass
    except KeyboardInterrupt:
        pass
    # ensure worker procs didnt die before master
    for proc in worker_procs:
        try:
            check_proc_running(proc)
        except subprocess.CalledProcessError as e:
            try:
                code = master_proc.wait(timeout=10)
                break
            except subprocess.TimeoutExpired:
                logging.error(
                    f"worker proc finished unexpectedly with ret code {e.returncode} (and master was still running)"
                )
                raise

logging.info(f"Load gen master process finished (return code {code})")
sys.exit(code)
