#!/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

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

setup_path()

from typing import Any, Dict, List, Optional

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:
    if 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) -> None:
    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}")
    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 queue_id, 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)
