#!/usr/bin/env python3

"""
Script to provide information about send-receive times.

It supports both munin and nagios outputs

It must be run on a machine that is using the live database for the
Django ORM.
"""
import argparse
import os
import random
import sys
import time
import traceback
from typing import Any, Dict, List, NoReturn, Optional

sys.path.append(".")
sys.path.append("/home/zulip/deployments/current")
from scripts.lib.setup_path import setup_path

setup_path()

import django
import zulip

usage = """Usage: send-receive.py [options] [config]

       'config' is optional, if present will return config info.
        Otherwise, returns the output data."""

parser = argparse.ArgumentParser(usage=usage)
parser.add_argument("--site", default="https://api.zulip.com")

parser.add_argument("--nagios", action="store_true")

parser.add_argument("--insecure", action="store_true")

parser.add_argument("--munin", action="store_true")

parser.add_argument("config", nargs="?")

options = parser.parse_args()

if not options.nagios and not options.munin:
    print("No output options specified! Please provide --munin or --nagios")
    sys.exit(0)

if options.munin and options.config == "config":
    print(
        """graph_title Send-Receive times
graph_info The number of seconds it takes to send and receive a message from the server
graph_args -u 5 -l 0
graph_vlabel RTT (seconds)
sendreceive.label Send-receive round trip time
sendreceive.warning 3
sendreceive.critical 5"""
    )
    sys.exit(0)

sys.path.append("/home/zulip/deployments/current")
os.environ["DJANGO_SETTINGS_MODULE"] = "zproject.settings"

django.setup()

from django.conf import settings

from zerver.models import get_realm, get_system_bot

states = {
    "OK": 0,
    "WARNING": 1,
    "CRITICAL": 2,
    "UNKNOWN": 3,
}


def report(state: str, timestamp: Any = None, msg: Optional[str] = None) -> NoReturn:
    now = int(time.time())
    if msg is None:
        msg = f"send time was {timestamp}"
    state_file_path = "/var/lib/nagios_state/check_send_receive_state"
    with open(state_file_path + ".tmp", "w") as f:
        f.write(f"{now}|{states[state]}|{state}|{msg}\n")
    os.rename(state_file_path + ".tmp", state_file_path)
    print(f"{state}: {msg}")
    sys.exit(states[state])


def send_zulip(sender: zulip.Client, message: Dict[str, Any]) -> None:
    result = sender.send_message(message)
    if result["result"] != "success" and options.nagios:
        report("CRITICAL", msg=f"Error sending Zulip, args were: {message}, {result}")


def get_zulips() -> List[Dict[str, Any]]:
    global last_event_id
    res = zulip_recipient.get_events(queue_id=queue_id, last_event_id=last_event_id)
    if "error" in res.get("result", {}):
        report("CRITICAL", msg="Error receiving Zulips, error was: {}".format(res["msg"]))
    for event in res["events"]:
        last_event_id = max(last_event_id, int(event["id"]))
    # If we get a heartbeat event, that means we've been hanging for
    # 40s, and we should bail.
    if "heartbeat" in (event["type"] for event in res["events"]):
        report("CRITICAL", msg="Got heartbeat waiting for Zulip, which means get_events is hanging")
    return [event["message"] for event in res["events"]]


internal_realm_id = get_realm(settings.SYSTEM_BOT_REALM).id
if (
    "staging" in options.site
    and settings.NAGIOS_STAGING_SEND_BOT is not None
    and settings.NAGIOS_STAGING_RECEIVE_BOT is not None
):
    sender = get_system_bot(settings.NAGIOS_STAGING_SEND_BOT, internal_realm_id)
    recipient = get_system_bot(settings.NAGIOS_STAGING_RECEIVE_BOT, internal_realm_id)
else:
    sender = get_system_bot(settings.NAGIOS_SEND_BOT, internal_realm_id)
    recipient = get_system_bot(settings.NAGIOS_RECEIVE_BOT, internal_realm_id)

zulip_sender = zulip.Client(
    email=sender.email,
    api_key=sender.api_key,
    verbose=True,
    insecure=options.insecure,
    client="ZulipMonitoring/0.1",
    site=options.site,
)

zulip_recipient = zulip.Client(
    email=recipient.email,
    api_key=recipient.api_key,
    verbose=True,
    insecure=options.insecure,
    client="ZulipMonitoring/0.1",
    site=options.site,
)

try:
    res = zulip_recipient.register(event_types=["message"])
    if "error" in res.get("result", {}):
        report("CRITICAL", msg="Error subscribing to Zulips: {}".format(res["msg"]))
    queue_id, last_event_id = (res["queue_id"], res["last_event_id"])
except Exception:
    report("CRITICAL", msg=f"Error subscribing to Zulips:\n{traceback.format_exc()}")
msg_to_send = str(random.getrandbits(64))
time_start = time.time()

send_zulip(
    zulip_sender,
    {
        "type": "private",
        "content": msg_to_send,
        "subject": "time to send",
        "to": recipient.email,
    },
)

msg_content: List[str] = []

while msg_to_send not in msg_content:
    messages = get_zulips()
    seconds_diff = time.time() - time_start

    msg_content = [m["content"] for m in messages]

zulip_recipient.deregister(queue_id)

if options.nagios:
    if seconds_diff > 12:
        report("CRITICAL", timestamp=seconds_diff)
    if seconds_diff > 3:
        report("WARNING", timestamp=seconds_diff)

if options.munin:
    print(f"sendreceive.value {seconds_diff}")
elif options.nagios:
    report("OK", timestamp=seconds_diff)
