#!/usr/bin/env python3
# Telegram (with graphs)
"""Send Checkmk notifications to a Telegram chat, including performance graphs.

Successor of the "telegram_notify" package (telegram.py, v1.4.3), ported to
Checkmk 2.4.0:

* parameters are defined with the new Rulesets API (cmk.rulesets.v1), see
  local/lib/python3/cmk_addons/plugins/telegram_graphs/rulesets/
* graphs are rendered by Checkmk itself (ajax_graph_images.py, the same
  mechanism used by the built-in HTML mail plug-in), authenticated with the
  site internal secret as required since 2.3
* besides the graphs of the notified object, the graphs of additional
  services of the same host can be attached (selected by regular expression)

Exit codes (as expected by Checkmk):
  0: notification sent, 1: temporary error (retry), 2: permanent error
"""

import html
import json
import re
import sys
from collections.abc import Sequence

import requests

from cmk.notification_plugins import utils

try:  # proxy support, available in every 2.4 site
    from cmk.utils.http_proxy_config import deserialize_http_proxy_config
except ImportError:  # pragma: no cover
    deserialize_http_proxy_config = None  # type: ignore[assignment]

TELEGRAM_API = "https://api.telegram.org/bot%s/%s"
TELEGRAM_MESSAGE_LEN_LIMIT = 4096
TELEGRAM_CAPTION_LEN_LIMIT = 1024
TELEGRAM_MEDIA_GROUP_LIMIT = 10
DEFAULT_MAX_GRAPHS = 10
REQUEST_TIMEOUT = 30

DEFAULT_HOST_TEMPLATE = (
    "<b>$NOTIFICATIONTYPE$: $LINKEDHOSTNAME$ $EVENT_TXT$</b>\n"
    "<code>\n"
    "Host:     $HOSTNAME$\n"
    "Alias:    $HOSTALIAS$\n"
    "Address:  $HOSTADDRESS$\n"
    "Event:    $EVENT_TXT$\n"
    "Output:   $HOSTOUTPUT$\n"
    "\n"
    "Detail:\n"
    "$LONGHOSTOUTPUT$\n"
    "</code>"
)

DEFAULT_SERVICE_TEMPLATE = (
    "<b>$NOTIFICATIONTYPE$: $LINKEDSERVICEDESC$ $EVENT_TXT$</b>\n"
    "<code>\n"
    "Host:     $HOSTNAME$\n"
    "Alias:    $HOSTALIAS$\n"
    "Address:  $HOSTADDRESS$\n"
    "Service:  $SERVICEDESC$\n"
    "Event:    $EVENT_TXT$\n"
    "Output:   $SERVICEOUTPUT$\n"
    "\n"
    "Detail:\n"
    "$LONGSERVICEOUTPUT$\n"
    "</code>"
)

SERVICE_STATES = {0: "ok", 1: "warn", 2: "crit", 3: "unknown"}
HOST_STATES = {0: "up", 1: "down", 2: "unreachable"}


class TelegramError(Exception):
    def __init__(self, message: str, exit_code: int = 2) -> None:
        super().__init__(message)
        self.exit_code = exit_code


# .--Context helpers------------------------------------------------------.


def is_service_notification(context: dict[str, str]) -> bool:
    return context["WHAT"] == "SERVICE"


def list_param(context: dict[str, str], key: str) -> list[str] | None:
    """Read a list/tuple parameter serialized by Checkmk.

    Lists end up as KEY_1, KEY_2, ... (plus KEY / KEYS joined). Returns None
    when the parameter is not configured at all, [] when it is empty.
    """
    if key not in context and f"{key}S" not in context and f"{key}_1" not in context:
        return None
    values = []
    nr = 1
    while f"{key}_{nr}" in context:
        values.append(context[f"{key}_{nr}"])
        nr += 1
    return values


def bool_param(context: dict[str, str], key: str, default: bool) -> bool:
    if key not in context:
        return default
    return context[key] in ("True", "true", "1", "yes")


def int_param(context: dict[str, str], key: str, default: int) -> int:
    try:
        return int(context[key])
    except (KeyError, ValueError):
        return default


def event_text_template(notification_type: str) -> str:
    # same wording as the built-in mail plug-in
    if notification_type in ("PROBLEM", "RECOVERY"):
        return "$PREVIOUS@HARDSHORTSTATE$ -> $@SHORTSTATE$"
    return {
        "FLAPPINGSTART": "Started Flapping",
        "FLAPPINGSTOP": "Stopped Flapping ($@SHORTSTATE$)",
        "FLAPPINGDISABLED": "Disabled Flapping ($@SHORTSTATE$)",
        "DOWNTIMESTART": "Downtime Start ($@SHORTSTATE$)",
        "DOWNTIMEEND": "Downtime End ($@SHORTSTATE$)",
        "DOWNTIMECANCELLED": "Downtime Cancelled ($@SHORTSTATE$)",
        "ACKNOWLEDGEMENT": "Acknowledged ($@SHORTSTATE$)",
        "CUSTOM": "Custom Notification ($@SHORTSTATE$)",
    }.get(notification_type, notification_type)


def extend_context(context: dict[str, str]) -> None:
    """Escape plugin output for Telegram HTML and add the plugin macros."""
    what = context["WHAT"]

    for key in (f"{what}OUTPUT", f"LONG{what}OUTPUT", "NOTIFICATIONCOMMENT"):
        if key in context:
            value = context[key].replace("\\n", "\n")
            context[key] = html.escape(value, quote=False)

    context["EVENT_TXT"] = utils.substitute_context(
        event_text_template(context["NOTIFICATIONTYPE"]).replace("@", what), context
    )
    context["LINKEDHOSTNAME"] = utils.format_link(
        '<a href="%s">%s</a>',
        utils.host_url_from_context(context),
        html.escape(context["HOSTNAME"]),
    )
    context["LINKEDSERVICEDESC"] = utils.format_link(
        '<a href="%s">%s</a>',
        utils.service_url_from_context(context),
        html.escape(context.get("SERVICEDESC", "")),
    )
    if utils._base_url(context):  # pylint: disable=protected-access
        context["GRAPH_URL"] = utils.graph_url_from_context(context)
    else:
        context["GRAPH_URL"] = ""


def build_message(context: dict[str, str]) -> str:
    """Fill the template and keep the message below Telegram's size limit by
    shortening the long output first, then the regular output."""
    if is_service_notification(context):
        template = context.get("PARAMETER_SERVICE_TEMPLATE") or DEFAULT_SERVICE_TEMPLATE
    else:
        template = context.get("PARAMETER_HOST_TEMPLATE") or DEFAULT_HOST_TEMPLATE

    message = utils.substitute_context(template, context)
    if len(message) <= TELEGRAM_MESSAGE_LEN_LIMIT:
        return message

    context = dict(context)
    what = context["WHAT"]
    for key in (f"LONG{what}OUTPUT", f"{what}OUTPUT"):
        value = context.get(key, "")
        if not value:
            continue
        excess = len(message) - TELEGRAM_MESSAGE_LEN_LIMIT
        keep = max(len(value) - excess - 1, 0)
        context[key] = value[:keep] + ("…" if keep else "")
        message = utils.substitute_context(template, context)
        if len(message) <= TELEGRAM_MESSAGE_LEN_LIMIT:
            return message

    return message[:TELEGRAM_MESSAGE_LEN_LIMIT]


# .--Graphs---------------------------------------------------------------.


def should_send_graphs(context: dict[str, str]) -> bool:
    if not bool_param(context, "PARAMETER_SEND_GRAPHS", True):
        return False
    if context.get("NOTIFICATIONTYPE", "").startswith("ALERTHANDLER"):
        return False

    if is_service_notification(context):
        allowed = list_param(context, "PARAMETER_GRAPHS_SERVICE_STATES")
        state = SERVICE_STATES.get(int_param(context, "SERVICESTATEID", 3), "unknown")
    else:
        allowed = list_param(context, "PARAMETER_GRAPHS_HOST_STATES")
        state = HOST_STATES.get(int_param(context, "HOSTSTATEID", 2), "unreachable")

    # Not configured -> always send graphs
    return allowed is None or state in allowed


def extra_services(context: dict[str, str]) -> list[str]:
    """Services of the notified host whose graphs are attached as well."""
    patterns = list_param(context, "PARAMETER_EXTRA_SERVICES") or []
    if not patterns:
        return []

    applies_to = list_param(context, "PARAMETER_EXTRA_SERVICES_ON")
    if applies_to is None:
        applies_to = ["host"]
    if ("service" if is_service_notification(context) else "host") not in applies_to:
        return []

    try:
        regexes = [re.compile(p) for p in patterns]
    except re.error as e:
        sys.stderr.write(f"Invalid regular expression in additional services: {e}\n")
        return []

    try:
        import livestatus  # pylint: disable=import-outside-toplevel

        rows = livestatus.LocalConnection().query(
            "GET services\n"
            "Columns: description perf_data\n"
            f"Filter: host_name = {livestatus.lqencode(context['HOSTNAME'])}\n"
        )
    except Exception as e:  # pylint: disable=broad-except
        sys.stderr.write(f"Could not query services from Livestatus: {e}\n")
        return []

    own = context.get("SERVICEDESC") if is_service_notification(context) else None
    return sorted(
        desc
        for desc, perf_data in rows
        if perf_data and desc != own and any(r.match(desc) for r in regexes)
    )


def fetch_graphs(context: dict[str, str], service: str | None, limit: int) -> list[utils.Graph]:
    """Render the graphs of one host (service=None) or service via Checkmk."""
    ctx = dict(context, PARAMETER_GRAPHS_PER_NOTIFICATION=str(limit))
    if service is None:
        ctx["WHAT"] = "HOST"
    else:
        ctx["WHAT"] = "SERVICE"
        ctx["SERVICEDESC"] = service
    return utils.render_cmk_graphs(ctx)


def collect_graphs(context: dict[str, str]) -> list[tuple[str, utils.Graph]]:
    """Return (caption, graph) for the notified object plus the additional services."""
    if not should_send_graphs(context):
        return []

    max_graphs = max(int_param(context, "PARAMETER_MAX_GRAPHS", DEFAULT_MAX_GRAPHS), 1)
    host = context["HOSTNAME"]
    result: list[tuple[str, utils.Graph]] = []

    targets: list[str | None] = [
        context["SERVICEDESC"] if is_service_notification(context) else None
    ]
    targets += extra_services(context)

    for service in targets:
        remaining = max_graphs - len(result)
        if remaining <= 0:
            break
        caption = f"{host} / {service}" if service else host
        result += [(caption, g) for g in fetch_graphs(context, service, remaining)]

    return result


# .--Telegram-------------------------------------------------------------.


class Telegram:
    def __init__(self, context: dict[str, str]) -> None:
        self._token = self._bot_token(context)
        self._chat_id = self._get_chat_id(context)
        self._thread_id = context.get("PARAMETER_MESSAGE_THREAD_ID", "")
        self._silent = bool_param(context, "PARAMETER_DISABLE_NOTIFICATION", False)
        self._proxies = self._get_proxies(context)

    @staticmethod
    def _bot_token(context: dict[str, str]) -> str:
        if not any(k.startswith("PARAMETER_BOT_TOKEN") for k in context):
            raise TelegramError("No Telegram bot token configured")
        return utils.get_password_from_env_or_context("PARAMETER_BOT_TOKEN", context)

    @staticmethod
    def _get_chat_id(context: dict[str, str]) -> str:
        # Notification parameter first, then the custom user attribute
        # TELEGRAM_CHAT_ID of the contact.
        for key in ("PARAMETER_CHAT_ID", "CONTACT_TELEGRAM_CHAT_ID"):
            value = context.get(key, "").strip()
            if value and value != "0":
                return value
        raise TelegramError(
            "No chat ID found (neither the notification parameter nor the "
            "contact attribute TELEGRAM_CHAT_ID is set)"
        )

    @staticmethod
    def _get_proxies(context: dict[str, str]) -> dict[str, str] | None:
        serialized = context.get("PARAMETER_PROXY_URL")
        if not serialized or deserialize_http_proxy_config is None:
            return None
        return deserialize_http_proxy_config(serialized).to_requests_proxies()

    def _call(
        self,
        method: str,
        data: dict[str, object],
        files: dict[str, tuple[str, bytes, str]] | None = None,
    ) -> dict:
        data = {"chat_id": self._chat_id, **data}
        if self._thread_id:
            data["message_thread_id"] = self._thread_id
        try:
            if files:
                # multipart/form-data: every value has to be a string
                form = {
                    k: v if isinstance(v, str) else json.dumps(v) for k, v in data.items()
                }
                response = requests.post(
                    TELEGRAM_API % (self._token, method),
                    data=form,
                    files=files,
                    proxies=self._proxies,
                    timeout=REQUEST_TIMEOUT,
                )
            else:
                response = requests.post(
                    TELEGRAM_API % (self._token, method),
                    json=data,
                    proxies=self._proxies,
                    timeout=REQUEST_TIMEOUT,
                )
        except requests.RequestException as e:
            # never leak the bot token (it is part of the URL)
            raise TelegramError(
                f"{method}: connection error: {str(e).replace(self._token, '****')}", 1
            ) from None

        try:
            answer = response.json()
        except ValueError:
            answer = {"ok": False, "description": response.text[:500]}

        if response.status_code == 429 or response.status_code >= 500:
            raise TelegramError(f"{method}: {response.status_code} {answer.get('description')}", 1)
        if not answer.get("ok"):
            raise TelegramError(f"{method}: {response.status_code} {answer.get('description')}", 2)
        return answer.get("result", {})

    def send_message(self, text: str) -> int | None:
        result = self._call(
            "sendMessage",
            {
                "text": text,
                "parse_mode": "HTML",
                "link_preview_options": {"is_disabled": True},
                "disable_notification": self._silent,
            },
        )
        return result.get("message_id") if isinstance(result, dict) else None

    def send_graphs(
        self, graphs: Sequence[tuple[str, utils.Graph]], reply_to: int | None
    ) -> None:
        for start in range(0, len(graphs), TELEGRAM_MEDIA_GROUP_LIMIT):
            chunk = graphs[start : start + TELEGRAM_MEDIA_GROUP_LIMIT]
            common: dict[str, object] = {"disable_notification": True}
            if reply_to:
                common["reply_parameters"] = {
                    "message_id": reply_to,
                    "allow_sending_without_reply": True,
                }

            if len(chunk) == 1:
                caption, graph = chunk[0]
                self._call(
                    "sendPhoto",
                    {**common, "caption": caption[:TELEGRAM_CAPTION_LEN_LIMIT]},
                    files={"photo": (graph.filename, graph.data, "image/png")},
                )
                continue

            media = []
            files = {}
            for nr, (caption, graph) in enumerate(chunk):
                name = f"graph{nr}"
                media.append(
                    {
                        "type": "photo",
                        "media": f"attach://{name}",
                        "caption": caption[:TELEGRAM_CAPTION_LEN_LIMIT],
                    }
                )
                files[name] = (graph.filename, graph.data, "image/png")
            self._call("sendMediaGroup", {**common, "media": media}, files=files)


def main() -> int:
    context = utils.collect_context()
    try:
        extend_context(context)
        telegram = Telegram(context)
        message_id = telegram.send_message(build_message(context))
        sys.stdout.write("Telegram message sent\n")
    except TelegramError as e:
        sys.stderr.write(f"Telegram: {e}\n")
        return e.exit_code

    # The text has been delivered; problems with the graphs must not lead to
    # a retry of the whole notification (the message would be sent twice).
    try:
        graphs = collect_graphs(context)
        if graphs:
            telegram.send_graphs(graphs, message_id)
            sys.stdout.write(f"{len(graphs)} graph(s) sent\n")
    except Exception as e:  # pylint: disable=broad-except
        sys.stderr.write(f"Telegram: graphs could not be sent: {e}\n")
    return 0


if __name__ == "__main__":
    sys.exit(main())
